import json
import logging
import re
import threading
import uuid
from datetime import datetime, timezone

import chromadb
from chromadb.config import Settings

from app.chunker import chunk_text
from app.config import COLLECTION_NAME, DATA_DIR, EMBEDDING_MODEL
from app.embeddings import Embedder
from app.seed import SAMPLES

logger = logging.getLogger("learn_rag")


class DisabledEmbedding:
    """占位嵌入函数。向量全部由 fastembed 预先算好，避免 Chroma 再下载一套模型。"""

    @staticmethod
    def name() -> str:
        return "precomputed"

    def __call__(self, input):
        raise RuntimeError("向量由 fastembed 预先计算")

    def embed_documents(self, input):
        raise RuntimeError("向量由 fastembed 预先计算")

    def embed_query(self, input):
        raise RuntimeError("向量由 fastembed 预先计算")


def to_score(distance: float) -> float:
    if distance <= 0:
        return 1.0
    if distance <= 1:
        return 1 - distance
    return 1 / (1 + distance)


class KnowledgeBase:
    def __init__(self, embedder: Embedder) -> None:
        self.embedder = embedder
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        self.path = DATA_DIR / "documents.json"
        self._lock = threading.Lock()
        self.client = chromadb.PersistentClient(
            path=str(DATA_DIR / "chroma"),
            settings=Settings(anonymized_telemetry=False),
        )
        self.collection = self._open_collection()
        self.docs: list[dict] = []

    def _open_collection(self):
        return self.client.get_or_create_collection(
            name=COLLECTION_NAME,
            embedding_function=DisabledEmbedding(),
            metadata={"hnsw:space": "cosine"},
        )

    def _reset_collection(self) -> None:
        try:
            self.client.delete_collection(COLLECTION_NAME)
        except Exception:
            logger.info("向量集合尚不存在，直接创建")
        self.collection = self._open_collection()

    def startup(self) -> None:
        needs_seed = not self.path.exists()
        payload = {"model": EMBEDDING_MODEL, "documents": []}
        if self.path.exists():
            payload = json.loads(self.path.read_text(encoding="utf-8"))
            self.docs = payload.get("documents", [])
        if needs_seed:
            logger.info("写入示例知识库：云栖书店")
            self._reset_collection()
            for sample in SAMPLES:
                self._index(sample["title"], sample["content"])
            self._save()
            return
        if payload.get("model") != EMBEDDING_MODEL:
            logger.info("嵌入模型已更换，重新构建向量库")
            previous = list(self.docs)
            self.docs = []
            self._reset_collection()
            for doc in previous:
                self._index(doc["title"], doc["content"], doc["id"], doc["created_at"])
            self._save()
            return
        self._reconcile()

    def _save(self) -> None:
        payload = {"model": EMBEDDING_MODEL, "documents": self.docs}
        temporary = self.path.with_suffix(".json.tmp")
        temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        temporary.replace(self.path)

    def _summary(self, doc: dict) -> dict:
        preview = re.sub(r"\s+", " ", doc["content"]).strip()[:80]
        return {
            "id": doc["id"],
            "title": doc["title"],
            "chunk_count": doc["chunk_count"],
            "created_at": doc["created_at"],
            "preview": preview,
        }

    def _index(
        self,
        title: str,
        content: str,
        doc_id: str | None = None,
        created_at: str | None = None,
    ) -> dict:
        chunks = chunk_text(content)
        if not chunks:
            raise ValueError("正文无法切分")
        doc_id = doc_id or uuid.uuid4().hex
        created_at = created_at or datetime.now(timezone.utc).isoformat(timespec="seconds")
        vectors = self.embedder.embed_documents(chunks)
        ids = [f"{doc_id}:{index}" for index in range(len(chunks))]
        metadatas = [
            {
                "doc_id": doc_id,
                "title": title,
                "chunk_index": index,
                "chunk_total": len(chunks),
            }
            for index in range(len(chunks))
        ]
        self.collection.add(
            ids=ids,
            documents=chunks,
            embeddings=vectors,
            metadatas=metadatas,
        )
        record = {
            "id": doc_id,
            "title": title,
            "content": content,
            "created_at": created_at,
            "chunk_count": len(chunks),
        }
        self.docs.append(record)
        return record

    def _reconcile(self) -> None:
        stored = self.collection.get(include=["metadatas"])
        present = {
            meta["doc_id"]
            for meta in (stored.get("metadatas") or [])
            if meta and meta.get("doc_id")
        }
        wanted = {doc["id"] for doc in self.docs}
        for doc_id in present - wanted:
            self.collection.delete(where={"doc_id": doc_id})
        for doc in self.docs:
            if doc["id"] not in present:
                logger.info("重新索引《%s》", doc["title"])
                self.docs = [item for item in self.docs if item["id"] != doc["id"]]
                self._index(doc["title"], doc["content"], doc["id"], doc["created_at"])
        self._save()

    def list_documents(self) -> list[dict]:
        with self._lock:
            return [self._summary(doc) for doc in self.docs]

    def counts(self) -> tuple[int, int]:
        with self._lock:
            return len(self.docs), int(self.collection.count())

    def add(self, title: str, content: str) -> dict:
        with self._lock:
            record = self._index(title, content)
            try:
                self._save()
            except Exception:
                self.docs.pop()
                self.collection.delete(where={"doc_id": record["id"]})
                raise
            logger.info("新增《%s》，%s 个切片", title, record["chunk_count"])
            return self._summary(record)

    def delete(self, doc_id: str) -> bool:
        with self._lock:
            if not any(doc["id"] == doc_id for doc in self.docs):
                return False
            self.collection.delete(where={"doc_id": doc_id})
            self.docs = [doc for doc in self.docs if doc["id"] != doc_id]
            self._save()
            return True

    def search(self, question: str, top_k: int) -> list[dict]:
        with self._lock:
            count = int(self.collection.count())
            if count == 0:
                return []
            vector = self.embedder.embed_query(question)
            found = self.collection.query(
                query_embeddings=[vector],
                n_results=min(top_k, count),
                include=["documents", "metadatas", "distances"],
            )
        documents = (found.get("documents") or [[]])[0] or []
        metadatas = (found.get("metadatas") or [[]])[0] or []
        distances = (found.get("distances") or [[]])[0] or []
        chunks = []
        for text, meta, distance in zip(documents, metadatas, distances):
            meta = meta or {}
            chunks.append(
                {
                    "id": f"{meta.get('doc_id', '')}:{meta.get('chunk_index', 0)}",
                    "title": meta.get("title") or "未命名资料",
                    "text": text or "",
                    "score": round(to_score(float(distance)), 4),
                    "chunk_index": int(meta.get("chunk_index") or 0) + 1,
                    "chunk_total": int(meta.get("chunk_total") or 1),
                }
            )
        chunks.sort(key=lambda item: item["score"], reverse=True)
        for number, chunk in enumerate(chunks, start=1):
            chunk["n"] = number
        return chunks