"""Document ingestion and retrieval for the DRIFT companion."""

import re
import uuid
from pathlib import Path
from typing import List, Optional

import chromadb

from infj_bot.core.config import PROJECT_ROOT, DATA_DIR
from infj_bot.core.embeddings import (
    get_default_embedding_function,
    LocalEmbeddingFunction,
)

SUPPORTED_TEXT = {
    ".txt",
    ".md",
    ".py",
    ".js",
    ".ts",
    ".jsx",
    ".tsx",
    ".json",
    ".yaml",
    ".yml",
    ".csv",
    ".sh",
    ".html",
    ".css",
    ".rs",
    ".go",
    ".java",
    ".c",
    ".cpp",
    ".h",
}
MAX_INGEST_FILE_BYTES = 2_000_000
MAX_DIRECTORY_FILES = 300


def _is_relative_to(child: Path, parent: Path) -> bool:
    try:
        child.relative_to(parent)
        return True
    except ValueError:
        return False


def _resolve_ingest_path(path: str) -> Path:
    target = Path(path).expanduser()
    if not target.is_absolute():
        target = PROJECT_ROOT / target
    target = target.resolve()
    allowed_roots = [PROJECT_ROOT.resolve(), Path.home().resolve()]
    if not any(_is_relative_to(target, root) for root in allowed_roots):
        raise PermissionError(f"Path {path} is outside the allowed ingestion roots.")
    return target


def _chunk_text(text: str, chunk_size: int = 800, overlap: int = 100) -> List[str]:
    """Split text into overlapping chunks by paragraphs."""
    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
    chunks = []
    current: List[str] = []
    current_len = 0
    for para in paragraphs:
        para_len = len(para)
        if current_len + para_len > chunk_size and current:
            chunks.append("\n\n".join(current))
            # Keep overlap
            overlap_text: List[str] = []
            overlap_len = 0
            for p in reversed(current):
                if overlap_len + len(p) > overlap:
                    break
                overlap_text.insert(0, p)
                overlap_len += len(p)
            current = overlap_text
            current_len = overlap_len
        current.append(para)
        current_len += para_len
    if current:
        chunks.append("\n\n".join(current))
    return chunks


def _read_pdf(path: Path) -> str:
    try:
        from pypdf import PdfReader

        reader = PdfReader(str(path))
        parts = []
        for page in reader.pages:
            text = page.extract_text()
            if text:
                parts.append(text)
        return "\n\n".join(parts)
    except Exception as exc:
        raise RuntimeError(f"PDF read failed: {exc}")


def _read_file(path: Path) -> str:
    if path.stat().st_size > MAX_INGEST_FILE_BYTES:
        raise ValueError(
            f"File too large for ingestion: {path} ({path.stat().st_size} bytes)"
        )
    suffix = path.suffix.lower()
    if suffix == ".pdf":
        return _read_pdf(path)
    return path.read_text(encoding="utf-8", errors="replace")


class DocumentStore:
    def __init__(
        self, persist_directory=None, embedding_function=None, use_semantic=True
    ):
        if persist_directory is None:
            persist_directory = str(DATA_DIR / "chroma_db")
        if embedding_function is None:
            if use_semantic:
                embedding_function = get_default_embedding_function()
            else:
                embedding_function = LocalEmbeddingFunction()
        self.embedding_function = embedding_function
        self.client = chromadb.PersistentClient(path=persist_directory)
        self.collection = self.client.get_or_create_collection(
            name="infj_documents",
            embedding_function=embedding_function,
        )

    def ingest(self, file_path: str, tags: Optional[List[str]] = None) -> int:
        path = _resolve_ingest_path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"Not found: {path}")
        if path.is_dir():
            return self.ingest_directory(path, tags=tags)

        text = _read_file(path)
        if not text.strip():
            return 0

        chunks = _chunk_text(text)
        if not chunks:
            return 0

        ids = [f"doc-{uuid.uuid4().hex[:12]}" for _ in chunks]
        metadatas = []
        for i, _chunk in enumerate(chunks):
            meta = {
                "source": str(path),
                "filename": path.name,
                "chunk_index": i,
                "total_chunks": len(chunks),
                "tags": ",".join(tags or []),
            }
            metadatas.append(meta)

        self.collection.add(
            documents=chunks,
            ids=ids,
            metadatas=metadatas,
        )
        return len(chunks)

    def ingest_directory(
        self, dir_path: Path, tags: Optional[List[str]] = None, recursive: bool = True
    ) -> int:
        total = 0
        scanned = 0
        pattern = "**/*" if recursive else "*"
        for child in Path(dir_path).glob(pattern):
            if child.is_file() and child.suffix.lower() in SUPPORTED_TEXT | {".pdf"}:
                scanned += 1
                if scanned > MAX_DIRECTORY_FILES:
                    raise ValueError(
                        f"Directory ingestion stopped after {MAX_DIRECTORY_FILES} supported files."
                    )
                try:
                    n = self.ingest(str(child), tags=tags)
                    total += n
                except Exception as exc:
                    print(f"[ingest skip] {child}: {exc}")
        return total

    def search(self, query: str, n_results: int = 5) -> List[dict]:
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results,
        )
        out = []
        for i, doc in enumerate(results["documents"][0]):
            meta = results["metadatas"][0][i]
            out.append(
                {
                    "document": doc,
                    "source": meta.get("source", "?"),
                    "filename": meta.get("filename", "?"),
                    "chunk_index": meta.get("chunk_index", 0),
                }
            )
        return out

    def list_sources(self) -> List[str]:
        results = self.collection.get(include=["metadatas"])
        sources = set()
        for meta in results.get("metadatas", []):
            if meta:
                sources.add(meta.get("source", "?"))
        return sorted(sources)

    def delete_source(self, source_path: str) -> int:
        results = self.collection.get(
            where={"source": source_path},
            include=[],
        )
        ids = results.get("ids", [])
        if ids:
            self.collection.delete(ids=ids)
        return len(ids)

    def count(self) -> int:
        return self.collection.count()


def format_doc_results(results: List[dict]) -> str:
    if not results:
        return "No matching documents found."
    lines = []
    for r in results:
        lines.append(
            f"[{r['filename']} chunk {r['chunk_index']}]\n{r['document'][:600]}"
        )
    return "\n---\n".join(lines)


if __name__ == "__main__":
    store = DocumentStore()
    print(f"Document store initialized. Documents: {store.count()}")