"""
Multi-tenant vector store: each school gets its own ChromaDB collection.
Collection naming: school_{school_id}_{board_id}
Within a collection, metadata filters handle class → subject → chapter → topic.
"""
import math
import re
import chromadb
from typing import List, Dict, Any, Optional
from langchain_core.documents import Document
from huggingface_hub import InferenceClient
from backend.config import CHROMA_DIR, EMBEDDING_MODEL_NAME, HINDI_EMBEDDING_MODEL_NAME, HF_TOKEN, USE_LOCAL_EMBEDDINGS


def _normalize_vector(vec: Any) -> List[float]:
    """L2 normalize vector for inner-product (ip) similarity."""
    if hasattr(vec, "tolist"):
        vec = vec.tolist()
    while len(vec) > 0 and isinstance(vec[0], list):
        vec = vec[0]
    norm = math.sqrt(sum(x * x for x in vec))
    return [x / norm for x in vec] if norm > 0 else vec


class RobustHFEmbeddings:
    """
    Lightweight embedding wrapper:
    - Default (Cloud): Uses Hugging Face Inference API (<70MB RAM, zero PyTorch load).
    - Local: Only loads PyTorch in RAM if USE_LOCAL_EMBEDDINGS=true.
    """
    def __init__(self, model_name: str):
        self.model_name = model_name
        self._local_emb = None
        self._client = InferenceClient(token=HF_TOKEN or None) if HF_TOKEN else None
        self._query_cache: dict = {}  # in-process cache: text -> embedding vector

        if USE_LOCAL_EMBEDDINGS or not self._client:
            try:
                from langchain_huggingface import HuggingFaceEmbeddings
                self._local_emb = HuggingFaceEmbeddings(
                    model_name=model_name,
                    model_kwargs={"device": "cpu"},
                    encode_kwargs={"normalize_embeddings": True}
                )
            except Exception:
                self._local_emb = None

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        if self._local_emb:
            try:
                return [_normalize_vector(v) for v in self._local_emb.embed_documents(texts)]
            except Exception:
                pass

        if not self._client:
            raise ValueError("No Hugging Face token (HF_TOKEN) or local embedding model available.")

        all_embeddings = []
        batch_size = 32
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            vecs = self._client.feature_extraction(batch, model=self.model_name)
            if hasattr(vecs, "tolist"):
                vecs = vecs.tolist()
            if len(batch) == 1 and isinstance(vecs[0], float):
                vecs = [vecs]
            all_embeddings.extend([_normalize_vector(v) for v in vecs])
        return all_embeddings

    def embed_query(self, text: str) -> List[float]:
        if not text or not text.strip():
            raise ValueError("embed_query received an empty string — check caller before embedding.")
        if text in self._query_cache:
            return self._query_cache[text]

        if self._local_emb:
            try:
                result = _normalize_vector(self._local_emb.embed_query(text))
                self._query_cache[text] = result
                return result
            except Exception:
                pass

        if not self._client:
            raise ValueError("No Hugging Face token (HF_TOKEN) or local embedding model available.")

        vec = self._client.feature_extraction(text, model=self.model_name)
        if hasattr(vec, "tolist"):
            vec = vec.tolist()
        result = _normalize_vector(vec)
        self._query_cache[text] = result
        return result


def _sanitize_collection_name(name: str) -> str:
    """ChromaDB collection names: 3-63 chars, alphanumeric + underscores + hyphens, must start/end with alphanum."""
    name = re.sub(r"[^a-zA-Z0-9_-]", "_", name)
    name = name.strip("_-") or "default"
    # Pad if too short
    if len(name) < 3:
        name = name + "_col"
    return name[:63]


