"""
mubit-agno — Agno memory DB and toolkit integration backed by MuBit.

Provides three integration surfaces for Agno agents:

1. MubitMemoryDb — Implements Agno's memory DB interface so built-in
   memory (automatic or agentic) persists to MuBit.
2. MubitToolkit — Agno Toolkit giving agents LLM-callable tools for
   remember, recall, reflect, context, checkpoint, and diagnose.
3. MubitAgnoMemory — Convenience wrapper bundling both surfaces plus
   full MAS extensions (handoff, feedback, outcomes, strategies, etc.).
"""

from __future__ import annotations

import json
import uuid
from typing import Any, Dict, List, Optional

from .mubit_client import MubitControlClient, normalize_metadata

try:
    from agno.memory.v2.db.base import MemoryDb
    from agno.memory.v2.schema import MemoryRow

    _HAS_AGNO_MEMORY = True
except ImportError:
    _HAS_AGNO_MEMORY = False

    class MemoryDb:  # type: ignore[no-redef]
        """Fallback base class when agno is not installed."""
        pass

    class MemoryRow:  # type: ignore[no-redef]
        """Fallback stub."""
        def __init__(self, **kwargs):
            for k, v in kwargs.items():
                setattr(self, k, v)

try:
    from agno.tools import Toolkit

    _HAS_AGNO_TOOLS = True
except ImportError:
    _HAS_AGNO_TOOLS = False

    class Toolkit:  # type: ignore[no-redef]
        """Fallback base class when agno is not installed."""
        def __init__(self, **kwargs):
            pass


# ---------------------------------------------------------------------------
# Intent mapping: Agno memory type → MuBit intent
# ---------------------------------------------------------------------------

def _memory_type_to_intent(memory_type: Optional[str]) -> str:
    if not memory_type:
        return "trace"
    t = memory_type.lower()
    if t in ("semantic", "fact"):
        return "fact"
    if t in ("episodic", "trace", "event"):
        return "trace"
    if t in ("procedural", "lesson", "rule", "skill"):
        return "lesson"
    return "trace"


# ---------------------------------------------------------------------------
# MubitMemoryDb — Agno MemoryDb interface
# ---------------------------------------------------------------------------

