from __future__ import annotations

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

from analogai.foundation.modules.fact.basefact import BaseFact
from analogai.foundation.modules.fact.fact_builder import build_fact
from analogai.foundation.modules.fact.enums.fact_reference_type import FactReferenceType
from analogai.foundation.modules.fact.repositories.fact_repository import FactRepository


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 InMemoryFactRepository(FactRepository):
    def __init__(self) -> None:
        self._facts: Dict[str, BaseFact] = {}
        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: BaseFact) -> 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[BaseFact]:
        fact = self._facts.get(id)
        return self._attach_resolver(fact) if fact else None

    def delete_fact(self, id: str) -> Optional[BaseFact]:
        if not id:
            return None
        fact = self._facts.pop(id, None)
        if fact is None:
            return None
        if fact.turn_id:
            remaining = [fid for fid in self._facts_by_turn.get(fact.turn_id, []) if fid != id]
            if remaining:
                self._facts_by_turn[fact.turn_id] = remaining
            else:
                self._facts_by_turn.pop(fact.turn_id, None)
        return self._attach_resolver(fact)

    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] = build_fact(
                id=fact.id,
                turn_id=fact.turn_id,
                statement_id=statement_id,
                notion_id=getattr(fact, "notion_id", None),
                reference_type=FactReferenceType.STATEMENT.value,
                user_id=fact.user_id,
                agent_id=fact.agent_id,
                text=fact.text,
                embedding=fact.embedding,
                source_id=fact.source_id,
                prev_id=fact.prev_id,
                next_id=fact.next_id,
                created_at=fact.created_at,
            )

    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] = build_fact(
                id=fact.id,
                turn_id=fact.turn_id,
                statement_id=getattr(fact, "statement_id", None),
                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,
                source_id=fact.source_id,
                prev_id=fact.prev_id,
                next_id=fact.next_id,
                created_at=fact.created_at,
            )

    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] = build_fact(
                id=fact.id,
                turn_id=fact.turn_id,
                statement_id=statement_id,
                notion_id=getattr(fact, "notion_id", None),
                reference_type=FactReferenceType.HYPOTHETICAL_STATEMENT.value,
                user_id=fact.user_id,
                agent_id=fact.agent_id,
                text=fact.text,
                embedding=fact.embedding,
                source_id=fact.source_id,
                prev_id=fact.prev_id,
                next_id=fact.next_id,
                created_at=fact.created_at,
            )

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

    def get_facts_by_source(self, source_id: str) -> List[BaseFact]:
        if not source_id:
            return []
        return [
            self._attach_resolver(fact)
            for fact in self._facts.values()
            if fact.source_id == source_id
        ]

    def list_source_ids(self) -> List[str]:
        sources: Dict[str, str] = {}
        for fact in self._facts.values():
            if not fact.source_id:
                continue
            sources.setdefault(fact.source_id, fact.source_id)
        return list(sources.keys())

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

    def get_statement_ids_for_turn(self, turn_id: str) -> List[str]:
        if not turn_id:
            return []
        return [key for key, tid in self._statement_links.items() if tid == turn_id]

    def get_belief_ids_for_turn(self, turn_id: str) -> List[str]:
        if not turn_id:
            return []
        return [key for key, tid in self._belief_links.items() if tid == turn_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[BaseFact]:
        query_lower = query.lower()
        matches: List[BaseFact] = []
        for fact in self._iter_user_facts(user_id, agent_id):
            text_value = str(getattr(fact, "text", "") or "")
            if query_lower in text_value.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[BaseFact]:
        scored: List[Tuple[float, BaseFact]] = []
        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 delete_facts_by_source(self, source_id: str) -> List[BaseFact]:
        matches = [fact for fact in self._facts.values() if fact.source_id == source_id]
        if not matches:
            return []
        deleted: List[BaseFact] = []
        fact_ids = set()
        turn_ids = set()
        for fact in matches:
            fact_ids.add(fact.id)
            if fact.turn_id:
                turn_ids.add(fact.turn_id)
            deleted.append(self._attach_resolver(fact))
        for fact_id in fact_ids:
            self._facts.pop(fact_id, None)
        for turn_id in turn_ids:
            remaining = [fid for fid in self._facts_by_turn.get(turn_id, []) if fid not in fact_ids]
            if remaining:
                self._facts_by_turn[turn_id] = remaining
            else:
                self._facts_by_turn.pop(turn_id, None)
        return deleted

    def _iter_user_facts(self, user_id: str, agent_id: str) -> Iterable[BaseFact]:
        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: 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)


