"""MemoryClient — the main entry point for the memory SDK.

Provides write(), retrieve(), and CRUD operations (get, get_all, delete, delete_all).
All dependencies are injected via constructor (LLM, Embeddings, Config).
"""

from __future__ import annotations

import logging
import uuid as uuid_mod
from dataclasses import dataclass, field
from dataclasses import fields as dc_fields
from datetime import datetime
from typing import Any

from sqlalchemy import delete as sa_delete
from sqlalchemy import select

from arandu.config import MemoryConfig
from arandu.db import create_engine, create_session_factory, init_db
from arandu.protocols import EmbeddingProvider, LLMProvider, LLMResult, TokenUsage
from arandu.tracing import PipelineTrace

logger = logging.getLogger(__name__)


@dataclass
class WriteResult:
    """Result of a memory write operation."""

    event_id: str = ""
    facts_added: list = field(default_factory=list)
    facts_updated: list = field(default_factory=list)
    facts_unchanged: list = field(default_factory=list)
    facts_deleted: list = field(default_factory=list)
    entities_resolved: list = field(default_factory=list)
    duration_ms: float = 0.0
    pipeline: PipelineTrace | None = None
    tokens_used: TokenUsage = field(default_factory=TokenUsage)
    success: bool = True
    error: str | None = None


@dataclass
class FactDetail:
    """Detailed view of a stored fact — used by CRUD operations (get, get_all).

    Unlike ``ScoredFact`` (which carries retrieval scores), ``FactDetail``
    exposes the full metadata of a fact as stored in the database.
    """

    fact_id: str = ""
    entity_name: str = ""
    entity_key: str = ""
    entity_type: str = ""
    attribute_key: str | None = None
    fact_text: str = ""
    category: str | None = None
    confidence: float = 0.0
    importance: float = 0.0
    valid_from: datetime | None = None
    created_at: datetime | None = None
    source_context: str | None = None
    speaker: str | None = None


@dataclass
class EntityDetail:
    """Detailed view of a known entity — used by the entities() method."""

    entity_id: str = ""
    canonical_key: str = ""
    display_name: str = ""
    entity_type: str = ""
    summary_text: str | None = None
    fact_count: int = 0
    importance_score: float | None = None
    first_seen_at: datetime | None = None
    last_seen_at: datetime | None = None
    profile_text: str | None = None


@dataclass
class ScoredFact:
    """A fact with retrieval scores.

    The ``scores`` dict contains the breakdown of individual scoring signals.
    Possible keys (not all are always present):

    - ``semantic`` — cosine similarity from pgvector (0.0–1.0)
    - ``keyword`` — fraction of query words found in fact text (0.0–1.0)
    - ``recency`` — exponential decay based on fact age (configurable half-life)
    - ``importance`` — base importance value of the fact (typically 0.5)
    - ``confidence`` — effective confidence with temporal decay
    - ``pattern`` — additive boost for reinforced facts (up to 0.1)
    - ``graph`` — score from BFS 2-hop knowledge graph traversal
    - ``reranker`` — LLM-based relevance score (0.0–1.0, only when reranker is enabled)
    """

    fact_id: str = ""
    entity_name: str = ""
    attribute_key: str = ""
    fact_text: str = ""
    score: float = 0.0
    scores: dict = field(default_factory=dict)
    speaker: str | None = None


@dataclass
class RetrieveResult:
    """Result of a memory retrieve operation.

    The ``warnings`` field contains user-facing strings for any ``entity_keys``
    the caller passed that did not resolve to a canonical entity or alias.
    When the filter matches fully, ``warnings`` is empty.
    """

    facts: list[ScoredFact] = field(default_factory=list)
    context: str = ""
    total_candidates: int = 0
    duration_ms: float = 0.0
    pipeline: PipelineTrace | None = None
    tokens_used: TokenUsage = field(default_factory=TokenUsage)
    config_effective: dict = field(default_factory=dict)
    warnings: list[str] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Config override helpers
# ---------------------------------------------------------------------------