class MubitMemoryDb(MemoryDb):
    """Agno MemoryDb backed by MuBit memory engine.

    Drop-in replacement for Agno's default memory storage, routing all
    memory operations through MuBit's control plane.

    Usage::

        from agno.agent import Agent
        from mubit_agno import MubitMemoryDb

        agent = Agent(
            model=...,
            memory=Memory(db=MubitMemoryDb(api_key="mbt_...")),
        )
    """

    def __init__(
        self,
        *,
        endpoint: str = "http://127.0.0.1:3000",
        api_key: str = "",
        session_id: str = "default",
        user_id: str = "",
        agent_id: str = "agno",
        mubit_client: Optional[Any] = None,
    ) -> None:
        self._session_id = session_id
        self._user_id = user_id
        self._agent_id = agent_id
        self._client = mubit_client or MubitControlClient(
            endpoint=endpoint.rstrip("/"), api_key=api_key
        )

    def create(self) -> None:
        """No-op — MuBit manages its own storage."""
        pass

    def memory_exists(self, memory: MemoryRow) -> bool:
        """Check if a memory with this ID exists by querying MuBit."""
        memory_id = getattr(memory, "id", None) or getattr(memory, "memory_id", None)
        if not memory_id:
            return False
        data = self._client.recall(
            session_id=self._session_id,
            user_id=self._user_id or None,
            query=getattr(memory, "memory", "") or str(memory_id),
            limit=5,
            agent_id=self._agent_id,
        )
        for ev in data.get("evidence", []):
            meta = normalize_metadata(ev)
            if meta.get("agno_memory_id") == str(memory_id):
                return True
        return False

    def read_memories(
        self,
        user_id: Optional[str] = None,
        limit: Optional[int] = None,
        sort: Optional[str] = None,
    ) -> List[MemoryRow]:
        """Retrieve memories from MuBit for the given user."""
        effective_user = user_id or self._user_id
        data = self._client.recall(
            session_id=self._session_id,
            user_id=effective_user or None,
            query="all memories",
            limit=limit or 100,
            agent_id=self._agent_id,
        )
        rows: List[MemoryRow] = []
        for ev in data.get("evidence", []):
            meta = normalize_metadata(ev)
            memory_id = meta.get("agno_memory_id", ev.get("id", str(uuid.uuid4())))
            topics = meta.get("topics")
            if isinstance(topics, str):
                try:
                    topics = json.loads(topics)
                except json.JSONDecodeError:
                    topics = [topics] if topics else []
            row_kwargs = {
                "id": memory_id,
                "user_id": effective_user or "",
                "memory": ev.get("content", ""),
                "topics": topics if isinstance(topics, list) else [],
            }
            if _HAS_AGNO_MEMORY:
                rows.append(MemoryRow(**row_kwargs))
            else:
                rows.append(type("MemoryRow", (), row_kwargs)())
        return rows

    def upsert_memory(self, memory: MemoryRow, create: bool = True) -> Optional[MemoryRow]:
        """Create or update a memory in MuBit."""
        memory_id = getattr(memory, "id", None) or getattr(memory, "memory_id", None) or str(uuid.uuid4())
        content = getattr(memory, "memory", "") or ""
        user_id = getattr(memory, "user_id", None) or self._user_id
        topics = getattr(memory, "topics", None) or []

        metadata: Dict[str, Any] = {
            "agno_memory_id": str(memory_id),
            "source": "agno",
        }
        if topics:
            metadata["topics"] = topics

        intent = _memory_type_to_intent(
            getattr(memory, "memory_type", None)
        )

        self._client.remember(
            session_id=self._session_id,
            user_id=user_id or None,
            text=content,
            agent_id=self._agent_id,
            intent=intent,
            upsert_key=str(memory_id),
            metadata=metadata,
        )
        return memory

    def delete_memory(self, memory_id: str) -> None:
        """Delete a memory by ID.

        MuBit does not have single-item delete via the control plane,
        so we overwrite the memory with an empty marker using upsert_key.
        """
        self._client.remember(
            session_id=self._session_id,
            user_id=self._user_id or None,
            text="[deleted]",
            agent_id=self._agent_id,
            intent="trace",
            upsert_key=str(memory_id),
            metadata={"agno_memory_id": str(memory_id), "deleted": True},
        )

    def drop_table(self) -> None:
        """Delete all memories for the current session."""
        self._client.delete_run(session_id=self._session_id)

    def table_exists(self) -> bool:
        """MuBit storage is always available."""
        return True

    def clear(self) -> None:
        """Clear all memories for the current session."""
        self._client.delete_run(session_id=self._session_id)


# ---------------------------------------------------------------------------
# MubitToolkit — Agno Toolkit with LLM-callable tools
# ---------------------------------------------------------------------------

