"""
cortex/sdk.py — Cortex: Clean Public Python SDK for evaOS memory.

Wraps CortexPlugin with a developer-friendly interface that maps human
concepts (wake/remember/retrieve/sleep) to the underlying machinery.

Usage::

    from evaos import Cortex

    cortex = Cortex(db_path="cortex.db", profile="companion")
    await cortex.initialize()
    await cortex.wake(session_id="session_001")
    await cortex.remember("User prefers dark mode")
    context = await cortex.retrieve("What are the user's preferences?")
    await cortex.sleep(session_id="session_001")

Design goals:
  - One import, one class, zero friction.
  - Every method maps cleanly to a human concept.
  - All implementation detail stays in CortexPlugin.
  - Typed return values wherever possible.
  - Graceful degradation — nothing raises unless critical.
"""
from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Profile → LLM provider priority presets
# ---------------------------------------------------------------------------

_PROFILE_PROVIDERS: Dict[str, List[str]] = {
    "companion": ["openai"],
    "developer": ["openai", "anthropic"],
    "local": ["ollama"],
    "minimal": ["openai"],
}


def _providers_for_profile(profile: str, llm_provider: str) -> List[str]:
    """Return LLM provider priority list for a given profile + explicit override."""
    if llm_provider != "openai":
        # Explicit provider always wins
        return [llm_provider]
    return _PROFILE_PROVIDERS.get(profile, ["openai"])


# ---------------------------------------------------------------------------
# Cortex SDK class
# ---------------------------------------------------------------------------


