"""
ChromaDB Vector Database utility for document embeddings and semantic search.
"""
import os
from typing import Dict, List, Optional, Tuple

from app.services.document_nodes import timer
from app.services.llm_service import LLMService
from app.utils.page_snapshot_mapping import (
    best_effort_span_for_semantic_chunk,
    build_page_chunks,
    pages_intersecting_span,
)

try:
    from langchain_experimental.text_splitter import SemanticChunker  # type: ignore
    from langchain_openai import OpenAIEmbeddings  # type: ignore
except Exception:  # optional deps
    SemanticChunker = None
    OpenAIEmbeddings = None


class ChromaVectorDB:
    """Manage ChromaDB vector database for document search."""

    def __init__(self, db_path: str = "chroma_db", openai_api_key: str = None):
        self.db_path = db_path
        self.openai_api_key = openai_api_key
        self.client = None
        self.collection = None
        self.init_chromadb()

    def init_chromadb(self):
        """Initialize ChromaDB with OpenAI embeddings"""
        try:
            import chromadb
            from chromadb.config import Settings
            from chromadb.utils import embedding_functions

            openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key=self.openai_api_key,
                model_name=LLMService.EMBEDDING_MODEL,
            )

            self.client = chromadb.PersistentClient(
                path=self.db_path,
                settings=Settings(
                    anonymized_telemetry=False,
                    allow_reset=True,
                ),
            )

            try:
                self.collection = self.client.get_collection(
                    name="documents",
                    embedding_function=openai_ef,
                )
            except Exception:
                self.collection = self.client.create_collection(
                    name="documents",
                    embedding_function=openai_ef,
                    metadata={"hnsw:space": "cosine"},
                )
        except ImportError as e:
            print(f"ChromaDB dependencies not available: {e}")

    def _semantic_chunking_enabled(self) -> bool:
        return str(os.getenv("SEMANTIC_CHUNKING", "0")).strip() in {
            "1",
            "true",
            "True",
            "yes",
            "YES",
        }

    def chunk_text(self, text: str, chunk_size: int = 1000, overlap: int = 200) -> List[str]:
        """Split text into chunks (fixed window or optional semantic chunking)."""
        use_semantic = self._semantic_chunking_enabled()
        if use_semantic and self.openai_api_key:
            try:
                if SemanticChunker is None or OpenAIEmbeddings is None:
                    raise ImportError("LangChain semantic chunking dependencies not available")

                breakpoint_type = (
                    os.getenv("SEMANTIC_CHUNK_BREAKPOINT", "percentile").strip() or "percentile"
                )
                with timer(f"chunking.semantic({breakpoint_type})"):
                    embeddings = OpenAIEmbeddings(
                        model=LLMService.EMBEDDING_MODEL,
                        openai_api_key=self.openai_api_key,
                    )
                    chunker = SemanticChunker(
                        embeddings=embeddings,
                        breakpoint_threshold_type=breakpoint_type,
                    )
                    docs = chunker.create_documents([text or ""])
                    chunks = [d.page_content for d in docs if getattr(d, "page_content", None)]

                min_chars = int(os.getenv("SEMANTIC_CHUNK_MIN_CHARS", "200") or "200")
                max_chars = int(os.getenv("SEMANTIC_CHUNK_MAX_CHARS", "2000") or "2000")
                guarded: List[str] = []
                for c in chunks:
                    c = c.strip()
                    if len(c) < min_chars:
                        continue
                    if len(c) <= max_chars:
                        guarded.append(c)
                        continue
                    start = 0
                    while start < len(c):
                        guarded.append(c[start : start + max_chars])
                        start += max_chars
                return guarded or []
            except Exception:
                pass

        with timer("chunking.fixed_window"):
            chunks: List[str] = []
            start = 0
            text = text or ""
            while start < len(text):
                end = start + chunk_size
                chunks.append(text[start:end])
                start += chunk_size - overlap
            return chunks

    def _chunk_spans_for_indexing(
        self,
        chunks: List[str],
        full_text: str,
        page_spans: Optional[List[Tuple[int, int]]],
        chunk_size: int = 1000,
        overlap: int = 200,
    ) -> List[Optional[Tuple[int, int]]]:
        if not page_spans:
            return [None] * len(chunks)

        if self._semantic_chunking_enabled():
            return [
                best_effort_span_for_semantic_chunk(chunk, full_text) for chunk in chunks
            ]

        out: List[Optional[Tuple[int, int]]] = []
        pos = 0
        for chunk in chunks:
            start = pos
            end = pos + len(chunk)
            out.append((start, end))
            pos += chunk_size - overlap
        return out

    def _attach_page_number_metadata(
        self,
        chunk_metadata: dict,
        char_span: Optional[Tuple[int, int]],
        page_spans: Optional[List[Tuple[int, int]]],
    ) -> None:
        if not char_span or not page_spans:
            return
        pages = pages_intersecting_span(char_span[0], char_span[1], page_spans)
        if pages:
            chunk_metadata["page_number"] = pages[0]

    def _build_chunk_payloads(
        self,
        file_id: str,
        content: str,
        metadata: dict,
        page_spans: Optional[List[Tuple[int, int]]] = None,
        page_texts: Optional[List[str]] = None,
    ) -> Tuple[List[str], List[str], List[dict]]:
        if page_texts:
            return build_page_chunks(file_id, page_texts, metadata)

        chunks = self.chunk_text(content or "")
        chunk_spans = self._chunk_spans_for_indexing(chunks, content or "", page_spans)

        ids: List[str] = []
        documents: List[str] = []
        metadatas: List[dict] = []

        for idx, chunk in enumerate(chunks):
            chunk_id = f"{file_id}_chunk_{idx}"
            ids.append(chunk_id)
            documents.append(chunk)

            chunk_metadata = metadata.copy() if isinstance(metadata, dict) else {}
            chunk_metadata.update(
                {
                    "chunk_index": idx,
                    "total_chunks": len(chunks),
                    "file_id": file_id,
                }
            )
            self._attach_page_number_metadata(
                chunk_metadata,
                chunk_spans[idx] if idx < len(chunk_spans) else None,
                page_spans,
            )
            metadatas.append(chunk_metadata)

        return ids, documents, metadatas

    def add_chunks(
        self,
        ids: List[str],
        documents: List[str],
        metadatas: List[dict],
    ) -> int:
        """Add pre-built chunks to ChromaDB."""
        if not self.collection or not ids:
            return 0

        batch_size = 100
        for i in range(0, len(ids), batch_size):
            batch_ids = ids[i : i + batch_size]
            batch_documents = documents[i : i + batch_size]
            batch_metadatas = metadatas[i : i + batch_size]
            try:
                with timer(f"chroma.add(batch={i // batch_size + 1},chunks={len(batch_ids)})"):
                    self.collection.add(
                        ids=batch_ids,
                        documents=batch_documents,
                        metadatas=batch_metadatas,
                    )
            except Exception as e:
                print(f"Error adding batch {i // batch_size + 1}: {e}")
        return len(ids)

    def add_document(
        self,
        file_id: str,
        content: str,
        metadata: dict,
        page_spans: Optional[List[Tuple[int, int]]] = None,
        page_texts: Optional[List[str]] = None,
    ) -> int:
        """Add document chunks to the vector database."""
        print(f"[FILE] chroma add_document start: file_id={file_id}")
        if not self.collection:
            return 0

        ids, documents, metadatas = self._build_chunk_payloads(
            file_id,
            content,
            metadata,
            page_spans=page_spans,
            page_texts=page_texts,
        )
        count = self.add_chunks(ids, documents, metadatas)
        print(f"[FILE] chroma add_document done: file_id={file_id}, chunks={count}")
        return count

    def chunk_document(
        self,
        file_id: str,
        content: str,
        metadata: dict,
        page_spans: Optional[List[Tuple[int, int]]] = None,
        page_texts: Optional[List[str]] = None,
    ) -> Tuple[List[str], List[str], List[dict]]:
        """Create (ids, documents, metadatas) without storing."""
        return self._build_chunk_payloads(
            file_id,
            content,
            metadata,
            page_spans=page_spans,
            page_texts=page_texts,
        )

    def search(
        self,
        query: str,
        n_results: int = 5,
        folder_path: str = None,
        where: Optional[dict] = None,
    ) -> dict:
        """Search for similar documents"""
        print(
            f"[FILE] chroma search: query={query[:80]}, n_results={n_results}, "
            f"folder={folder_path}, where={where}"
        )
        if not self.collection:
            return {"ids": [], "documents": [], "metadatas": []}

        where_filter = dict(where) if isinstance(where, dict) else None
        if folder_path:
            if where_filter:
                where_filter = {"$and": [where_filter, {"folder_path": folder_path}]}
            else:
                where_filter = {"folder_path": folder_path}

        with timer(f"chroma.query(n_results={n_results})"):
            results = self.collection.query(
                query_texts=[query],
                n_results=n_results,
                where=where_filter,
            )

        return results

    def get_document(self, file_id: str, agent_id: str = None) -> Optional[str]:
        """Get full document content from chunks, optionally filtered by agent_id."""
        print(f"[FILE] chroma get_document: file_id={file_id}, agent_id={agent_id}")
        if not self.collection:
            return None

        try:
            if agent_id:
                where_filter = {
                    "$and": [
                        {"file_id": file_id},
                        {"agent_id": str(agent_id)},
                    ]
                }
            else:
                where_filter = {"file_id": file_id}

            with timer("chroma.get_document"):
                results = self.collection.get(
                    where=where_filter,
                    include=["documents", "metadatas"],
                )

            if results["documents"]:
                chunks_with_meta = list(zip(results["documents"], results["metadatas"]))
                chunks_with_meta.sort(key=lambda x: x[1].get("chunk_index", 0))
                return "\n".join([chunk for chunk, _ in chunks_with_meta])

        except Exception as e:
            print(f"Error getting document: {e}")

        return None

    def delete_document(self, file_id: str):
        """Delete all chunks for a document"""
        if not self.collection:
            print(f"[CHROMA] No collection available for delete_document({file_id})")
            return

        try:
            print(f"[CHROMA] Deleting chunks for file_id: {file_id}")
            results = self.collection.get(where={"file_id": file_id})
            count = len(results["ids"]) if results and "ids" in results else 0

            if count > 0:
                self.collection.delete(where={"file_id": file_id})
                print(f"[CHROMA] Successfully deleted {count} chunks for file_id: {file_id}")
            else:
                print(f"[CHROMA] No chunks found in Chroma for file_id: {file_id}")

        except Exception as e:
            print(f"[CHROMA] Error deleting document {file_id}: {e}")

    def count(self) -> int:
        """Get total number of chunks in the collection"""
        if not self.collection:
            return 0
        return self.collection.count()

    def clear_all(self):
        """Clear all documents from the collection"""
        if not self.collection:
            return

        try:
            self.client.delete_collection("documents")
            self.init_chromadb()
        except Exception as e:
            print(f"Error clearing collection: {e}")
