from __future__ import annotations

from typing import List, Optional, Tuple
import math
import operator
import re

from analogai.foundation.settings import config as infra_config
from analogai.foundation.modules.fact.basefact import BaseFact
from analogai.foundation.modules.fact.source.fact_source import FactSource, SourceType
from analogai.foundation.modules.fact.enums.fact_reference_type import FactReferenceType
from analogai.foundation.modules.fact.repositories.fact_repository import FactRepository
from analogai.foundation.modules.fact.repositories.infrastructure.fact_storage_mapper import (
    FactStorageMapper,
)


class _RepositoryFactResolver:
    def __init__(self, repository: "FactRepository") -> None:
        self._repository = repository

    def resolve(self, fact_id: str) -> Optional[BaseFact]:
        return self._repository.get_fact(fact_id)


class MongoFactRepository(FactRepository):
    def __init__(
        self,
        uri: Optional[str] = None,
        database: Optional[str] = None,
        facts_collection: str = "facts",
        links_collection: str = "fact_links",
        turn_index_collection: str = "fact_turn_index",
    ) -> None:
        try:
            from pymongo import MongoClient
        except Exception as exc:  # pragma: no cover - optional dependency
            raise ImportError("pymongo is required for MongoFactRepository") from exc
        mongo_uri = uri or infra_config.get("MONGODB_URI", "")
        mongo_db = database or infra_config.get("MONGODB_DATABASE", "")
        if not mongo_uri or not mongo_db:
            raise ValueError("MongoDB configuration missing (MONGODB_URI/MONGODB_DATABASE)")
        client = MongoClient(mongo_uri)
        db = client[mongo_db]
        self._facts = db[facts_collection]
        self._links = db[links_collection]
        self._turn_index = db[turn_index_collection]
        self._resolver = _RepositoryFactResolver(self)
        self._ensure_indexes()

    def _ensure_indexes(self) -> None:
        try:
            self._facts.create_index([("_id", 1)], unique=True, background=True)
            self._facts.create_index([("user_id", 1), ("agent_id", 1)], background=True)
            self._facts.create_index([("turn_id", 1)], background=True)
            self._facts.create_index([("statement_id", 1)], background=True)
            self._facts.create_index([("notion_id", 1)], background=True)
            self._facts.create_index([("reference_type", 1)], background=True)

            self._links.create_index([("kind", 1), ("entity_id", 1)], unique=True, background=True)
            self._links.create_index([("turn_id", 1)], background=True)

            self._turn_index.create_index([("_id", 1)], unique=True, background=True)
            self._turn_index.create_index([("turn_ids", 1)], background=True)
        except Exception:
            pass

    def register_turn(self, user_id: str, agent_id: str, turn_id: str) -> None:
        key = f"{user_id}:{agent_id}"
        self._turn_index.update_one(
            {"_id": key},
            {"$push": {"turn_ids": turn_id}},
            upsert=True,
        )

    def save_fact(self, fact: BaseFact) -> None:
        doc = FactStorageMapper.to_document(fact)
        self._facts.update_one({"_id": fact.id}, {"$set": doc}, upsert=True)

    def get_fact(self, id: str) -> Optional[BaseFact]:
        doc = self._facts.find_one({"_id": id})
        fact = FactStorageMapper.from_document(doc)
        if fact is None:
            return None
        return self._attach_resolver(fact)

    def delete_fact(self, id: str) -> Optional[BaseFact]:
        if not id:
            return None
        doc = self._facts.find_one({"_id": id})
        if doc is None:
            return None
        self._facts.delete_one({"_id": id})
        fact = FactStorageMapper.from_document(doc)
        return self._attach_resolver(fact) if fact else None

    def update_facts_statement_id(self, turn_id: str, statement_id: str) -> None:
        if not turn_id or not statement_id:
            return
        self._facts.update_many(
            {"turn_id": turn_id},
            {"$set": {"statement_id": statement_id, "reference_type": FactReferenceType.STATEMENT.value}},
        )

    def update_facts_notion_id(self, turn_id: str, notion_id: str) -> None:
        if not turn_id or not notion_id:
            return
        self._facts.update_many(
            {"turn_id": turn_id},
            {"$set": {"notion_id": notion_id, "reference_type": FactReferenceType.NOTION.value}},
        )

    def update_facts_hypothetical_statement_id(self, turn_id: str, statement_id: str) -> None:
        if not turn_id or not statement_id:
            return
        self._facts.update_many(
            {"turn_id": turn_id},
            {"$set": {"statement_id": statement_id, "reference_type": FactReferenceType.HYPOTHETICAL_STATEMENT.value}},
        )

    def get_facts_for_turn(self, turn_id: str) -> List[BaseFact]:
        cursor = self._facts.find({"turn_id": turn_id})
        results: List[BaseFact] = []
        for doc in cursor:
            fact = FactStorageMapper.from_document(doc)
            if fact is None:
                continue
            results.append(self._attach_resolver(fact))
        return results

    def get_facts_by_source(self, source_id: str) -> List[BaseFact]:
        if not source_id:
            return []
        cursor = self._facts.find({"source_id": source_id})
        results: List[BaseFact] = []
        for doc in cursor:
            fact = FactStorageMapper.from_document(doc)
            if fact is None:
                continue
            results.append(self._attach_resolver(fact))
        return results

    def list_source_ids(self) -> List[str]:
        cursor = self._facts.find(
            {"$or": [{"source_id": {"$ne": None}}, {"source_type": {"$ne": None}}]},
            {"source_id": 1, "source_type": 1, "source_value": 1},
        )
        sources: dict[str, str] = {}
        for doc in cursor:
            source_id = _source_id_from_doc(doc)
            if not source_id:
                continue
            sources.setdefault(source_id, source_id)
        return list(sources.keys())

    def get_turn_ids(self, user_id: str, agent_id: str) -> List[str]:
        key = f"{user_id}:{agent_id}"
        doc = self._turn_index.find_one({"_id": key})
        if not doc:
            return []
        return [str(item) for item in doc.get("turn_ids", []) if item]

    def get_statement_ids_for_turn(self, turn_id: str) -> List[str]:
        if not turn_id:
            return []
        cursor = self._links.find({"kind": "statement", "turn_id": turn_id})
        return [str(doc.get("entity_id")) for doc in cursor if doc.get("entity_id")]

    def get_belief_ids_for_turn(self, turn_id: str) -> List[str]:
        if not turn_id:
            return []
        cursor = self._links.find({"kind": "belief", "turn_id": turn_id})
        return [str(doc.get("entity_id")) for doc in cursor if doc.get("entity_id")]

    def link_statement(self, statement_id: str, turn_id: str) -> None:
        if statement_id and turn_id:
            self._set_link("statement", statement_id, turn_id)

    def link_belief(self, belief_id: str, turn_id: str) -> None:
        if belief_id and turn_id:
            self._set_link("belief", belief_id, turn_id)

    def link_notion(self, notion_id: str, turn_id: str) -> None:
        if notion_id and turn_id:
            self._set_link("notion", notion_id, turn_id)

    def link_hypothetical_statement(self, statement_id: str, turn_id: str) -> None:
        if statement_id and turn_id:
            self._set_link("hypothetical_statement", statement_id, turn_id)

    def get_turn_id_for_statement(self, statement_id: str) -> Optional[str]:
        return self._get_link("statement", statement_id)

    def get_turn_id_for_belief(self, belief_id: str) -> Optional[str]:
        return self._get_link("belief", belief_id)

    def get_turn_id_for_notion(self, notion_id: str) -> Optional[str]:
        return self._get_link("notion", notion_id)

    def get_turn_id_for_hypothetical_statement(self, statement_id: str) -> Optional[str]:
        return self._get_link("hypothetical_statement", statement_id)

    def search_text(self, user_id: str, agent_id: str, query: str, top_k: int) -> List[BaseFact]:
        safe = re.escape(query)
        cursor = self._facts.find(
            {
                "user_id": user_id,
                "agent_id": agent_id,
                "text": {"$regex": safe, "$options": "i"},
            }
        ).limit(max(1, int(top_k)))
        results: List[BaseFact] = []
        for doc in cursor:
            fact = FactStorageMapper.from_document(doc)
            if fact is None:
                continue
            results.append(self._attach_resolver(fact))
        return results

    def search_similar(
        self,
        user_id: str,
        agent_id: str,
        embedding: List[float],
        top_k: int,
        similarity_cutoff: Optional[float] = None,
    ) -> List[BaseFact]:
        cursor = self._facts.find(
            {
                "user_id": user_id,
                "agent_id": agent_id,
                "embedding": {"$ne": None},
            }
        )
        scored: List[Tuple[float, BaseFact]] = []
        for doc in cursor:
            fact = FactStorageMapper.from_document(doc)
            if fact is None:
                continue
            fact = self._attach_resolver(fact)
            if not fact.embedding:
                raise RuntimeError("Fact missing embedding")
            score = _cosine_similarity(embedding, fact.embedding)
            if similarity_cutoff is not None and score < similarity_cutoff:
                continue
            scored.append((score, fact))
        scored.sort(key=operator.itemgetter(0), reverse=True)
        return [fact for _, fact in scored[: max(1, int(top_k))]]

    def delete_facts_by_source(self, source_id: str) -> List[BaseFact]:
        if not source_id:
            return []
        query = {"source_id": source_id}
        cursor = self._facts.find(query)
        deleted: List[BaseFact] = []
        for doc in cursor:
            fact = FactStorageMapper.from_document(doc)
            if fact is None:
                continue
            deleted.append(self._attach_resolver(fact))
        if deleted:
            self._facts.delete_many(query)
        return deleted

    def _set_link(self, kind: str, entity_id: str, turn_id: str) -> None:
        key = f"{kind}:{entity_id}"
        self._links.update_one(
            {"_id": key},
            {"$set": {"turn_id": turn_id, "kind": kind, "entity_id": entity_id}},
            upsert=True,
        )

    def _get_link(self, kind: str, entity_id: str) -> Optional[str]:
        key = f"{kind}:{entity_id}"
        doc = self._links.find_one({"_id": key})
        if not doc:
            return None
        return doc.get("turn_id")

    def _attach_resolver(self, fact: BaseFact) -> BaseFact:
        fact._resolver = self._resolver
        fact._deleter = self.delete_fact
        return fact


def _cosine_similarity(a: List[float], b: List[float]) -> float:
    if not a or not b:
        return 0.0
    if len(a) != len(b):
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(y * y for y in b))
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0
    return dot / (norm_a * norm_b)


def _source_id_from_doc(doc: dict) -> Optional[str]:
    if not doc:
        return None
    source_id = doc.get("source_id")
    if source_id:
        return str(source_id)
    agent_id = doc.get("agent_id")
    source_type = doc.get("source_type")
    source_value = doc.get("source_value")
    if source_type is None or source_value is None:
        return None
    try:
        stype = SourceType(str(source_type))
    except Exception:
        try:
            stype = SourceType[str(source_type).upper()]
        except Exception:
            stype = SourceType.OTHER
    source = FactSource(stype, str(source_value), str(agent_id) if agent_id else None)
    return source.id