class Cortex:
    """
    The evaOS Python SDK. Clean public API for memory operations.

    Wraps :class:`cortex.api.plugin.CortexPlugin` with a simpler surface
    designed for application developers (not framework integrators).

    Usage::

        from evaos import Cortex

        cortex = Cortex(db_path="my_memory.db", profile="companion")
        await cortex.initialize()
        await cortex.wake(session_id="session_001")
        await cortex.remember("User prefers dark mode")
        context = await cortex.retrieve("What are the user's preferences?")
        await cortex.sleep(session_id="session_001")

    Args:
        db_path:        Path to the SQLite memory database (created if absent).
        config_path:    Optional path to a ``cortex.toml`` config file.
        profile:        Memory profile preset — 'companion' | 'developer' | 'local' | 'minimal'.
        llm_provider:   Primary LLM provider — 'openai' | 'anthropic' | 'ollama'.
        embed_model:    Override the embedding model (defaults from config/profile).
        extract_model:  Override the extraction/completion model.
    """

    def __init__(
        self,
        db_path: str = "cortex.db",
        config_path: Optional[str] = None,
        profile: str = "companion",
        llm_provider: str = "openai",
        embed_model: Optional[str] = None,
        extract_model: Optional[str] = None,
    ) -> None:
        self._db_path = db_path
        self._config_path = config_path
        self._profile = profile
        self._llm_provider = llm_provider
        self._embed_model = embed_model
        self._extract_model = extract_model
        self._plugin: Any = None
        self._active_session_id: Optional[str] = None

        # Build the plugin eagerly (config + module wiring happen in __init__)
        self._plugin = self._build_plugin()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _build_plugin(self) -> Any:
        """Instantiate CortexPlugin with SDK-level parameter mapping."""
        from cortex.api.plugin import CortexPlugin

        plugin = CortexPlugin(
            config_path=self._config_path,
            db_path=self._db_path,
        )

        # Apply profile / model overrides on top of the loaded config
        providers = _providers_for_profile(self._profile, self._llm_provider)
        plugin.config.llm_provider_priority = providers

        if self._embed_model:
            plugin.config.embedding_model = self._embed_model
        if self._extract_model:
            plugin.config.extraction_model = self._extract_model

        return plugin

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def initialize(self) -> None:
        """Initialize storage backend and apply pending schema migrations.

        Call once before any other SDK method when operating outside a
        context manager.  Safe to call multiple times (idempotent).
        """
        if self._plugin is None:
            self._plugin = self._build_plugin()
        if self._plugin.storage is not None:
            try:
                await self._plugin.storage.initialize()
                logger.info("Cortex storage initialized (db=%s)", self._db_path)
            except Exception as exc:
                logger.warning("Storage initialize failed (may already be ready): %s", exc)

    async def wake(
        self,
        session_id: Optional[str] = None,
        owner_id: str = "default",
    ) -> Any:
        """Start a memory session.

        Loads identity cornerstones into memory and resets the per-session
        LLM budget counter.

        Args:
            session_id: Caller-supplied session identifier.  Auto-generated
                        if omitted.
            owner_id:   Memory owner namespace (e.g. user UUID).

        Returns:
            WakeReport with ``session_id`` and ``cornerstones_loaded``.
        """
        if owner_id != "default":
            self._plugin.config.owner_id = owner_id
        report = await self._plugin.wake(session_id=session_id)
        self._active_session_id = report.session_id
        return report

    async def remember(
        self,
        content: Any,
        session_id: Optional[str] = None,
        owner_id: str = "default",
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Any:
        """Store information. Claim extraction happens automatically.

        Accepts a plain string, a list of ``{"role": ..., "content": ...}``
        message dicts, or any form CortexPlugin accepts.

        Args:
            content:    Text or conversation to remember.
            session_id: Session to associate with; uses active session if omitted.
            owner_id:   Memory owner namespace.
            metadata:   Reserved for future metadata tagging (ignored for now).

        Returns:
            RememberResult with extraction + storage counts and any soft errors.
        """
        if owner_id != "default":
            self._plugin.config.owner_id = owner_id
        sid = session_id or self._active_session_id
        return await self._plugin.remember(conversation=content, session_id=sid)

    async def retrieve(
        self,
        query: str,
        owner_id: str = "default",
        top_k: int = 10,
        token_budget: Optional[int] = None,
    ) -> Any:
        """Retrieve relevant memories for a query.

        Returns a RetrievalResult whose ``.context_block`` field is ready for
        direct injection into an LLM prompt.

        Args:
            query:        Natural-language retrieval query.
            owner_id:     Memory owner namespace.
            top_k:        Maximum number of memories to surface (best-effort).
            token_budget: Hard token ceiling for the assembled context block.

        Returns:
            RetrievalResult, or None on retrieval failure.
        """
        if owner_id != "default":
            self._plugin.config.owner_id = owner_id
        if top_k != 10:
            # Map top_k → token budget if not explicitly set
            token_budget = token_budget or (top_k * 100)
        return await self._plugin.retrieve(
            query=query,
            token_budget=token_budget,
        )

    async def ask(
        self,
        query: str,
        owner_id: str = "default",
        history: Optional[List[Dict[str, Any]]] = None,
    ) -> Any:
        """Ask a question and get a synthesised answer from memory.

        Dialectic fast path: retrieves relevant memories and uses an LLM to
        synthesise a grounded answer.  Replaces the former stub that fell back
        to bare retrieval.

        Args:
            query:    Natural-language question.
            owner_id: Memory owner namespace.
            history:  Optional conversation history for follow-up questions.
                      List of ``{role, content}`` dicts (older messages first).

        Returns:
            AskResult with ``answer``, ``sources``, ``confidence``, and ``reasoning``.
        """
        if owner_id != "default":
            self._plugin.config.owner_id = owner_id
        return await self._plugin.ask(
            query=query,
            owner_id=owner_id if owner_id != "default" else None,
            history=history,
        )

    async def sleep(
        self,
        session_id: Optional[str] = None,
        owner_id: str = "default",
    ) -> Any:
        """End the current session and trigger post-session consolidation.

        Args:
            session_id: Session to close; uses active session if omitted.
            owner_id:   Memory owner namespace.

        Returns:
            SleepReport with consolidation counts and any soft errors.
        """
        if owner_id != "default":
            self._plugin.config.owner_id = owner_id
        sid = session_id or self._active_session_id
        report = await self._plugin.sleep(session_id=sid)
        self._active_session_id = None
        return report

    async def dream(self, owner_id: str = "default") -> Dict[str, Any]:
        """Trigger a dream cycle (overnight memory reconsolidation).

        Runs Ebbinghaus decay, consolidation, and identity drift checks.
        Safe to call manually; normally invoked by the circadian scheduler.

        Args:
            owner_id: Memory owner namespace.

        Returns:
            Dict with ``ok``, ``claims_decayed``, ``drift_reports``, ``errors``.
        """
        if owner_id != "default":
            self._plugin.config.owner_id = owner_id

        # Delegate to CortexPlugin.dream() which wires the full CircadianEngine
        # pipeline: orphan recovery, decay, consolidation, journal generation.
        report = await self._plugin.dream(owner_id=owner_id)

        # Normalise to a plain dict so SDK callers get a consistent return type.
        step_fields = [
            "orphans_recovered", "claims_consolidated", "contradictions_resolved",
            "claims_decayed", "claims_archived", "cornerstone_proposals",
            "entity_merge_suggestions", "brain_graphs_refreshed", "feedback_applied",
        ]
        return {
            "ok": True,
            "claims_decayed": getattr(report, "claims_decayed", 0),
            "claims_consolidated": getattr(report, "claims_consolidated", 0),
            "contradictions_resolved": getattr(report, "contradictions_resolved", 0),
            "cornerstone_proposals": getattr(report, "cornerstone_proposals", 0),
            "orphans_recovered": getattr(report, "orphans_recovered", 0),
            "journal": getattr(report, "journal", ""),
            "duration_ms": getattr(report, "duration_ms", 0),
            "steps_completed": sum(1 for f in step_fields if getattr(report, f, 0)),
            "drift_reports": [],
            "errors": getattr(report, "errors", []),
        }

    # ------------------------------------------------------------------
    # Identity / cornerstones
    # ------------------------------------------------------------------

    async def seal_cornerstone(
        self,
        label: str,
        content: str,
        owner_id: str = "default",
    ) -> Any:
        """Seal an identity cornerstone — a memory that must not drift.

        Cornerstones are protected from deletion and tracked for drift
        against their original content.

        Args:
            label:    Short human-readable label (e.g. "core-identity").
            content:  The canonical content to protect.
            owner_id: Memory owner namespace.

        Returns:
            CornerstoneMemory on success, or None on failure.
        """
        if owner_id != "default":
            self._plugin.config.owner_id = owner_id
        return await self._plugin.seal_cornerstone(
            memory_id=label,
            label=label,
            content=content,
        )

    async def check_drift(self, owner_id: str = "default") -> List[Any]:
        """Check drift scores for all cornerstones.

        Compares current cornerstone content against their sealed originals
        and returns a drift report per cornerstone.

        Args:
            owner_id: Memory owner namespace.

        Returns:
            List of DriftReport objects (empty if no cornerstones exist).
        """
        if owner_id != "default":
            self._plugin.config.owner_id = owner_id

        try:
            from cortex.identity.drift_calculator import DriftCalculator
            if self._plugin.storage is None or self._plugin.llm is None:
                return []
            drift_calc = DriftCalculator(
                storage=self._plugin.storage,
                llm=self._plugin.llm,
                config=self._plugin.config,
            )
            return await drift_calc.check_all(owner_id=owner_id)
        except Exception as exc:
            logger.warning("check_drift failed: %s", exc)
            return []

    # ------------------------------------------------------------------
    # Feedback
    # ------------------------------------------------------------------

    async def feedback(
        self,
        claim_id: str,
        feedback_type: str,
        value: Optional[Any] = None,
    ) -> Any:
        """Submit feedback on a retrieved memory.

        Args:
            claim_id:      ID of the memory claim being rated.
            feedback_type: 'helpful' | 'unhelpful' | 'incorrect'.
            value:         Optional extra value (score, note, etc.).

        Returns:
            MemoryFeedback record, or None on failure.
        """
        helpful = feedback_type in ("helpful", "correct", "good")
        context = str(value) if value is not None else None
        return await self._plugin.feedback(
            memory_id=claim_id,
            helpful=helpful,
            context=context,
        )

    # ------------------------------------------------------------------
    # Diagnostics
    # ------------------------------------------------------------------

    async def health(self) -> Dict[str, Any]:
        """Return a system health report.

        Checks all internal modules (storage, LLM, extraction, retrieval, etc.)
        and returns a flat dict with per-module status strings.

        Returns:
            Dict with ``ok`` (bool), ``storage``, ``modules``, ``errors``.
        """
        report = await self._plugin.health()
        return {
            "ok": report.ok,
            "storage": report.storage,
            "modules": report.modules,
            "errors": report.errors,
            "timestamp": report.timestamp,
        }

    async def stats(self, owner_id: str = "default") -> Dict[str, Any]:
        """Get memory statistics for an owner.

        Args:
            owner_id: Memory owner namespace.

        Returns:
            Dict with ``claims``, ``entities``, ``cornerstones``, ``sessions``.
        """
        if owner_id != "default":
            self._plugin.config.owner_id = owner_id

        result: Dict[str, Any] = {
            "owner_id": owner_id,
            "claims": 0,
            "entities": 0,
            "cornerstones": 0,
            "sessions": 0,
            "errors": [],
        }

        if self._plugin.storage is None:
            result["errors"].append("storage unavailable")
            return result

        try:
            result["claims"] = await self._plugin.storage.count_claims(owner_id=owner_id)
        except Exception as exc:
            result["errors"].append(f"claims count failed: {exc}")

        try:
            result["entities"] = await self._plugin.storage.count_entities(owner_id=owner_id)
        except Exception as exc:
            result["errors"].append(f"entities count failed: {exc}")

        try:
            result["sessions"] = await self._plugin.storage.count_sessions(owner_id=owner_id)
        except Exception as exc:
            result["errors"].append(f"sessions count failed: {exc}")

        try:
            cornerstones = await self._plugin.get_cornerstones()
            result["cornerstones"] = len(cornerstones)
        except Exception as exc:
            result["errors"].append(f"cornerstones count failed: {exc}")

        return result

    async def export(
        self,
        format: str = "json",  # noqa: A002 (shadow builtin intentional for API clarity)
        owner_id: str = "default",
    ) -> Any:
        """Export memories to a portable format.

        Args:
            format:   'json' (default) — additional formats planned.
            owner_id: Memory owner namespace.

        Returns:
            List of claim dicts (JSON-serialisable) or empty list on failure.
        """
        if owner_id != "default":
            self._plugin.config.owner_id = owner_id

        if self._plugin.storage is None:
            logger.error("export: storage not available")
            return []

        try:
            claims = await self._plugin.storage.get_claims_by_status(
                status="active",
                owner_id=owner_id,
            )
            if format == "json":
                return [
                    {
                        "id": c.id,
                        "subject": c.subject,
                        "predicate": c.predicate,
                        "object_value": c.object_value,
                        "confidence": getattr(c, "confidence", 1.0),
                        "created_at": getattr(c, "created_at", None),
                    }
                    for c in claims
                ]
            return claims
        except Exception as exc:
            logger.error("export failed: %s", exc, exc_info=True)
            return []

    # ------------------------------------------------------------------
    # Class methods
    # ------------------------------------------------------------------

    @classmethod
    def from_config(cls, config_path: str) -> "Cortex":
        """Create a Cortex instance from a TOML config file.

        Reads ``db_path``, ``profile``, ``llm_provider``, ``embed_model``,
        and ``extract_model`` from the config file; all other fields use
        SDK defaults.

        Args:
            config_path: Path to a ``cortex.toml`` config file.

        Returns:
            A new :class:`Cortex` instance configured from the file.

        Example::

            cortex = Cortex.from_config("~/.config/cortex.toml")
        """
        # Import TOML parser (stdlib 3.11+, else tomli)
        try:
            import tomllib  # type: ignore
        except ImportError:
            try:
                import tomli as tomllib  # type: ignore
            except ImportError:
                raise ImportError(
                    "No TOML parser available. Install 'tomli' for Python < 3.11."
                )

        import os

        resolved = os.path.expanduser(config_path)
        with open(resolved, "rb") as fh:
            raw = tomllib.load(fh)

        cortex_section = raw.get("cortex", {})
        storage_section = cortex_section.get("storage", {})
        llm_section = cortex_section.get("llm", {})
        extraction_section = cortex_section.get("extraction", {})
        embeddings_section = cortex_section.get("embeddings", {})

        db_path = storage_section.get("db_path", "cortex.db")
        profile = cortex_section.get("profile", "companion")
        llm_provider = (
            llm_section.get("providers", ["openai"])[0]
            if isinstance(llm_section.get("providers"), list)
            else llm_section.get("provider", "openai")
        )
        embed_model = embeddings_section.get("model", None)
        extract_model = extraction_section.get("model", None)

        return cls(
            db_path=db_path,
            config_path=config_path,
            profile=profile,
            llm_provider=llm_provider,
            embed_model=embed_model,
            extract_model=extract_model,
        )

    # ------------------------------------------------------------------
    # Repr
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        return (
            f"Cortex(db_path={self._db_path!r}, profile={self._profile!r}, "
            f"llm_provider={self._llm_provider!r})"
        )
