"""Google ADK adapter for Molecule AI workspace runtime.

Wraps Google's Agent Development Kit (google-adk v1.x) as a Molecule AI
WorkspaceAdapter, bridging the A2A protocol to Google ADK's runner/session
model.

Google ADK concepts used
------------------------
- ``google.adk.agents.LlmAgent``  — An LLM-backed agent with instructions and
  optional tools.  Declared with ``model``, ``name``, and ``instruction``.
- ``google.adk.runners.Runner``   — Drives one or more agents inside a session;
  ``run_async()`` streams ``Event`` objects, including the final response text.
- ``google.adk.sessions.InMemorySessionService`` — Manages session state in
  memory.  Each ``Runner`` owns a single ``InMemorySessionService`` instance.

Runtime-config keys (all optional)
------------------------------------
``max_output_tokens`` — int, default 8192.  Forwarded to the ADK ``GenerateContentConfig``.
``temperature``       — float, default 1.0.
``agent_name``        — str, default ``"molecule-adk-agent"``.

Environment variables
---------------------
``GOOGLE_API_KEY``   — Google AI Studio key (required for ``gemini-*`` models).
``GOOGLE_GENAI_USE_VERTEXAI`` — set to ``"1"`` to use Vertex AI instead of AI
                                Studio.  In that case supply
                                ``GOOGLE_CLOUD_PROJECT`` and
                                ``GOOGLE_CLOUD_LOCATION`` as well.
"""

from __future__ import annotations

import logging
import os
from typing import TYPE_CHECKING, Any

from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.helpers import new_text_message

from adapter_base import AdapterConfig, BaseAdapter

# Import sanitize_agent_error from the workspace package. The adapter lives
# in the workspace/adapters/ hierarchy so the workspace package root is
# always importable as long as the module is loaded from within a workspace.
# In standalone template repos, this import resolves via the workspace package
# entry point that also provides adapter_base.
try:
    from executor_helpers import sanitize_agent_error  # type: ignore[attr-defined]
except ImportError:  # pragma: no cover
    sanitize_agent_error = None  # fallback: below handler falls back to class-name only

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_DEFAULT_AGENT_NAME = "molecule-adk-agent"
_DEFAULT_MAX_OUTPUT_TOKENS = 8192
_DEFAULT_TEMPERATURE = 1.0
_NO_TEXT_MSG = "Error: message contained no text content."
_NO_RESPONSE_MSG = "(no response generated)"


# ---------------------------------------------------------------------------
# GoogleADKA2AExecutor
# ---------------------------------------------------------------------------


