import asyncio
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from datetime import datetime
from difflib import SequenceMatcher
from pathlib import Path
from typing import Self, TypedDict

import aiosqlite
from pydantic import BaseModel, ConfigDict

import ntrp.database as database
from ntrp.channel import Channel
from ntrp.constants import (
    FACT_DEDUP_EMBEDDING_SIMILARITY,
    FACT_DEDUP_TEXT_RATIO,
    FORGET_SEARCH_LIMIT,
    FORGET_SIMILARITY_THRESHOLD,
    RECALL_SEARCH_LIMIT,
    SYSTEM_PROMPT_OBSERVATION_LIMIT,
    USER_ENTITY_NAME,
)
from ntrp.embedder import Embedder, EmbeddingConfig
from ntrp.events.internal import FactCreated, FactDeleted
from ntrp.logging import get_logger
from ntrp.memory.consolidation_runner import ConsolidationRunner
from ntrp.memory.decay import decay_score
from ntrp.memory.extraction import Extractor
from ntrp.memory.models import ExtractionResult, Fact, FactContext, Observation
from ntrp.memory.retrieval import retrieve_with_observations
from ntrp.memory.store.base import GraphDatabase
from ntrp.memory.store.dreams import DreamRepository
from ntrp.memory.store.facts import FactRepository
from ntrp.memory.store.observations import ObservationRepository

_logger = get_logger(__name__)