def _apply_config_overrides(
    base: MemoryConfig,
    overrides: dict[str, Any],
) -> MemoryConfig:
    """Merge overrides into a base MemoryConfig, with validation.

    - Unknown keys emit a warning and are ignored.
    - Type mismatches for scalar fields (int, float, bool, str) raise ValueError.
    - Valid keys are merged; unset keys inherit from base.
    """
    valid_fields = {f.name for f in dc_fields(MemoryConfig)}
    merged = {f.name: getattr(base, f.name) for f in dc_fields(MemoryConfig)}

    for k, v in overrides.items():
        if k not in valid_fields:
            logger.warning("config_overrides: unknown field '%s' ignored", k)
            continue

        # Type check for scalar fields
        expected = type(getattr(base, k))
        if expected in (int, float, bool, str) and not isinstance(v, expected):
            # Allow int → float promotion
            if expected is float and isinstance(v, int):
                v = float(v)
            else:
                raise ValueError(
                    f"config_overrides: field '{k}' expects {expected.__name__}, got {type(v).__name__}: {v!r}"
                )

        merged[k] = v

    return MemoryConfig(**merged)


def _serialize_config(config: MemoryConfig) -> dict[str, Any]:
    """Serialize a MemoryConfig to a JSON-friendly dict (all fields)."""
    result: dict[str, Any] = {}
    for f in dc_fields(MemoryConfig):
        val = getattr(config, f.name)
        if isinstance(val, set):
            result[f.name] = sorted(val)
        else:
            result[f.name] = val
    return result


class _TokenTrackingLLM:
    """Thin wrapper around an LLMProvider that accumulates token usage."""

    def __init__(self, inner: LLMProvider) -> None:
        self._inner = inner
        self.total_usage = TokenUsage()

    async def complete(
        self,
        messages: list[dict],
        temperature: float = 0,
        response_format: dict | None = None,
        max_tokens: int | None = None,
    ) -> LLMResult:
        result = await self._inner.complete(
            messages=messages,
            temperature=temperature,
            response_format=response_format,
            max_tokens=max_tokens,
        )
        if result.usage:
            self.total_usage = self.total_usage + result.usage
        return result


def _fact_to_detail(fact: Any) -> FactDetail:
    """Convert a MemoryFact ORM instance to a FactDetail dataclass."""
    return FactDetail(
        fact_id=str(fact.id),
        entity_name=fact.entity_name or "",
        entity_key=fact.entity_key,
        entity_type=fact.entity_type,
        attribute_key=fact.attribute_key,
        fact_text=fact.fact_text,
        category=fact.category,
        confidence=fact.confidence,
        importance=fact.importance,
        valid_from=fact.valid_from,
        created_at=fact.created_at,
        source_context=fact.source_context,
        speaker=fact.speaker,
    )


def _entity_to_detail(entity: Any) -> EntityDetail:
    """Convert a MemoryEntity ORM instance to an EntityDetail dataclass."""
    return EntityDetail(
        entity_id=str(entity.id),
        canonical_key=entity.canonical_key,
        display_name=entity.display_name or "",
        entity_type=entity.entity_type,
        summary_text=entity.summary_text,
        fact_count=entity.fact_count,
        importance_score=entity.importance_score,
        first_seen_at=entity.first_seen_at,
        last_seen_at=entity.last_seen_at,
        profile_text=entity.profile_text,
    )