class MubitToolkit(Toolkit):
    """Agno Toolkit giving agents direct access to MuBit memory operations.

    Usage::

        from agno.agent import Agent
        from mubit_agno import MubitToolkit

        agent = Agent(
            model=...,
            tools=[MubitToolkit(api_key="mbt_...")],
        )
    """

    def __init__(
        self,
        *,
        endpoint: str = "http://127.0.0.1:3000",
        api_key: str = "",
        session_id: str = "default",
        user_id: str = "",
        agent_id: str = "agno",
        mubit_client: Optional[Any] = None,
    ) -> None:
        super().__init__(name="mubit_memory")
        self._session_id = session_id
        self._user_id = user_id
        self._agent_id = agent_id
        self._client = mubit_client or MubitControlClient(
            endpoint=endpoint.rstrip("/"), api_key=api_key
        )
        self.register(self.mubit_remember)
        self.register(self.mubit_recall)
        self.register(self.mubit_reflect)
        self.register(self.mubit_get_context)
        self.register(self.mubit_checkpoint)
        self.register(self.mubit_diagnose)
        self.register(self.mubit_memory_health)

    def mubit_remember(self, content: str, intent: str = "trace") -> str:
        """Store a memory in MuBit.

        Args:
            content: The text content to remember.
            intent: Memory type — one of 'fact', 'trace', 'lesson', 'rule',
                    'observation', 'tool_output'. Defaults to 'trace'.

        Returns:
            JSON string with the ingestion result.
        """
        result = self._client.remember(
            session_id=self._session_id,
            user_id=self._user_id or None,
            text=content,
            agent_id=self._agent_id,
            intent=intent,
        )
        return json.dumps(result, default=str)

    def mubit_recall(self, query: str, limit: int = 5) -> str:
        """Search memories in MuBit using semantic search.

        Args:
            query: Natural language search query.
            limit: Maximum number of results to return.

        Returns:
            JSON string with evidence list and confidence scores.
        """
        result = self._client.recall(
            session_id=self._session_id,
            user_id=self._user_id or None,
            query=query,
            limit=limit,
            agent_id=self._agent_id,
        )
        return json.dumps(result, default=str)

    def mubit_reflect(self) -> str:
        """Extract lessons from the current session's evidence.

        Triggers MuBit's reflection pipeline to discover success/failure
        patterns, observations, and rules from the session history.

        Returns:
            JSON string with extracted lessons.
        """
        result = self._client.reflect(
            session_id=self._session_id,
            user_id=self._user_id or None,
            agent_id=self._agent_id,
        )
        return json.dumps(result, default=str)

    def mubit_get_context(self, query: str, max_token_budget: int = 2048) -> str:
        """Get pre-assembled context from MuBit memory.

        Retrieves a formatted context block containing relevant memories,
        lessons, and facts for the given query.

        Args:
            query: What context is needed for.
            max_token_budget: Maximum tokens for the context block.

        Returns:
            JSON string with context block and sources.
        """
        result = self._client.get_context(
            session_id=self._session_id,
            user_id=self._user_id or None,
            query=query,
            agent_id=self._agent_id,
            max_token_budget=max_token_budget,
        )
        return json.dumps(result, default=str)

    def mubit_checkpoint(self, label: str, snapshot: str) -> str:
        """Create a checkpoint of the current memory state.

        Args:
            label: Short label for the checkpoint.
            snapshot: Context snapshot text to preserve.

        Returns:
            JSON string with checkpoint confirmation.
        """
        result = self._client.checkpoint(
            session_id=self._session_id,
            user_id=self._user_id or None,
            snapshot=snapshot,
            label=label,
            agent_id=self._agent_id,
        )
        return json.dumps(result, default=str)

    def mubit_diagnose(self, error_text: str) -> str:
        """Surface failure-path lessons relevant to an error.

        Args:
            error_text: Description of the error or failure to diagnose.

        Returns:
            JSON string with relevant failure lessons and suggestions.
        """
        result = self._client.diagnose(
            session_id=self._session_id,
            user_id=self._user_id or None,
            error_text=error_text,
        )
        return json.dumps(result, default=str)

    def mubit_memory_health(self) -> str:
        """Get a quality and staleness report for session memories.

        Returns:
            JSON string with memory health metrics.
        """
        result = self._client.memory_health(
            session_id=self._session_id,
            user_id=self._user_id or None,
        )
        return json.dumps(result, default=str)


# ---------------------------------------------------------------------------
# MubitAgnoMemory — Convenience wrapper with full MAS extensions
# ---------------------------------------------------------------------------