class RememberFactResult(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    fact: Fact
    entities_extracted: list[str]


class ReembedProgress(TypedDict):
    total: int
    done: int


class FactMemory:
    def __init__(
        self,
        conn: aiosqlite.Connection,
        embedding: EmbeddingConfig,
        model: str,
        channel: Channel,
        embedder: Embedder | None = None,
        extractor: Extractor | None = None,
    ):
        self.db = GraphDatabase(conn, embedding.dim)
        self.facts = FactRepository(conn)
        self.observations = ObservationRepository(conn)
        self.dreams = DreamRepository(conn)
        self.embedder = embedder or Embedder(embedding)
        self.extractor = extractor or Extractor(model)
        self.channel = channel
        self._db_lock = asyncio.Lock()
        self._reembed_task: asyncio.Task | None = None
        self._reembed_progress: ReembedProgress | None = None

        self._consolidation = ConsolidationRunner(
            facts=self.facts,
            observations=self.observations,
            dreams=self.dreams,
            embedder=self.embedder,
            model_fn=lambda: self.model,
            publish=self.channel.publish,
            transaction=self.transaction,
            db_lock=self._db_lock,
            db_conn=conn,
        )

    @asynccontextmanager
    async def transaction(self) -> AsyncGenerator[None]:
        async with self._db_lock:
            try:
                yield
                await self.db.conn.commit()
            except Exception:
                await self.db.conn.rollback()
                raise

    @property
    def is_consolidating(self) -> bool:
        return self._consolidation.running

    @classmethod
    async def create(
        cls,
        db_path: Path,
        embedding: EmbeddingConfig,
        model: str,
        channel: Channel,
    ) -> Self:
        conn = await database.connect(db_path, vec=True)
        instance = cls(conn, embedding, model, channel=channel)
        await instance.db.init_schema()
        if instance.db.dim_changed:
            _logger.info("Embedding dimension changed — starting background re-embed")
            instance.start_reembed(embedding)
        return instance

    # --- Consolidation delegation ---

    @property
    def dreams_enabled(self) -> bool:
        return self._consolidation.dreams_enabled

    @dreams_enabled.setter
    def dreams_enabled(self, value: bool) -> None:
        self._consolidation.dreams_enabled = value

    async def run_consolidation(self) -> str:
        return await self._consolidation.run_once()

    # --- Re-embedding ---

    @property
    def reembed_running(self) -> bool:
        return self._reembed_task is not None and not self._reembed_task.done()

    @property
    def reembed_progress(self) -> ReembedProgress | None:
        return self._reembed_progress

    @property
    def model(self) -> str:
        return self.extractor.model

    def update_model(self, model: str) -> None:
        self.extractor.model = model

    def start_reembed(self, embedding: EmbeddingConfig, *, rebuild: bool = False) -> None:
        if self._reembed_task and not self._reembed_task.done():
            self._reembed_task.cancel()
        self._reembed_task = asyncio.create_task(self._run_reembed(embedding, rebuild=rebuild))

    async def _run_reembed(self, embedding: EmbeddingConfig, *, rebuild: bool = False, batch_size: int = 100) -> None:
        try:
            new_embedder = Embedder(embedding)
            if rebuild:
                await self.db.rebuild_vec_tables(embedding.dim)

            facts = await self.facts.list_all_with_embeddings()
            observations = await self.observations.list_all_with_embeddings()
            total = len(facts) + len(observations)
            self._reembed_progress = {"total": total, "done": 0}

            done = 0
            for i in range(0, len(facts), batch_size):
                batch = facts[i : i + batch_size]
                embeddings = await new_embedder.embed([f.text for f in batch])
                for fact, emb in zip(batch, embeddings):
                    await self.facts.update_embedding(fact.id, emb)
                done += len(batch)
                self._reembed_progress["done"] = done

            for i in range(0, len(observations), batch_size):
                batch = observations[i : i + batch_size]
                embeddings = await new_embedder.embed([o.summary for o in batch])
                for obs, emb in zip(batch, embeddings):
                    await self.observations.update_embedding(obs.id, emb)
                done += len(batch)
                self._reembed_progress["done"] = done

            await self.db.conn.commit()
            self.embedder = new_embedder
            _logger.info("Re-embedded %d memory vectors", total)
        except asyncio.CancelledError:
            raise
        except Exception:
            _logger.warning("Memory re-embed failed", exc_info=True)
        finally:
            self._reembed_progress = None

    # --- Lifecycle ---

    async def close(self) -> None:
        if self._reembed_task:
            self._reembed_task.cancel()
            try:
                await self._reembed_task
            except asyncio.CancelledError:
                pass
            self._reembed_task = None
        await self.db.conn.close()

    # --- Core API ---

    async def remember(
        self,
        text: str,
        source_type: str = "explicit",
        source_ref: str | None = None,
        happened_at: datetime | None = None,
    ) -> RememberFactResult | None:
        if not text or not text.strip():
            return None

        embedding, extraction = await asyncio.gather(
            self.embedder.embed_one(text),
            self.extractor.extract(text),
        )

        async with self.transaction():
            if similar := await self.facts.search_facts_vector(embedding, limit=1):
                existing_fact, similarity = similar[0]
                text_ratio = SequenceMatcher(None, text.lower(), existing_fact.text.lower()).ratio()
                is_dup = text_ratio >= FACT_DEDUP_TEXT_RATIO or similarity >= FACT_DEDUP_EMBEDDING_SIMILARITY
                _logger.info(
                    "Dedup check: fact %d text_ratio=%.3f sim=%.3f dup=%s — %r",
                    existing_fact.id,
                    text_ratio,
                    similarity,
                    is_dup,
                    existing_fact.text[:80],
                )
                if is_dup:
                    await self.facts.reinforce([existing_fact.id])
                    return None

            fact = await self.facts.create(
                text=text,
                source_type=source_type,
                source_ref=source_ref,
                embedding=embedding,
                happened_at=happened_at,
            )
            entities_extracted = await self._process_extraction(fact.id, extraction)

        self.channel.publish(FactCreated(fact_id=fact.id, text=text))
        return RememberFactResult(fact=fact, entities_extracted=entities_extracted)

    async def _process_extraction(self, fact_id: int, extraction: ExtractionResult) -> list[str]:
        seen: set[str] = set()
        for entity in extraction.entities:
            name = entity.name.strip()
            if not name or name.lower() in seen:
                continue
            seen.add(name.lower())
            entity_id = await self._resolve_entity(name)
            await self.facts.add_entity_ref(fact_id, name, entity_id)
        return list(seen)

    async def _resolve_entity(self, name: str) -> int | None:
        if existing := await self.facts.get_entity_by_name(name):
            return existing.id
        entity = await self.facts.create_entity(name=name)
        return entity.id

    async def recall(
        self,
        query: str,
        limit: int = RECALL_SEARCH_LIMIT,
        query_time: datetime | None = None,
    ) -> FactContext:
        query_embedding = await self.embedder.embed_one(query)

        context = await retrieve_with_observations(
            self.facts,
            self.observations,
            query,
            query_embedding,
            seed_limit=limit,
            query_time=query_time,
        )

        async with self.transaction():
            if context.facts:
                await self.facts.reinforce([f.id for f in context.facts])
            if context.observations:
                await self.observations.reinforce([o.id for o in context.observations])
                displayed_fact_ids = [f.id for facts in context.bundled_sources.values() for f in facts]
                if displayed_fact_ids:
                    await self.facts.reinforce(displayed_fact_ids)

        return context

    async def forget(self, query: str) -> int:
        query_embedding = await self.embedder.embed_one(query)
        results = await self.facts.search_facts_vector(query_embedding, limit=FORGET_SEARCH_LIMIT)

        async with self.transaction():
            count = 0
            deleted_ids = []
            for fact, score in results:
                if score >= FORGET_SIMILARITY_THRESHOLD:
                    await self.facts.delete(fact.id)
                    deleted_ids.append(fact.id)
                    count += 1
                    self.channel.publish(FactDeleted(fact_id=fact.id))
            if count > 0:
                await self.observations.remove_source_facts(deleted_ids)
                await self.dreams.remove_source_facts(deleted_ids)
                await self.facts.cleanup_orphaned_entities()
            return count

    async def merge_entities(self, names: list[str], canonical_name: str | None = None) -> int:
        if len(names) < 2:
            return 0

        async with self.transaction():
            entities = []
            for name in names:
                if entity := await self.facts.get_entity_by_name(name):
                    entities.append(entity)

            if len(entities) < 2:
                return 0

            if canonical_name:
                keep = next((e for e in entities if e.name.lower() == canonical_name.lower()), entities[0])
            else:
                keep = entities[0]

            merge_ids = [e.id for e in entities if e.id != keep.id]
            count = await self.facts.merge_entities(keep.id, merge_ids)
            await self.observations.merge_entity_refs(keep.id, merge_ids)
            _logger.info("Merged entities %s → '%s' (%d merged)", [e.name for e in entities], keep.name, count)
            return count

    async def count(self) -> int:
        return await self.facts.count()

    async def get_context(self, user_limit: int = 10) -> tuple[list[Observation], list[Fact]]:
        if user_entity := await self.facts.get_entity_by_name(USER_ENTITY_NAME):
            raw_obs = await self.observations.get_for_entity(user_entity.id, limit=20)
            scored_obs = sorted(
                raw_obs,
                key=lambda o: decay_score(o.last_accessed_at, o.access_count),
                reverse=True,
            )
            observations = scored_obs[:SYSTEM_PROMPT_OBSERVATION_LIMIT]
        else:
            observations = []

        exclude_ids: set[int] = set()
        for obs in observations:
            exclude_ids.update(obs.source_fact_ids)

        all_user_facts = await self.facts.get_facts_for_entity(USER_ENTITY_NAME, limit=user_limit + len(exclude_ids))
        user_facts = [f for f in all_user_facts if f.id not in exclude_ids][:user_limit]

        return observations, user_facts

    async def clear_observations(self) -> dict[str, int]:
        async with self.transaction():
            obs_count = await self.observations.clear_all()
            facts_reset = await self.facts.reset_consolidated()
            return {"observations_deleted": obs_count, "facts_reset": facts_reset}

    async def clear(self) -> dict[str, int]:
        async with self.transaction():
            counts = {
                "facts": await self.facts.count(),
                "observations": await self.observations.count(),
            }
            await self.db.clear_all()
            return counts
