"""
mubit-adk — Google ADK BaseMemoryService backed by MuBit.

This adapter now routes MuBit calls through the canonical Python SDK transport
instead of maintaining its own raw HTTP client.
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import Any, Optional

from .mubit_client import MubitControlClient, normalize_metadata

try:
    from google.adk.memory.base_memory_service import (
        BaseMemoryService,
        SearchMemoryResponse,
    )
    from google.adk.memory.memory_entry import MemoryEntry
    from google.adk.sessions.session import Session
    from google.adk.events.event import Event
    from google.genai import types as genai_types
except ImportError as e:
    raise ImportError("google-adk is required: pip install mubit-adk[adk]") from e


def _event_to_text(event: Event) -> str:
    parts: list[str] = []
    if hasattr(event, "content") and event.content:
        content = event.content
        if hasattr(content, "parts") and content.parts:
            for part in content.parts:
                if hasattr(part, "text") and part.text:
                    parts.append(part.text)
    if not parts and hasattr(event, "text") and event.text:
        parts.append(event.text)
    return "\n".join(parts) if parts else str(event)


def _memory_entry_to_text(entry: MemoryEntry) -> str:
    if hasattr(entry, "content") and entry.content:
        content = entry.content
        if hasattr(content, "parts") and content.parts:
            texts = [part.text for part in content.parts if getattr(part, "text", None)]
            if texts:
                return "\n".join(texts)
    return str(entry)


class MubitMemoryService(BaseMemoryService):
    def __init__(
        self,
        *,
        endpoint: str = "http://127.0.0.1:3000",
        api_key: str = "",
        mubit_client: Optional[Any] = None,
    ) -> None:
        self._client = mubit_client or MubitControlClient(
            endpoint=endpoint.rstrip("/"), api_key=api_key
        )

    async def add_session_to_memory(self, session: Session) -> None:
        if not session.events:
            return

        items: list[dict[str, Any]] = []
        for idx, event in enumerate(session.events):
            text = _event_to_text(event)
            if not text.strip():
                continue
            author = getattr(event, "author", None) or "unknown"
            item: dict[str, Any] = {
                "item_id": f"adk-{session.id}-{idx}",
                "content_type": "text/plain",
                "text": text,
                "metadata": {
                    "author": author,
                    "event_index": idx,
                    "session_id": session.id,
                },
                "user_id": session.user_id or "",
            }
            if author == "user":
                item["intent"] = "context"
            elif author == "tool":
                item["intent"] = "tool_output"
            else:
                item["intent"] = "trace"
            items.append(item)

        if not items:
            return

        self._client.ingest_items(
            session_id=session.id,
            agent_id=session.app_name or "adk-agent",
            idempotency_key=f"session-{session.id}",
            items=items,
        )

    async def search_memory(
        self,
        *,
        app_name: str,
        user_id: str,
        query: str,
    ) -> SearchMemoryResponse:
        data = self._client.recall(
            session_id="",
            user_id=user_id,
            query=query,
            limit=10,
            include_working_memory=True,
        )
        memories: list[MemoryEntry] = []
        for ev in data.get("evidence", []):
            metadata = normalize_metadata(ev)
            content = genai_types.Content(
                parts=[genai_types.Part(text=ev.get("content", ""))],
                role="model",
            )
            memories.append(
                MemoryEntry(
                    content=content,
                    custom_metadata={
                        "score": ev.get("score", 0),
                        "evidence_type": ev.get("evidence_type", ""),
                        **metadata,
                    },
                    id=ev.get("id"),
                    author=metadata.get("author", "memory"),
                )
            )
        return SearchMemoryResponse(memories=memories)

    async def add_events_to_memory(
        self,
        *,
        app_name: str,
        user_id: str,
        events: Sequence[Event],
        session_id: Optional[str] = None,
        custom_metadata: Optional[Mapping[str, object]] = None,
    ) -> None:
        items: list[dict[str, Any]] = []
        for idx, event in enumerate(events):
            text = _event_to_text(event)
            if not text.strip():
                continue
            author = getattr(event, "author", None) or "unknown"
            metadata: dict[str, Any] = {"author": author, "event_index": idx}
            if custom_metadata:
                metadata.update(custom_metadata)
            item: dict[str, Any] = {
                "item_id": f"adk-evt-{idx}",
                "content_type": "text/plain",
                "text": text,
                "metadata": metadata,
                "user_id": user_id,
            }
            if author == "tool":
                item["intent"] = "tool_output"
            elif author == "user":
                item["intent"] = "context"
            else:
                item["intent"] = "trace"
            items.append(item)

        if not items:
            return

        self._client.ingest_items(
            session_id=session_id or "default",
            agent_id=app_name or "adk-agent",
            idempotency_key="",
            items=items,
        )

    async def add_memory(
        self,
        *,
        app_name: str,
        user_id: str,
        memories: Sequence[MemoryEntry],
        custom_metadata: Optional[Mapping[str, object]] = None,
    ) -> None:
        items: list[dict[str, Any]] = []
        for idx, entry in enumerate(memories):
            text = _memory_entry_to_text(entry)
            if not text.strip():
                continue
            metadata: dict[str, Any] = dict(entry.custom_metadata or {})
            if custom_metadata:
                metadata.update(custom_metadata)
            if entry.author:
                metadata["author"] = entry.author
            items.append(
                {
                    "item_id": entry.id or f"adk-mem-{idx}",
                    "content_type": "text/plain",
                    "text": text,
                    "metadata": metadata,
                    "user_id": user_id,
                    "intent": "fact",
                }
            )

        if not items:
            return

        self._client.ingest_items(
            session_id="direct",
            agent_id=app_name or "adk-agent",
            idempotency_key="",
            items=items,
        )

    async def checkpoint(
        self,
        *,
        app_name: str,
        user_id: str,
        snapshot: str,
        session_id: Optional[str] = None,
        label: str = "",
        metadata: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        return self._client.checkpoint(
            session_id=session_id or "default",
            user_id=user_id,
            snapshot=snapshot,
            label=label,
            metadata=metadata,
            agent_id=app_name or "adk-agent",
        )

    async def record_outcome(
        self,
        *,
        app_name: str,
        user_id: str,
        reference_id: str,
        outcome: str,
        signal: Optional[float] = None,
        rationale: str = "",
        session_id: Optional[str] = None,
    ) -> dict[str, Any]:
        return self._client.record_outcome(
            session_id=session_id or "default",
            user_id=user_id,
            reference_id=reference_id,
            outcome=outcome,
            signal=signal,
            rationale=rationale,
            agent_id=app_name or "adk-agent",
        )

    async def archive(
        self,
        *,
        app_name: str,
        user_id: str,
        content: str,
        artifact_kind: str,
        session_id: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
        agent_id: Optional[str] = None,
        origin_agent_id: Optional[str] = None,
        source_attempt_id: Optional[str] = None,
        source_tool: Optional[str] = None,
        labels: Optional[list[str]] = None,
        family: Optional[str] = None,
        importance: Optional[str] = None,
    ) -> dict[str, Any]:
        return self._client.archive(
            session_id=session_id or "default",
            user_id=user_id,
            content=content,
            artifact_kind=artifact_kind,
            metadata=metadata,
            agent_id=agent_id or app_name or "adk-agent",
            origin_agent_id=origin_agent_id,
            source_attempt_id=source_attempt_id,
            source_tool=source_tool,
            labels=labels,
            family=family,
            importance=importance,
        )

    async def dereference(
        self,
        *,
        user_id: str,
        reference_id: str,
        session_id: Optional[str] = None,
        agent_id: Optional[str] = None,
    ) -> dict[str, Any]:
        return self._client.dereference(
            session_id=session_id or "default",
            user_id=user_id,
            reference_id=reference_id,
            agent_id=agent_id,
        )

    async def surface_strategies(
        self,
        *,
        session_id: Optional[str] = None,
        lesson_types: Optional[list[str]] = None,
        max_strategies: int = 5,
    ) -> dict[str, Any]:
        return self._client.surface_strategies(
            session_id=session_id or "default",
            lesson_types=lesson_types,
            max_strategies=max_strategies,
        )

    async def get_context(self, *, app_name: str, user_id: str, query: str, session_id: Optional[str] = None, **kwargs):
        return self._client.get_context(
            session_id=session_id or "default",
            user_id=user_id,
            query=query,
            agent_id=app_name or "adk-agent",
            mode=kwargs.get("mode"),
            sections=kwargs.get("sections"),
            entry_types=kwargs.get("entry_types"),
            limit=kwargs.get("limit"),
            max_token_budget=kwargs.get("max_token_budget"),
        )

    async def memory_health(self, *, user_id: str, session_id: Optional[str] = None, limit: int = 100):
        return self._client.memory_health(
            session_id=session_id or "default",
            user_id=user_id,
            limit=limit,
        )

    async def diagnose(self, *, user_id: str, error_text: str, session_id: Optional[str] = None, error_type: Optional[str] = None, limit: int = 10):
        return self._client.diagnose(
            session_id=session_id or "default",
            user_id=user_id,
            error_text=error_text,
            error_type=error_type,
            limit=limit,
        )

    async def register_agent(self, *, user_id: str, session_id: Optional[str] = None, agent_id: str, role: Optional[str] = None, **kwargs):
        return self._client.register_agent(
            session_id=session_id or "default",
            user_id=user_id,
            agent_id=agent_id,
            role=role,
            status=kwargs.get("status"),
            read_scopes=kwargs.get("read_scopes"),
            write_scopes=kwargs.get("write_scopes"),
            shared_memory_lanes=kwargs.get("shared_memory_lanes"),
            capabilities=kwargs.get("capabilities"),
        )

    async def list_agents(self, *, user_id: str, session_id: Optional[str] = None):
        return self._client.list_agents(
            session_id=session_id or "default",
            user_id=user_id,
        )

    async def handoff(self, *, user_id: str, session_id: Optional[str] = None, from_agent_id: str, to_agent_id: str, content: str, **kwargs):
        return self._client.handoff(
            session_id=session_id or "default",
            user_id=user_id,
            from_agent_id=from_agent_id,
            to_agent_id=to_agent_id,
            content=content,
            requested_action=kwargs.get("requested_action"),
            metadata=kwargs.get("metadata"),
        )

    async def feedback(self, *, user_id: str, session_id: Optional[str] = None, handoff_id: str, verdict: str, **kwargs):
        return self._client.feedback(
            session_id=session_id or "default",
            user_id=user_id,
            handoff_id=handoff_id,
            verdict=verdict,
            comments=kwargs.get("comments"),
            from_agent_id=kwargs.get("from_agent_id"),
            metadata=kwargs.get("metadata"),
        )

    async def close(self) -> None:
        return None