class VectorStoreManager:
    """
    Manages persistent Chroma vector storage with per-tenant collection isolation.
    Each school+board combination gets its own collection.
    """
    def __init__(self, persist_directory: str = str(CHROMA_DIR)):
        self.persist_directory = persist_directory
        self.embeddings = RobustHFEmbeddings(EMBEDDING_MODEL_NAME)
        self._hindi_embeddings = None
        self.client = chromadb.PersistentClient(path=self.persist_directory)

    @property
    def hindi_embeddings(self):
        """Lazy-loaded multilingual model for Hindi & Indian languages (saves RAM until needed)."""
        if self._hindi_embeddings is None:
            self._hindi_embeddings = RobustHFEmbeddings(HINDI_EMBEDDING_MODEL_NAME)
        return self._hindi_embeddings

    def get_embeddings(self, subject: Optional[str] = None, text_sample: str = ""):
        """Route to Hindi multilingual model if subject is Hindi/Sanskrit or text is Devanagari."""
        is_indic = bool(subject and str(subject).strip().lower() in ("hindi", "हिन्दी", "sanskrit", "संस्कृत"))
        has_devanagari = bool(re.search(r"[\u0900-\u097F]", text_sample or ""))
        if is_indic or has_devanagari:
            return self.hindi_embeddings
        return self.embeddings

    def _collection_name(self, school_id: str, board_id: str) -> str:
        return _sanitize_collection_name(f"school_{school_id}_{board_id}")

    def get_tenant_collection(self, school_id: str, board_id: str):
        """Get or create the ChromaDB collection for a specific school+board tenant."""
        name = self._collection_name(school_id, board_id)
        return self.client.get_or_create_collection(
            name=name,
            metadata={"hnsw:space": "ip", "hnsw:search_ef": 30}
        )

    def add_documents(
        self,
        documents: List[Document],
        school_id: str,
        board_id: str
    ) -> int:
        if not documents:
            return 0
        collection = self.get_tenant_collection(school_id, board_id)
        texts = [doc.page_content for doc in documents]
        metadatas = [doc.metadata for doc in documents]
        ids = [
            f"{doc.metadata.get('class_id', 'x')}_{doc.metadata.get('subject', 'x')}_"
            f"{doc.metadata.get('chapter', 'x')}_{i}_{hash(doc.page_content)}"
            for i, doc in enumerate(documents)
        ]
        # Build embedding texts with breadcrumb context — stored doc text stays original
        texts_to_embed = []
        for doc in documents:
            breadcrumb = doc.metadata.get("breadcrumb", "")
            embed_text = (breadcrumb + "\n" + doc.page_content).strip() if breadcrumb else doc.page_content
            texts_to_embed.append(embed_text)

        # Dynamically route embeddings based on subject and script
        subj = documents[0].metadata.get("subject", "") if documents else ""
        emb_engine = self.get_embeddings(subject=subj, text_sample=texts_to_embed[0] if texts_to_embed else "")
        embeddings = emb_engine.embed_documents(texts_to_embed)
        collection.upsert(
            documents=texts,
            embeddings=embeddings,
            metadatas=metadatas,
            ids=ids
        )
        return len(documents)

    def delete_chapter(
        self,
        school_id: str,
        board_id: str,
        class_id: str,
        subject: str,
        chapter: str,
        section_id: Optional[str] = None,
        sub_subject: Optional[str] = None,
        topic: Optional[str] = None
    ) -> int:
        """Delete chunks for a specific chapter/lesson, optionally scoped by section, sub_subject, or topic."""
        collection = self.get_tenant_collection(school_id, board_id)
        conds = [
            {"class_id": str(class_id)},
            {"subject": subject},
            {"chapter": chapter}
        ]
        if section_id:
            conds.append({"section_id": str(section_id)})
        if sub_subject:
            conds.append({"sub_subject": str(sub_subject)})
        if topic:
            conds.append({"topic": str(topic)})

        where_filter = {"$and": conds} if len(conds) > 1 else conds[0]
        try:
            existing = collection.get(where=where_filter)
            if existing and existing.get("ids"):
                collection.delete(ids=existing["ids"])
                return len(existing["ids"])
        except Exception:
            pass
        return 0

    def query_context(
        self,
        query: str,
        school_id: str,
        board_id: str,
        class_id: Optional[str] = None,
        section_id: Optional[str] = None,
        subject: Optional[str] = None,
        sub_subject: Optional[str] = None,
        chapter: Optional[str] = None,
        lesson: Optional[str] = None,
        topic: Optional[str] = None,
        n_results: int = 4
    ) -> List[Dict[str, Any]]:
        collection = self.get_tenant_collection(school_id, board_id)

        effective_lesson = lesson or chapter
        clean_topic = str(topic).strip() if topic else ""
        is_generic_topic = (
            not clean_topic
            or clean_topic.lower() in ("all", "all topics", "all topics (full lesson)", "general", "none")
            or (effective_lesson and clean_topic.lower() == str(effective_lesson).strip().lower())
        )

        filter_conditions = []
        if class_id:
            filter_conditions.append({"class_id": str(class_id)})
        if section_id and section_id not in ("All", ""):
            filter_conditions.append({"section_id": str(section_id)})
        if subject:
            filter_conditions.append({"subject": str(subject)})
        if sub_subject and sub_subject not in ("General", ""):
            filter_conditions.append({"sub_subject": str(sub_subject)})
        if effective_lesson:
            filter_conditions.append({"chapter": str(effective_lesson)})
        if not is_generic_topic:
            filter_conditions.append({"topic": clean_topic})

        where_filter = None
        if len(filter_conditions) == 1:
            where_filter = filter_conditions[0]
        elif len(filter_conditions) > 1:
            where_filter = {"$and": filter_conditions}

        emb_engine = self.get_embeddings(subject=subject, text_sample=query)
        query_embedding = emb_engine.embed_query(query)

        try:
            count = collection.count()
            if count == 0:
                return []
            results = collection.query(
                query_embeddings=[query_embedding],
                n_results=min(n_results, count),
                where=where_filter
            )
            # If no results matched with strict filters, fall back to class/subject/lesson (widest lesson scope)
            if (not results or not results.get("documents") or not results["documents"][0]):
                fallback_conditions = []
                if class_id:
                    fallback_conditions.append({"class_id": str(class_id)})
                if subject:
                    fallback_conditions.append({"subject": str(subject)})
                if effective_lesson:
                    fallback_conditions.append({"chapter": str(effective_lesson)})
                fb_filter = {"$and": fallback_conditions} if len(fallback_conditions) > 1 else (fallback_conditions[0] if fallback_conditions else None)
                if fb_filter:
                    results = collection.query(
                        query_embeddings=[query_embedding],
                        n_results=min(n_results, count),
                        where=fb_filter
                    )
        except Exception:
            return []

        matched_docs = []
        if results and results.get("documents") and results["documents"][0]:
            docs = results["documents"][0]
            metas = results["metadatas"][0] if results.get("metadatas") else [{}] * len(docs)
            for doc_text, meta in zip(docs, metas):
                matched_docs.append({
                    "content": doc_text,
                    "metadata": meta
                })
        return matched_docs

    def get_curriculum(
        self,
        school_id: str,
        board_id: str,
        class_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """Extracts available hierarchy (Subjects -> Chapters/Lessons -> Topics) for a tenant."""
        collection = self.get_tenant_collection(school_id, board_id)
        where_filter = {"class_id": str(class_id)} if class_id else None

        try:
            data = collection.get(where=where_filter, include=["metadatas"])
        except Exception:
            return {}

        hierarchy: Dict[str, Dict[str, set]] = {}
        if data and data.get("metadatas"):
            for meta in data["metadatas"]:
                if not meta:
                    continue
                subj = meta.get("subject", "General")
                chap = meta.get("lesson") or meta.get("chapter", "General")
                top = meta.get("topic", "General")
                if subj not in hierarchy:
                    hierarchy[subj] = {}
                if chap not in hierarchy[subj]:
                    hierarchy[subj][chap] = set()
                hierarchy[subj][chap].add(top)

        result = {}
        for s, chaps in hierarchy.items():
            result[s] = {}
            for c, tops in chaps.items():
                clean_tops = [t for t in tops if t.strip().lower() != c.strip().lower()]
                result[s][c] = sorted(clean_tops)
        return result

    def get_erp_curriculum(
        self,
        school_id: str,
        board_id: str,
        class_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Extracts full ERP hierarchy:
        Class -> Section -> Subject -> Sub-Subject -> Lesson -> Topics
        """
        collection = self.get_tenant_collection(school_id, board_id)
        where_filter = {"class_id": str(class_id)} if class_id else None

        try:
            data = collection.get(where=where_filter, include=["metadatas"])
        except Exception:
            return {}

        tree: Dict[str, Dict[str, Dict[str, Dict[str, Dict[str, set]]]]] = {}

        if data and data.get("metadatas"):
            for meta in data["metadatas"]:
                if not meta:
                    continue
                cls = str(meta.get("class_id", class_id or "General"))
                sec = str(meta.get("section_id", "A"))
                subj = meta.get("subject", "General")
                sub_subj = meta.get("sub_subject", "General")
                lesson = meta.get("lesson") or meta.get("chapter", "General")
                top = meta.get("topic", "General")

                tree.setdefault(cls, {}).setdefault(sec, {}).setdefault(subj, {}).setdefault(sub_subj, {}).setdefault(lesson, set()).add(top)

        # Convert sets to sorted lists, stripping redundant lesson names
        clean_tree = {}
        for c, secs in tree.items():
            clean_tree[c] = {}
            for sec, subjs in secs.items():
                clean_tree[c][sec] = {}
                for s, sub_subjs in subjs.items():
                    clean_tree[c][sec][s] = {}
                    for ss, lessons in sub_subjs.items():
                        clean_tree[c][sec][s][ss] = {}
                        for les, tops in lessons.items():
                            clean_tops = [t for t in tops if t.strip().lower() != les.strip().lower()]
                            clean_tree[c][sec][s][ss][les] = sorted(clean_tops)

        return clean_tree

    def get_training_status(
        self,
        school_id: str,
        board_id: str
    ) -> Dict[str, Any]:
        """Returns chunk counts grouped by class/subject/chapter for the teacher dashboard."""
        collection = self.get_tenant_collection(school_id, board_id)
        try:
            data = collection.get(include=["metadatas"])
        except Exception:
            return {}

        status: Dict[str, Dict[str, Dict[str, Any]]] = {}
        topics_map: Dict[str, set] = {}

        if data and data.get("metadatas"):
            for meta in data["metadatas"]:
                if not meta:
                    continue
                cls = meta.get("class_id", "?")
                subj = meta.get("subject", "?")
                chap = meta.get("chapter", "?")
                key = f"Class {cls}"
                ckey = f"{key}_{subj}_{chap}"

                if key not in status:
                    status[key] = {}
                if subj not in status[key]:
                    status[key][subj] = {}
                if chap not in status[key][subj]:
                    status[key][subj][chap] = {
                        "chunk_count": 0,
                        "source": meta.get("source", ""),
                        "topics": []
                    }
                    topics_map[ckey] = set()

                status[key][subj][chap]["chunk_count"] += 1
                top = meta.get("topic")
                if top and top.strip():
                    topics_map[ckey].add(top.strip())

        # Finalize topics: convert sets to sorted lists, filtering out chapter title if subtopics exist
        for cls_label, subjects in status.items():
            for subj, chapters in subjects.items():
                for chap, info in chapters.items():
                    ckey = f"{cls_label}_{subj}_{chap}"
                    all_tops = topics_map.get(ckey, set())
                    # Filter out redundant chapter name if other subtopics are present
                    clean_tops = [t for t in all_tops if t.lower() != chap.lower()]
                    info["topics"] = sorted(clean_tops) if clean_tops else sorted(list(all_tops))

        return status

    def list_tenants(self) -> List[str]:
        """List all tenant collections (admin utility)."""
        collections = self.client.list_collections()
        return [c.name for c in collections if c.name.startswith("school_")]


vector_manager = VectorStoreManager()