from __future__ import annotations

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

from analogai.foundation.settings import config as infra_config
from analogai.deepthink.messaging.fact.fact import Fact
from analogai.deepthink.messaging.fact.enums.fact_reference_type import FactReferenceType
from analogai.deepthink.messaging.fact.repositories.fact_repository import FactRepositoryPort
from analogai.deepthink.messaging.fact.repositories.infrastructure.fact_storage_mapper import FactStorageMapper


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

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


class InMemoryFactRepository(FactRepositoryPort):
    def __init__(self) -> None:
        self._facts: Dict[str, Fact] = {}
        self._facts_by_turn: Dict[str, List[str]] = {}
        self._turn_index: Dict[Tuple[str, str], List[str]] = {}
        self._statement_links: Dict[str, str] = {}
        self._belief_links: Dict[str, str] = {}
        self._notion_links: Dict[str, str] = {}
        self._hypothetical_links: Dict[str, str] = {}
        self._resolver = _RepositoryFactResolver(self)

    def register_turn(self, user_id: str, agent_id: str, turn_id: str) -> None:
        self._turn_index.setdefault((user_id, agent_id), []).append(turn_id)

    def save_fact(self, fact: Fact) -> None:
        self._facts[fact.id] = fact
        if fact.turn_id:
            self._facts_by_turn.setdefault(fact.turn_id, []).append(fact.id)

    def get_fact(self, id: str) -> Optional[Fact]:
        fact = self._facts.get(id)
        return self._attach_resolver(fact) if fact else None

    def update_facts_statement_id(self, turn_id: str, statement_id: str) -> None:
        ids = self._facts_by_turn.get(turn_id, [])
        for fact_key in ids:
            fact = self._facts.get(fact_key)
            if fact is None:
                continue
            self._facts[fact_key] = Fact(
                id=fact.id,
                turn_id=fact.turn_id,
                statement_id=statement_id,
                notion_id=fact.notion_id,
                reference_type=FactReferenceType.STATEMENT.value,
                user_id=fact.user_id,
                agent_id=fact.agent_id,
                text=fact.text,
                embedding=fact.embedding,
                prev_id=fact.prev_id,
                next_id=fact.next_id,
            )

    def update_facts_notion_id(self, turn_id: str, notion_id: str) -> None:
        ids = self._facts_by_turn.get(turn_id, [])
        for fact_key in ids:
            fact = self._facts.get(fact_key)
            if fact is None:
                continue
            self._facts[fact_key] = Fact(
                id=fact.id,
                turn_id=fact.turn_id,
                statement_id=fact.statement_id,
                notion_id=notion_id,
                reference_type=FactReferenceType.NOTION.value,
                user_id=fact.user_id,
                agent_id=fact.agent_id,
                text=fact.text,
                embedding=fact.embedding,
                prev_id=fact.prev_id,
                next_id=fact.next_id,
            )

    def update_facts_hypothetical_statement_id(self, turn_id: str, statement_id: str) -> None:
        ids = self._facts_by_turn.get(turn_id, [])
        for fact_key in ids:
            fact = self._facts.get(fact_key)
            if fact is None:
                continue
            self._facts[fact_key] = Fact(
                id=fact.id,
                turn_id=fact.turn_id,
                statement_id=statement_id,
                notion_id=fact.notion_id,
                reference_type=FactReferenceType.HYPOTHETICAL_STATEMENT.value,
                user_id=fact.user_id,
                agent_id=fact.agent_id,
                text=fact.text,
                embedding=fact.embedding,
                prev_id=fact.prev_id,
                next_id=fact.next_id,
            )

    def get_facts_for_turn(self, turn_id: str) -> List[Fact]:
        return [
            self._attach_resolver(self._facts[cid])
            for cid in self._facts_by_turn.get(turn_id, [])
            if cid in self._facts
        ]

    def get_turn_ids(self, user_id: str, agent_id: str) -> List[str]:
        return list(self._turn_index.get((user_id, agent_id), []))

    def link_statement(self, statement_id: str, turn_id: str) -> None:
        if statement_id and turn_id:
            self._statement_links[statement_id] = turn_id

    def link_belief(self, belief_id: str, turn_id: str) -> None:
        if belief_id and turn_id:
            self._belief_links[belief_id] = turn_id

    def link_notion(self, notion_id: str, turn_id: str) -> None:
        if notion_id and turn_id:
            self._notion_links[notion_id] = turn_id

    def link_hypothetical_statement(self, statement_id: str, turn_id: str) -> None:
        if statement_id and turn_id:
            self._hypothetical_links[statement_id] = turn_id

    def get_turn_id_for_statement(self, statement_id: str) -> Optional[str]:
        return self._statement_links.get(statement_id)

    def get_turn_id_for_belief(self, belief_id: str) -> Optional[str]:
        return self._belief_links.get(belief_id)

    def get_turn_id_for_notion(self, notion_id: str) -> Optional[str]:
        return self._notion_links.get(notion_id)

    def get_turn_id_for_hypothetical_statement(self, statement_id: str) -> Optional[str]:
        return self._hypothetical_links.get(statement_id)

    def search_text(self, user_id: str, agent_id: str, query: str, top_k: int) -> List[Fact]:
        query_lower = query.lower()
        matches: List[Fact] = []
        for fact in self._iter_user_facts(user_id, agent_id):
            if query_lower in (fact.text or "").lower():
                matches.append(self._attach_resolver(fact))
        return matches[: max(1, int(top_k))]

    def search_similar(
        self,
        user_id: str,
        agent_id: str,
        embedding: List[float],
        top_k: int,
        similarity_cutoff: Optional[float] = None,
    ) -> List[Fact]:
        scored: List[Tuple[float, Fact]] = []
        for fact in self._iter_user_facts(user_id, agent_id):
            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, self._attach_resolver(fact)))
        scored.sort(key=operator.itemgetter(0), reverse=True)
        return [fact for _, fact in scored[: max(1, int(top_k))]]

    def _iter_user_facts(self, user_id: str, agent_id: str) -> Iterable[Fact]:
        for fact in self._facts.values():
            if fact.user_id == user_id and fact.agent_id == agent_id:
                yield fact

    def _attach_resolver(self, fact: Fact) -> Fact:
        fact._resolver = self._resolver
        return fact


class MongoFactRepository(FactRepositoryPort):
    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)

    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: Fact) -> None:
        doc = FactStorageMapper.to_document(fact)
        self._facts.update_one({"_id": fact.id}, {"$set": doc}, upsert=True)

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

    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[Fact]:
        cursor = self._facts.find({"turn_id": turn_id})
        results: List[Fact] = []
        for doc in cursor:
            fact = FactStorageMapper.from_document(doc)
            if fact is None:
                continue
            results.append(self._attach_resolver(fact))
        return results

    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 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[Fact]:
        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[Fact] = []
        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[Fact]:
        cursor = self._facts.find(
            {
                "user_id": user_id,
                "agent_id": agent_id,
                "embedding": {"$ne": None},
            }
        )
        scored: List[Tuple[float, Fact]] = []
        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 _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: Fact) -> Fact:
        fact._resolver = self._resolver
        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)
