"""RAG knowledge base: ingest, search, no hallucination."""

from __future__ import annotations

import json
import logging
import uuid
from pathlib import Path
from typing import Any

from bs4 import BeautifulSoup
from pypdf import PdfReader

from app.config import get_settings
from app.services.embeddings import (
    chunk_text,
    cosine_similarity,
    embed_texts,
    token_overlap_score,
)

logger = logging.getLogger(__name__)


class KnowledgeStore:
    """File-backed vector store (Chroma when available)."""

    def __init__(self) -> None:
        self.settings = get_settings()
        self._collection = None
        self._fallback_path = (
            Path(self.settings.chroma_path) / "fallback.json"
        )
        self._fallback: list[dict[str, Any]] = []
        self._init_store()

    def _init_store(self) -> None:
        path = Path(self.settings.chroma_path)
        path.mkdir(parents=True, exist_ok=True)
        if self.settings.use_json_store:
            self._collection = None
            self._load_fallback()
            logger.info("JSON knowledge store ready")
            return
        try:
            import chromadb

            client = chromadb.PersistentClient(path=str(path))
            self._collection = client.get_or_create_collection(
                name="replyflow_kb",
                metadata={"hnsw:space": "cosine"},
            )
            logger.info("Chroma knowledge store ready")
        except Exception as exc:  # noqa: BLE001
            logger.warning(
                "Chroma unavailable (%s); using JSON fallback",
                exc,
            )
            self._collection = None
            self._load_fallback()

    def _load_fallback(self) -> None:
        if self._fallback_path.exists():
            self._fallback = json.loads(
                self._fallback_path.read_text(encoding="utf-8")
            )
        else:
            self._fallback = []

    def _save_fallback(self) -> None:
        self._fallback_path.parent.mkdir(
            parents=True, exist_ok=True
        )
        self._fallback_path.write_text(
            json.dumps(self._fallback, ensure_ascii=False),
            encoding="utf-8",
        )

    @property
    def chunk_count(self) -> int:
        if self._collection is not None:
            return int(self._collection.count())
        return len(self._fallback)

    async def add_chunks(
        self,
        chunks: list[str],
        metadata: dict[str, Any],
    ) -> int:
        """Embed and store text chunks."""
        if not chunks:
            return 0
        embeddings = await embed_texts(
            chunks,
            force_local=self.settings.seed_local_embeddings,
        )
        ids = [str(uuid.uuid4()) for _ in chunks]
        metadatas = [
            {**metadata, "chunk_index": i}
            for i in range(len(chunks))
        ]

        if self._collection is not None:
            self._collection.add(
                ids=ids,
                documents=chunks,
                embeddings=embeddings,
                metadatas=metadatas,
            )
        else:
            for i, chunk in enumerate(chunks):
                self._fallback.append(
                    {
                        "id": ids[i],
                        "document": chunk,
                        "embedding": embeddings[i],
                        "metadata": metadatas[i],
                    }
                )
            self._save_fallback()
        return len(chunks)

    async def search(
        self,
        query: str,
        top_k: int = 4,
        min_score: float = 0.12,
    ) -> list[dict[str, Any]]:
        """Return relevant chunks above similarity threshold."""
        if self.chunk_count == 0:
            return []

        query_emb = (
            await embed_texts(
                [query],
                force_local=self.settings.seed_local_embeddings,
            )
        )[0]
        use_openai = (
            get_settings().has_openai
            and not self.settings.seed_local_embeddings
        )

        if self._collection is not None:
            result = self._collection.query(
                query_embeddings=[query_emb],
                n_results=min(top_k, max(1, self.chunk_count)),
            )
            docs = result.get("documents", [[]])[0]
            metas = result.get("metadatas", [[]])[0]
            dists = result.get("distances", [[]])[0]
            hits: list[dict[str, Any]] = []
            for doc, meta, dist in zip(docs, metas, dists):
                vector_score = 1.0 - float(dist)
                keyword = token_overlap_score(query, doc)
                score = (
                    vector_score
                    if use_openai
                    else 0.45 * vector_score + 0.55 * keyword
                )
                if score >= min_score:
                    hits.append(
                        {
                            "text": doc,
                            "metadata": meta or {},
                            "score": score,
                        }
                    )
            hits.sort(key=lambda x: x["score"], reverse=True)
            return hits[:top_k]

        scored = []
        for item in self._fallback:
            vector_score = cosine_similarity(
                query_emb, item["embedding"]
            )
            keyword = token_overlap_score(query, item["document"])
            score = 0.35 * vector_score + 0.65 * keyword
            if score >= min_score:
                scored.append(
                    {
                        "text": item["document"],
                        "metadata": item["metadata"],
                        "score": score,
                    }
                )
        scored.sort(key=lambda x: x["score"], reverse=True)
        return scored[:top_k]


_store: KnowledgeStore | None = None


def get_knowledge_store() -> KnowledgeStore:
    """Singleton knowledge store."""
    global _store
    if _store is None:
        _store = KnowledgeStore()
    return _store


def extract_text_from_file(path: Path) -> str:
    """Extract text from pdf, html, md, or txt."""
    suffix = path.suffix.lower()
    raw = path.read_bytes()
    if suffix == ".pdf":
        reader = PdfReader(str(path))
        parts = []
        for page in reader.pages:
            parts.append(page.extract_text() or "")
        return "\n".join(parts)
    text = raw.decode("utf-8", errors="ignore")
    if suffix in {".html", ".htm"}:
        soup = BeautifulSoup(text, "html.parser")
        return soup.get_text(separator="\n")
    return text


async def ingest_path(path: Path, source_type: str) -> int:
    """Ingest a single file into the knowledge store."""
    text = extract_text_from_file(path)
    chunks = chunk_text(text)
    store = get_knowledge_store()
    return await store.add_chunks(
        chunks,
        {
            "filename": path.name,
            "source_type": source_type,
        },
    )


async def ingest_directory(directory: Path) -> int:
    """Ingest all supported files from a directory."""
    total = 0
    patterns = ("*.md", "*.txt", "*.pdf", "*.html", "*.htm")
    files: list[Path] = []
    for pattern in patterns:
        files.extend(directory.glob(pattern))
    for path in sorted(set(files)):
        source = path.suffix.lower().lstrip(".") or "text"
        count = await ingest_path(path, source)
        total += count
        logger.info("Ingested %s (%s chunks)", path.name, count)
    return total