#!/usr/bin/env python3
import os
import uuid
from pathlib import Path

import requests
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, PointStruct, VectorParams

DOCS_DIR = Path(os.getenv("RAG_DOCS_DIR", "./docs"))
QDRANT_URL = os.getenv("QDRANT_URL", "http://127.0.0.1:6333")
COLLECTION = os.getenv("QDRANT_COLLECTION", "docs")
EMBED_MODEL = os.getenv("EMBED_MODEL", "nomic-embed-text")
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "1000"))
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "120"))


def read_text(path: Path) -> str:
    if path.suffix.lower() == ".pdf":
        from pypdf import PdfReader

        reader = PdfReader(str(path))
        return "\n".join((p.extract_text() or "") for p in reader.pages)
    return path.read_text(encoding="utf-8", errors="ignore")


def chunk_text(text: str, size: int, overlap: int) -> list[str]:
    chunks = []
    i = 0
    step = max(1, size - overlap)
    while i < len(text):
        chunks.append(text[i : i + size])
        i += step
    return chunks


def embed(text: str) -> list[float]:
    r = requests.post(
        "http://127.0.0.1:11434/api/embeddings",
        json={"model": EMBED_MODEL, "prompt": text},
        timeout=120,
    )
    r.raise_for_status()
    return r.json()["embedding"]


def main() -> None:
    if not DOCS_DIR.exists():
        print(f"Docs directory not found: {DOCS_DIR}")
        return

    client = QdrantClient(url=QDRANT_URL)
    test_vec = embed("ping")
    dim = len(test_vec)

    if not client.collection_exists(COLLECTION):
        client.create_collection(
            collection_name=COLLECTION,
            vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
        )

    points: list[PointStruct] = []
    supported = {".md", ".txt", ".pdf"}

    for path in DOCS_DIR.rglob("*"):
        if not path.is_file() or path.suffix.lower() not in supported:
            continue
        raw = read_text(path).strip()
        if not raw:
            continue
        for idx, chunk in enumerate(chunk_text(raw, CHUNK_SIZE, CHUNK_OVERLAP)):
            vec = embed(chunk)
            points.append(
                PointStruct(
                    id=str(uuid.uuid4()),
                    vector=vec,
                    payload={
                        "source": str(path),
                        "chunk_id": idx,
                        "text": chunk[:4000],
                    },
                )
            )

    if points:
        client.upsert(collection_name=COLLECTION, points=points)
        print(f"Indexed chunks: {len(points)}")
    else:
        print("No documents indexed.")


if __name__ == "__main__":
    main()