class GoogleADKA2AExecutor(AgentExecutor):
    """A2A executor backed by a Google ADK ``Runner``.

    Each executor instance owns a single ``Runner`` and ``InMemorySessionService``.
    Sessions are created on first use and reused across subsequent turns
    (the session_id is derived from the A2A context_id so each task gets a
    stable, isolated session).

    Parameters
    ----------
    model:
        ADK model identifier, e.g. ``"gemini-2.0-flash"`` or
        ``"gemini-1.5-pro"``.
    system_prompt:
        Optional instruction prepended to every conversation.  Passed to
        ``LlmAgent(instruction=...)``.
    agent_name:
        Internal ADK agent name.  Defaults to ``_DEFAULT_AGENT_NAME``.
    max_output_tokens:
        Token cap forwarded to ``GenerateContentConfig``.
    temperature:
        Sampling temperature forwarded to ``GenerateContentConfig``.
    heartbeat:
        Optional ``HeartbeatLoop`` instance (unused directly but stored for
        future heartbeat integration).
    _runner:
        Inject a pre-built ``Runner`` — for testing only.  When provided,
        the real ADK ``Runner`` is never constructed.
    """

    def __init__(
        self,
        model: str,
        system_prompt: str | None = None,
        agent_name: str = _DEFAULT_AGENT_NAME,
        max_output_tokens: int = _DEFAULT_MAX_OUTPUT_TOKENS,
        temperature: float = _DEFAULT_TEMPERATURE,
        heartbeat: Any = None,
        _runner: Any = None,
    ) -> None:
        self.model = model
        self.system_prompt = system_prompt
        self.agent_name = agent_name
        self.max_output_tokens = max_output_tokens
        self.temperature = temperature
        self._heartbeat = heartbeat
        self._sessions_created: set[str] = set()

        if _runner is not None:
            # Test injection — skip building the real ADK objects.
            self._runner = _runner
        else:
            self._runner = self._build_runner()

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

    def _build_runner(self) -> Any:  # pragma: no cover — requires real ADK
        """Construct a Google ADK ``Runner`` with an ``LlmAgent``.

        Lazy-imports ``google.adk`` so the rest of the workspace runtime
        doesn't pull in google-adk on startup (it's only needed when this
        executor is actually instantiated by ``GoogleADKAdapter.create_executor``).
        """
        from google.adk.agents import LlmAgent
        from google.adk.runners import Runner
        from google.adk.sessions import InMemorySessionService

        agent = LlmAgent(
            name=self.agent_name,
            model=self.model,
            instruction=self.system_prompt or "",
        )

        session_service = InMemorySessionService()
        runner = Runner(
            agent=agent,
            app_name=self.agent_name,
            session_service=session_service,
        )
        return runner

    async def _ensure_session(self, session_id: str, user_id: str) -> None:
        """Create a session in the service if it doesn't exist yet."""
        if session_id in self._sessions_created:
            return
        session_service = self._runner.session_service
        existing = await session_service.get_session(
            app_name=self.agent_name,
            user_id=user_id,
            session_id=session_id,
        )
        if existing is None:
            await session_service.create_session(
                app_name=self.agent_name,
                user_id=user_id,
                session_id=session_id,
            )
        self._sessions_created.add(session_id)

    def _extract_text(self, context: RequestContext) -> str:
        """Pull plain text out of the A2A message parts."""
        from shared_runtime import extract_message_text
        return extract_message_text(context)

    def _build_content(self, user_text: str) -> Any:
        """Wrap user text in an ADK-compatible ``Content`` object."""
        from google.genai.types import Content, Part
        return Content(role="user", parts=[Part(text=user_text)])

    # ------------------------------------------------------------------
    # AgentExecutor interface
    # ------------------------------------------------------------------

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Run a single ADK turn and enqueue the reply as an A2A Message.

        Sequence:
        1. Extract user text from A2A message parts.
        2. Ensure an ADK session exists for this context_id.
        3. Call ``runner.run_async()`` and collect all response events.
        4. Concatenate final-response text; fall back to ``_NO_RESPONSE_MSG``
           when the model produces no output.
        5. Enqueue the reply via ``event_queue``.
        """
        user_text = self._extract_text(context)
        if not user_text:
            parts = getattr(getattr(context, "message", None), "parts", None)
            logger.warning("GoogleADKA2AExecutor: no text in message parts: %s", parts)
            await event_queue.enqueue_event(new_text_message(_NO_TEXT_MSG))
            return

        session_id = getattr(context, "context_id", None) or "default-session"
        user_id = "molecule-user"

        try:
            await self._ensure_session(session_id, user_id)

            content = self._build_content(user_text)
            response_parts: list[str] = []

            async for event in self._runner.run_async(
                session_id=session_id,
                user_id=user_id,
                new_message=content,
            ):
                # Collect text from final-response events
                if not getattr(event, "is_final_response", lambda: False)():
                    continue
                candidate_response = getattr(event, "response", None)
                if candidate_response is None:
                    continue
                for part in getattr(
                    getattr(candidate_response, "content", None) or MissingContent(),
                    "parts", []
                ):
                    text = getattr(part, "text", None)
                    if text:
                        response_parts.append(text)

            final_text = "".join(response_parts).strip() or _NO_RESPONSE_MSG
            await event_queue.enqueue_event(new_text_message(final_text))

        except Exception as exc:
            logger.error(
                "GoogleADKA2AExecutor: execution error [model=%s]: %s",
                self.model,
                type(exc).__name__,
                exc_info=True,
            )
            # Include exception detail (first ~1 KB) in the A2A error response so
            # callers get actionable context without needing workspace log access.
            # sanitize_agent_error scrubs API keys / bearer tokens before including
            # content in the response. Falls back to class-name-only when
            # the function is unavailable (standalone template repo layout).
            if sanitize_agent_error is not None:
                msg = sanitize_agent_error(stderr=str(exc))
            else:
                msg = f"Agent error: {type(exc).__name__}"
            await event_queue.enqueue_event(new_text_message(msg))

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Cancel a running task — emits canceled state per A2A protocol."""
        from a2a.types import TaskState, TaskStatus, TaskStatusUpdateEvent

        await event_queue.enqueue_event(
            TaskStatusUpdateEvent(
                status=TaskStatus(state=TaskState.TASK_STATE_CANCELED),
                final=True,
            )
        )


