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

Provides write() and retrieve() as the two primary operations.
All dependencies are injected via constructor (LLM, Embeddings, Config).
"""

from __future__ import annotations

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

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 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)


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

    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)


# ---------------------------------------------------------------------------
# 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


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(user_id="user_123", message="I live in São Paulo")
        context = await memory.retrieve(user_id="user_123", query="where does the user live?")
    """

    def __init__(
        self,
        database_url: str,
        llm: LLMProvider,
        embeddings: EmbeddingProvider,
        config: MemoryConfig | None = None,
        reranker_llm: LLMProvider | None = None,
    ) -> None:
        self._config = config or MemoryConfig()
        self._llm = llm
        self._reranker_llm = reranker_llm or 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,
        user_id: str,
        message: str,
        source: str = "api",
        *,
        config_overrides: dict | None = None,
        dry_run: bool = False,
        recent_messages: list[str] | None = None,
        verbose: bool = False,
    ) -> WriteResult:
        """Extract facts from a message and store them in memory.

        Args:
            user_id: Unique identifier for the user (any string: UUID, email, numeric, etc.).
            message: The user's message text.
            source: Source channel identifier (default: "api").
            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.

        Returns:
            WriteResult with details of extracted and stored facts.
        """
        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,
                user_id=str(user_id),
                message=message,
                llm=tracked_llm,
                embeddings=self._embeddings,
                config=effective_config,
                source=source,
                recent_messages=recent_messages,
                trace=trace,
                dry_run=dry_run,
            )
            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,
        user_id: str,
        query: str,
        *,
        config_overrides: dict | None = None,
        verbose: bool = False,
    ) -> RetrieveResult:
        """Retrieve relevant memory context for a query.

        Args:
            user_id: Unique identifier for the user (any string: UUID, email, numeric, etc.).
            query: The query to search memory for.
            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.

        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)
        tracked_reranker = _TokenTrackingLLM(self._reranker_llm)

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

        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,
                )
                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 + tracked_reranker.total_usage,
            config_effective=config_dict,
        )

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