class MemoryClient:
    """Main entry point for the arandu SDK.

    Usage::

        from arandu import MemoryClient, MemoryConfig
        from arandu.providers.openai import OpenAIProvider

        provider = OpenAIProvider(api_key="sk-...")
        memory = MemoryClient(
            database_url="postgresql+psycopg://user:pass@localhost/mydb",
            llm=provider,
            embeddings=provider,
        )
        await memory.initialize()

        result = await memory.write(agent_id="my_agent", message="I live in São Paulo", speaker_name="Pedro")
        context = await memory.retrieve(agent_id="my_agent", query="where does Pedro live?")
    """

    def __init__(
        self,
        database_url: str,
        llm: LLMProvider,
        embeddings: EmbeddingProvider,
        config: MemoryConfig | None = None,
    ) -> None:
        self._config = config or MemoryConfig()
        self._llm = llm
        self._embeddings = embeddings
        self._engine = create_engine(database_url)
        self._session_factory = create_session_factory(self._engine)
        self._initialized = False

    async def initialize(self) -> None:
        """Create memory tables in the database (idempotent)."""
        await init_db(self._engine)
        self._initialized = True

    @property
    def config(self) -> MemoryConfig:
        return self._config

    async def write(
        self,
        agent_id: str,
        message: str,
        speaker_name: str,
        source: str = "api",
        *,
        session_id: str = "default",
        config_overrides: dict | None = None,
        dry_run: bool = False,
        recent_messages: list[str] | None = None,
        verbose: bool = False,
        occurred_at: datetime | None = None,
    ) -> WriteResult:
        """Extract facts from a message and store them in memory.

        Args:
            agent_id: Identifier for the agent whose memory this is. All facts,
                entities, and relationships are scoped to this agent. Use the same
                agent_id to share memory across sessions.
            message: The message text to extract knowledge from.
            speaker_name: Name of the speaker (the person talking). Required.
                This identifies WHO is speaking — the person whose words are being
                memorized. First-person pronouns ("I", "me", "my") resolve to
                this name. Example: ``speaker_name="Pedro Menezes"``.
            source: Source channel identifier (default: "api").
            session_id: Identifier for the conversation session (default: "default").
                Use different session_ids to separate conversation contexts.
                The agent's memory persists across sessions, but session_id helps
                track which conversation each event came from.
            config_overrides: Optional dict of MemoryConfig field overrides for this request.
                Only the provided keys are changed; all others inherit from the client config.
                Example: ``{"extraction_timeout_sec": 60.0}``
            dry_run: If True, run extraction + resolution + reconciliation but skip
                the upsert stage. No data is persisted. Useful for benchmarking.
            recent_messages: Optional conversation context (last N messages)
                for resolving pronouns and anaphora references in extraction.
            verbose: If True, include pipeline trace with intermediate data.
            occurred_at: Optional timestamp for when this message was sent.
                Defaults to now. Use for importing historical conversations
                so the extraction can resolve relative dates correctly.

        Returns:
            WriteResult with details of extracted and stored facts.

        Raises:
            ValueError: If speaker_name is empty or whitespace.
        """
        if not speaker_name or not speaker_name.strip():
            raise ValueError("speaker_name is required and cannot be empty")

        from arandu.write.pipeline import run_write_pipeline

        effective_config = self._config
        if config_overrides:
            effective_config = _apply_config_overrides(self._config, config_overrides)

        trace = PipelineTrace() if verbose else None
        tracked_llm = _TokenTrackingLLM(self._llm)

        async with self._session_factory() as session:
            result = await run_write_pipeline(
                session=session,
                agent_id=str(agent_id),
                message=message,
                llm=tracked_llm,
                embeddings=self._embeddings,
                config=effective_config,
                speaker_name=speaker_name.strip(),
                source=source,
                session_id=session_id,
                recent_messages=recent_messages,
                trace=trace,
                dry_run=dry_run,
                occurred_at=occurred_at,
            )
            if not dry_run:
                await session.commit()

        return WriteResult(
            event_id=result.event_id,
            facts_added=result.facts_added,
            facts_updated=result.facts_updated,
            facts_unchanged=result.facts_unchanged,
            facts_deleted=result.facts_deleted,
            entities_resolved=result.entities_resolved,
            duration_ms=result.duration_ms,
            pipeline=trace,
            tokens_used=tracked_llm.total_usage,
            success=result.success,
            error=result.error,
        )

    async def retrieve(
        self,
        agent_id: str,
        query: str,
        *,
        session_id: str | None = None,
        config_overrides: dict | None = None,
        verbose: bool = False,
        entity_keys: list[str] | None = None,
    ) -> RetrieveResult:
        """Retrieve relevant memory context for a query.

        Args:
            agent_id: Identifier for the agent whose memory to search.
            query: The query to search memory for.
            session_id: Optional session filter. If None (default), searches
                across all sessions. If provided, may be used for context
                prioritization in future versions.
            config_overrides: Optional dict of MemoryConfig field overrides for this request.
                Only the provided keys are changed; all others inherit from the client config.
                Example: ``{"enable_reranker": False, "topk_facts": 10}``
            verbose: If True, include pipeline trace with intermediate data.
            entity_keys: Optional list of entity keys to filter by (OR logic).
                When provided, only facts linked to at least one of the specified
                entities are returned.  Example: ``["person:pedro", "person:ana"]``

        Returns:
            RetrieveResult with ranked facts and formatted context string.
        """
        from arandu.read.pipeline import run_read_pipeline

        effective_config = self._config
        if config_overrides:
            effective_config = _apply_config_overrides(self._config, config_overrides)

        trace = PipelineTrace() if verbose else None
        tracked_llm = _TokenTrackingLLM(self._llm)

        async with self._session_factory() as session:
            result = await run_read_pipeline(
                session=session,
                agent_id=str(agent_id),
                query=query,
                llm=tracked_llm,
                embeddings=self._embeddings,
                config=effective_config,
                trace=trace,
                entity_keys=entity_keys,
            )

        config_dict = _serialize_config(effective_config)

        return RetrieveResult(
            facts=[
                ScoredFact(
                    fact_id=f.fact_id,
                    entity_name=f.entity_name,
                    attribute_key=f.attribute_key,
                    fact_text=f.fact_text,
                    score=f.score,
                    scores=f.scores,
                    speaker=f.speaker,
                )
                for f in result.facts
            ],
            context=result.context,
            total_candidates=result.total_candidates,
            duration_ms=result.duration_ms,
            pipeline=trace,
            tokens_used=tracked_llm.total_usage,
            config_effective=config_dict,
            warnings=list(result.warnings),
        )

    # ------------------------------------------------------------------
    # CRUD operations
    # ------------------------------------------------------------------

    async def get(self, agent_id: str, fact_id: str) -> FactDetail | None:
        """Return a single fact by ID, or ``None`` if not found.

        Fetches any fact belonging to *agent_id* regardless of its
        ``valid_to`` status — even facts that have been soft-deleted by
        the reconciliation pipeline are returned when looked up by ID.

        Args:
            agent_id: Agent that owns the fact.
            fact_id: UUID of the fact to retrieve.

        Returns:
            A ``FactDetail`` with the fact's metadata, or ``None``
            if no matching fact exists.
        """
        from arandu.models import MemoryFact

        try:
            parsed_id = uuid_mod.UUID(fact_id)
        except ValueError:
            return None

        async with self._session_factory() as session:
            stmt = select(MemoryFact).where(
                MemoryFact.id == parsed_id,
                MemoryFact.agent_id == str(agent_id),
            )
            result = await session.execute(stmt)
            fact = result.scalar_one_or_none()

        if fact is None:
            return None

        return _fact_to_detail(fact)

    async def get_all(
        self,
        agent_id: str,
        *,
        limit: int = 100,
        offset: int = 0,
        entity_keys: list[str] | None = None,
    ) -> list[FactDetail]:
        """List active facts for an agent with pagination.

        Only facts with ``valid_to IS NULL`` (i.e. not soft-deleted by the
        reconciliation pipeline) are returned.  Results are ordered by
        ``created_at`` descending (newest first).

        Args:
            agent_id: Agent whose facts to list.
            limit: Maximum number of facts to return (default 100).
            offset: Number of facts to skip (default 0).
            entity_keys: Optional list of entity keys to filter by (OR logic).
                When provided, only facts linked to at least one of the specified
                entities are returned.  Aliases registered in
                ``memory_entity_aliases`` are resolved automatically to their
                canonical key before filtering. Keys that resolve to nothing
                are silently dropped from the filter (the CRUD path does not
                return warnings — use ``retrieve()`` for that).

        Returns:
            A list of ``FactDetail`` objects.
        """
        from arandu._fact_query import FactQuery
        from arandu.read.retrieval import _resolve_entity_keys

        async with self._session_factory() as session:
            query = FactQuery.active(str(agent_id)).order_by_created().limit(limit).offset(offset)

            if entity_keys is not None:
                resolved_entity_keys, _invalid = await _resolve_entity_keys(
                    session,
                    str(agent_id),
                    entity_keys,
                )
                # Apply the resolved list (may be empty → filter matches nothing,
                # consistent with the existing "all invalid" behavior).
                query = query.for_entities(resolved_entity_keys)

            result = await session.execute(query.build())
            facts = list(result.scalars().all())

        return [_fact_to_detail(f) for f in facts]

    async def delete(self, agent_id: str, fact_id: str) -> bool:
        """Hard-delete a single fact by ID.

        The row is physically removed from the database.  Associated
        ``MemoryFactEntityLink`` rows are removed automatically via the
        ``CASCADE`` foreign-key constraint.

        Args:
            agent_id: Agent that owns the fact.
            fact_id: UUID of the fact to delete.

        Returns:
            ``True`` if a fact was deleted, ``False`` if no matching
            fact was found.
        """
        from arandu.models import MemoryFact

        try:
            parsed_id = uuid_mod.UUID(fact_id)
        except ValueError:
            return False

        async with self._session_factory() as session:
            stmt = sa_delete(MemoryFact).where(
                MemoryFact.id == parsed_id,
                MemoryFact.agent_id == str(agent_id),
            )
            cursor = await session.execute(stmt)
            await session.commit()
            row_count: int = cursor.rowcount  # type: ignore[attr-defined]

        return row_count > 0

    async def delete_all(self, agent_id: str) -> int:
        """Hard-delete **all** memory for an agent.

        Removes all facts, entities, aliases, and relationships belonging
        to *agent_id*. ``MemoryFactEntityLink`` rows are removed
        automatically via the ``CASCADE`` foreign-key constraint on
        ``MemoryFact``.  ``MemoryEvent`` rows are preserved (immutable
        audit log).

        Args:
            agent_id: Agent whose memory to delete.

        Returns:
            The number of facts deleted.
        """
        from arandu.models import (
            MemoryEntity,
            MemoryEntityAlias,
            MemoryEntityRelationship,
            MemoryFact,
        )

        async with self._session_factory() as session:
            # Delete facts first (entity links cascade)
            stmt_facts = sa_delete(MemoryFact).where(MemoryFact.agent_id == str(agent_id))
            cursor = await session.execute(stmt_facts)
            fact_count: int = cursor.rowcount  # type: ignore[attr-defined]

            # Delete relationships, aliases, and entities
            await session.execute(
                sa_delete(MemoryEntityRelationship).where(MemoryEntityRelationship.agent_id == str(agent_id))
            )
            await session.execute(sa_delete(MemoryEntityAlias).where(MemoryEntityAlias.agent_id == str(agent_id)))
            await session.execute(sa_delete(MemoryEntity).where(MemoryEntity.agent_id == str(agent_id)))

            await session.commit()

        return fact_count

    async def entities(
        self,
        agent_id: str,
        *,
        limit: int = 100,
        offset: int = 0,
    ) -> list[EntityDetail]:
        """List known entities for an agent.

        Returns active entities (``is_active = True``) ordered by
        ``last_seen_at`` descending (most recently seen first), with
        ``id`` as tiebreaker for stable pagination.

        Args:
            agent_id: Agent whose entities to list.
            limit: Maximum number of entities to return (default 100).
            offset: Number of entities to skip (default 0).

        Returns:
            A list of ``EntityDetail`` objects.
        """
        from arandu.models import MemoryEntity

        async with self._session_factory() as session:
            stmt = (
                select(MemoryEntity)
                .where(
                    MemoryEntity.agent_id == str(agent_id),
                    MemoryEntity.is_active.is_(True),
                )
                .order_by(MemoryEntity.last_seen_at.desc(), MemoryEntity.id)
                .limit(limit)
                .offset(offset)
            )
            result = await session.execute(stmt)
            entities = list(result.scalars().all())

        return [_entity_to_detail(e) for e in entities]

    async def close(self) -> None:
        """Dispose of the database engine and release connections."""
        await self._engine.dispose()