class MissingContent:
    """Sentinel to avoid AttributeError when response.content is None."""
    parts: list = []


# ---------------------------------------------------------------------------
# GoogleADKAdapter
# ---------------------------------------------------------------------------


class GoogleADKAdapter(BaseAdapter):
    """Molecule AI workspace adapter for Google ADK (google-adk v1.x).

    Implements the full ``BaseAdapter`` lifecycle:
    - ``setup()``           — validates config and runs ``_common_setup()``.
    - ``create_executor()`` — returns a ``GoogleADKA2AExecutor`` configured
                             from ``AdapterConfig``.
    """

    # Stored by setup(); consumed by create_executor()
    _setup_result: Any = None

    # ------------------------------------------------------------------
    # Identity
    # ------------------------------------------------------------------

    @staticmethod
    def name() -> str:
        """Runtime identifier — matches the ``runtime`` field in config.yaml."""
        return "google-adk"

    @staticmethod
    def display_name() -> str:
        """Human-readable name shown in the Molecule AI UI."""
        return "Google ADK"

    @staticmethod
    def description() -> str:
        """Short description of this adapter's capabilities."""
        return (
            "Google Agent Development Kit (ADK) adapter. "
            "Runs LLM agents via Google Gemini models using the official "
            "google-adk Python SDK (Apache-2.0)."
        )

    @staticmethod
    def get_config_schema() -> dict:
        """JSON Schema for runtime_config fields rendered in the Config tab."""
        return {
            "type": "object",
            "properties": {
                "agent_name": {
                    "type": "string",
                    "default": _DEFAULT_AGENT_NAME,
                    "description": "Internal ADK agent name",
                },
                "max_output_tokens": {
                    "type": "integer",
                    "default": _DEFAULT_MAX_OUTPUT_TOKENS,
                    "description": "Maximum output tokens for the Gemini model",
                },
                "temperature": {
                    "type": "number",
                    "default": _DEFAULT_TEMPERATURE,
                    "minimum": 0.0,
                    "maximum": 2.0,
                    "description": "Sampling temperature",
                },
            },
            "additionalProperties": False,
        }

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

    async def setup(self, config: AdapterConfig) -> None:
        """Validate config and run the shared platform setup pipeline.

        Raises ``RuntimeError`` if the required API key is not set and
        Vertex AI mode is not active.

        Args:
            config: ``AdapterConfig`` populated by the workspace runtime.
        """
        use_vertex = os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "").strip() in ("1", "true", "True")
        api_key = os.environ.get("GOOGLE_API_KEY", "").strip()

        if not use_vertex and not api_key:
            raise RuntimeError(
                "GoogleADKAdapter requires GOOGLE_API_KEY (for AI Studio) or "
                "GOOGLE_GENAI_USE_VERTEXAI=1 with GOOGLE_CLOUD_PROJECT set."
            )

        logger.info(
            "GoogleADKAdapter.setup: model=%s vertex=%s", config.model, use_vertex
        )

        self._setup_result = await self._common_setup(config)

    async def create_executor(self, config: AdapterConfig) -> GoogleADKA2AExecutor:
        """Build and return a ``GoogleADKA2AExecutor`` for A2A integration.

        Uses the system prompt assembled by ``_common_setup()`` in ``setup()``.
        Runtime-config keys ``agent_name``, ``max_output_tokens``, and
        ``temperature`` are respected when present.

        Args:
            config: ``AdapterConfig`` populated by the workspace runtime.

        Returns:
            A ready-to-use ``GoogleADKA2AExecutor`` instance.
        """
        rc = config.runtime_config or {}

        # Strip provider prefix from model, e.g. "google:gemini-2.0-flash" → "gemini-2.0-flash"
        model = config.model
        if ":" in model:
            model = model.split(":", 1)[1]

        system_prompt = (
            self._setup_result.system_prompt
            if self._setup_result is not None
            else config.system_prompt or ""
        )

        return GoogleADKA2AExecutor(
            model=model,
            system_prompt=system_prompt,
            agent_name=rc.get("agent_name", _DEFAULT_AGENT_NAME),
            max_output_tokens=int(rc.get("max_output_tokens", _DEFAULT_MAX_OUTPUT_TOKENS)),
            temperature=float(rc.get("temperature", _DEFAULT_TEMPERATURE)),
            heartbeat=config.heartbeat,
        )


# ---------------------------------------------------------------------------
# Module-level alias required by the adapter autodiscovery loader
# ---------------------------------------------------------------------------

Adapter = GoogleADKAdapter