class MubitAgnoMemory:
    """Convenience class for integrating MuBit memory with Agno.

    Bundles MubitMemoryDb + MubitToolkit and exposes the full set of
    MuBit MAS extensions (handoff, feedback, outcomes, strategies, etc.).

    Usage::

        from agno.agent import Agent
        from agno.memory.v2.memory import Memory
        from mubit_agno import MubitAgnoMemory

        mubit = MubitAgnoMemory(api_key="mbt_...", session_id="run-1")

        agent = Agent(
            model=...,
            memory=Memory(db=mubit.as_memory_db()),
            tools=[mubit.as_toolkit()],
        )

        # Extended MuBit features
        mubit.checkpoint("Phase 1 done", "Completed research")
        mubit.record_outcome("task-1", "success", rationale="All checks passed")
    """

    def __init__(
        self,
        *,
        endpoint: str = "http://127.0.0.1:3000",
        api_key: str = "",
        session_id: str = "default",
        user_id: str = "",
        agent_id: str = "agno",
        mubit_client: Optional[Any] = None,
    ) -> None:
        self._session_id = session_id
        self._user_id = user_id
        self._agent_id = agent_id
        self._client = mubit_client or MubitControlClient(
            endpoint=endpoint.rstrip("/"), api_key=api_key
        )
        self._memory_db = MubitMemoryDb(
            session_id=session_id,
            user_id=user_id,
            agent_id=agent_id,
            mubit_client=self._client,
        )
        self._toolkit = MubitToolkit(
            session_id=session_id,
            user_id=user_id,
            agent_id=agent_id,
            mubit_client=self._client,
        )

    def as_memory_db(self) -> MubitMemoryDb:
        """Return the MubitMemoryDb instance for Agno's Memory(db=...)."""
        return self._memory_db

    def as_toolkit(self) -> MubitToolkit:
        """Return the MubitToolkit instance for Agent(tools=[...])."""
        return self._toolkit

    # -- Context & Retrieval ------------------------------------------------

    def get_context(
        self,
        query: str,
        *,
        mode: Optional[str] = None,
        sections: Optional[List[str]] = None,
        entry_types: Optional[List[str]] = None,
        max_token_budget: int = 2048,
    ) -> str:
        """Fetch assembled context from MuBit.

        Returns:
            Pre-formatted context string.
        """
        result = self._client.get_context(
            session_id=self._session_id,
            user_id=self._user_id or None,
            query=query,
            mode=mode,
            sections=sections,
            entry_types=entry_types,
            max_token_budget=max_token_budget,
            agent_id=self._agent_id,
        )
        return result.get("context_block", "")

    # -- Reflection & Learning ----------------------------------------------

    def reflect(self, *, trigger_reflections: bool = True) -> Dict[str, Any]:
        """Extract lessons from session evidence."""
        return self._client.reflect(
            session_id=self._session_id,
            user_id=self._user_id or None,
            trigger_reflections=trigger_reflections,
            agent_id=self._agent_id,
        )

    def lessons(
        self,
        *,
        scope: Optional[str] = None,
        lesson_type: Optional[str] = None,
        importance: Optional[str] = None,
        limit: int = 50,
    ) -> Dict[str, Any]:
        """List lessons with optional filtering."""
        return self._client.lessons(
            session_id=self._session_id,
            user_id=self._user_id or None,
            scope=scope,
            lesson_type=lesson_type,
            importance=importance,
            limit=limit,
        )

    # -- Checkpoints & Outcomes ---------------------------------------------

    def checkpoint(
        self,
        label: str,
        snapshot: str,
        *,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Create a checkpoint of current memory state."""
        return self._client.checkpoint(
            session_id=self._session_id,
            user_id=self._user_id or None,
            snapshot=snapshot,
            label=label,
            metadata=metadata,
            agent_id=self._agent_id,
        )

    def record_outcome(
        self,
        reference_id: str,
        outcome: str,
        *,
        signal: Optional[float] = None,
        rationale: str = "",
    ) -> Dict[str, Any]:
        """Record outcome feedback for a task or decision."""
        return self._client.record_outcome(
            session_id=self._session_id,
            user_id=self._user_id or None,
            reference_id=reference_id,
            outcome=outcome,
            signal=signal,
            rationale=rationale,
            agent_id=self._agent_id,
        )

    def record_step_outcome(
        self,
        step_id: str,
        *,
        signal: float = 0.0,
        rationale: str = "",
        directive: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Record per-step process reward."""
        return self._client.record_step_outcome(
            session_id=self._session_id,
            user_id=self._user_id or None,
            step_id=step_id,
            signal=signal,
            rationale=rationale,
            directive=directive,
            agent_id=self._agent_id,
        )

    def surface_strategies(
        self,
        *,
        lesson_types: Optional[List[str]] = None,
        max_strategies: int = 5,
    ) -> Dict[str, Any]:
        """Surface strategy clusters from lessons."""
        return self._client.surface_strategies(
            session_id=self._session_id,
            user_id=self._user_id or None,
            lesson_types=lesson_types,
            max_strategies=max_strategies,
        )

    # -- Agent Coordination (MAS) -------------------------------------------

    def register_agent(
        self,
        agent_id: str,
        *,
        role: Optional[str] = None,
        status: Optional[str] = None,
        read_scopes: Optional[List[str]] = None,
        write_scopes: Optional[List[str]] = None,
        shared_memory_lanes: Optional[List[str]] = None,
        capabilities: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Register an agent for MAS coordination."""
        return self._client.register_agent(
            session_id=self._session_id,
            user_id=self._user_id or None,
            agent_id=agent_id,
            role=role,
            status=status,
            read_scopes=read_scopes,
            write_scopes=write_scopes,
            shared_memory_lanes=shared_memory_lanes,
            capabilities=capabilities,
        )

    def list_agents(self) -> Dict[str, Any]:
        """List registered agents for the current session."""
        return self._client.list_agents(
            session_id=self._session_id,
            user_id=self._user_id or None,
        )

    def handoff(
        self,
        from_agent_id: str,
        to_agent_id: str,
        content: str,
        *,
        requested_action: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Create a handoff between agents."""
        return self._client.handoff(
            session_id=self._session_id,
            user_id=self._user_id or None,
            from_agent_id=from_agent_id,
            to_agent_id=to_agent_id,
            content=content,
            requested_action=requested_action,
            metadata=metadata,
        )

    def feedback(
        self,
        handoff_id: str,
        verdict: str,
        *,
        comments: Optional[str] = None,
        from_agent_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Submit feedback on a handoff."""
        return self._client.feedback(
            session_id=self._session_id,
            user_id=self._user_id or None,
            handoff_id=handoff_id,
            verdict=verdict,
            comments=comments,
            from_agent_id=from_agent_id,
            metadata=metadata,
        )

    # -- Archive & Reference ------------------------------------------------

    def archive(
        self,
        content: str,
        artifact_kind: str,
        *,
        metadata: Optional[Dict[str, Any]] = None,
        labels: Optional[List[str]] = None,
        family: Optional[str] = None,
        importance: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Archive an exact reusable artifact with a stable reference ID."""
        return self._client.archive(
            session_id=self._session_id,
            user_id=self._user_id or None,
            content=content,
            artifact_kind=artifact_kind,
            metadata=metadata,
            agent_id=self._agent_id,
            labels=labels,
            family=family,
            importance=importance,
        )

    def dereference(self, reference_id: str) -> Dict[str, Any]:
        """Fetch exact content by reference ID."""
        return self._client.dereference(
            session_id=self._session_id,
            user_id=self._user_id or None,
            reference_id=reference_id,
            agent_id=self._agent_id,
        )

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

    def diagnose(
        self,
        error_text: str,
        *,
        error_type: Optional[str] = None,
        limit: int = 10,
    ) -> Dict[str, Any]:
        """Surface failure-path lessons for debugging."""
        return self._client.diagnose(
            session_id=self._session_id,
            user_id=self._user_id or None,
            error_text=error_text,
            error_type=error_type,
            limit=limit,
        )

    def memory_health(self, *, limit: int = 100) -> Dict[str, Any]:
        """Get quality and staleness report for session memories."""
        return self._client.memory_health(
            session_id=self._session_id,
            user_id=self._user_id or None,
            limit=limit,
        )
