import asyncio
from collections.abc import AsyncGenerator, AsyncIterator
from dataclasses import dataclass
from typing import Annotated, Any
from uuid import UUID, uuid4
import msgspec
from google.adk.agents import LlmAgent
from google.adk.agents.context_cache_config import ContextCacheConfig
from google.adk.agents.readonly_context import ReadonlyContext
from google.adk.apps import App
from google.adk.apps.app import EventsCompactionConfig
from google.adk.events import Event, EventActions
from google.adk.runners import Runner
from google.adk.tools import FunctionTool
from google.adk.tools.tool_context import ToolContext as AdkToolContext
from google.genai import types
from litestar import Controller, Litestar, post
from litestar.di import Provide
from litestar.exceptions import ServiceUnavailableException
from litestar.params import Dependency
from litestar.response import ServerSentEvent
from litestar.response.sse import ServerSentEventMessage
from litestar.status_codes import HTTP_200_OK
from sqlspec import SQLSpec
from sqlspec.adapters.asyncpg import AsyncpgConfig, AsyncpgPoolConfig
from sqlspec.adapters.asyncpg.adk import AsyncpgADKMemoryStore, AsyncpgADKStore
from sqlspec.extensions.adk import SQLSpecMemoryService, SQLSpecSessionService
from sqlspec.extensions.litestar import SQLSpecPlugin
RUN_CONTEXT_KEY = "temp:run_context_token"
DYNAMIC_INSTRUCTION_KEY = "temp:dynamic_instruction"
class ChatPromptRequest(msgspec.Struct, kw_only=True):
    """Incoming request payload for an agent conversation turn."""
    message: str
    session_id: str | None = None
    stream: bool = True
class ChatResponse(msgspec.Struct, kw_only=True):
    """Non-streaming response payload for an agent conversation turn."""
    session_id: str
    message: str
    finish_reason: str
    usage: dict[str, int]
class StreamEvent(msgspec.Struct, kw_only=True):
    """Normalized outward event emitted over SSE."""
    event_type: str
    data: dict[str, Any]
@dataclass(frozen=True, slots=True)
class DomainToolContext:
    """Request-scoped context bound for the duration of a single agent turn."""
    workspace_id: UUID
    user_id: str
class RunContextRegistry:
    """Store request-scoped tool contexts behind opaque single-turn tokens."""
    def __init__(self) -> None:
        self._contexts: dict[str, DomainToolContext] = {}
    def bind(self, context: DomainToolContext) -> str:
        """Register a tool context and return an opaque token for temp state."""
        token = uuid4().hex
        self._contexts[token] = context
        return token
    def resolve(self, token: str) -> DomainToolContext:
        """Resolve an active tool context from its opaque token."""
        context = self._contexts.get(token)
        if context is None:
            msg = "Agent tool context is missing or expired"
            raise RuntimeError(msg)
        return context
    def release(self, token: str) -> None:
        """Remove a token when the turn completes or disconnects."""
        self._contexts.pop(token, None)
def _dynamic_instruction(ctx: ReadonlyContext) -> str:
    """Read invocation-scoped context from temporary session state."""
    return str(ctx.state.get(DYNAMIC_INSTRUCTION_KEY, ""))
def build_bound_tools(registry: RunContextRegistry) -> list[FunctionTool]:
    """Create ADK FunctionTools that resolve request-scoped state via RunContextRegistry."""
    async def get_workspace_summary(
        environment: str,
        tool_context: AdkToolContext | None = None,
    ) -> dict[str, Any]:
        """Fetch workload summary metrics for the active workspace environment."""
        if tool_context is None:
            return {"error": {"code": "missing_context", "message": "Tool context unavailable."}}
        token = str(tool_context.state.get(RUN_CONTEXT_KEY, ""))
        domain_ctx = registry.resolve(token)
        return {
            "workspace_id": str(domain_ctx.workspace_id),
            "environment": environment,
            "status": "healthy",
        }
    return [FunctionTool(get_workspace_summary)]
class AgentRuntimeService:
    """Manage the singleton ADK Runner and stream turn execution."""
    def __init__(
        self,
        runner: Runner,
        session_service: SQLSpecSessionService,
        run_contexts: RunContextRegistry,
        app_name: str = "litestar_assistant",
    ) -> None:
        self._runner = runner
        self._session_service = session_service
        self._run_contexts = run_contexts
        self._app_name = app_name
    async def stream_turn(
        self,
        *,
        workspace_id: UUID,
        user_id: str,
        message: str,
        session_id: str | None = None,
    ) -> AsyncGenerator[StreamEvent, None]:
        """Execute one conversation turn and yield normalized StreamEvents."""
        resolved_id = session_id or uuid4().hex
        existing = await self._session_service.get_session(
            app_name=self._app_name,
            user_id=user_id,
            session_id=resolved_id,
        )
        if existing is None:
            await self._session_service.create_session(
                app_name=self._app_name,
                user_id=user_id,
                session_id=resolved_id,
                state={"workspace_id": str(workspace_id)},
            )
        token = self._run_contexts.bind(DomainToolContext(workspace_id=workspace_id, user_id=user_id))
        state_delta: dict[str, object] = {
            RUN_CONTEXT_KEY: token,
            DYNAMIC_INSTRUCTION_KEY: f"Active workspace: {workspace_id}",
        }
        content = types.Content(role="user", parts=[types.Part.from_text(text=message)])
        usage = {"prompt_tokens": 0, "output_tokens": 0, "total_tokens": 0, "cached_tokens": 0}
        try:
            yield StreamEvent(event_type="session", data={"session_id": resolved_id})
            async for event in self._runner.run_async(
                user_id=user_id,
                session_id=resolved_id,
                new_message=content,
                state_delta=state_delta,
            ):
                if event.usage_metadata is not None:
                    usage["prompt_tokens"] += int(event.usage_metadata.prompt_token_count or 0)
                    usage["output_tokens"] += int(event.usage_metadata.candidates_token_count or 0)
                    usage["total_tokens"] += int(event.usage_metadata.total_token_count or 0)
                    usage["cached_tokens"] += int(event.usage_metadata.cached_content_token_count or 0)
                if event.content and event.content.parts and event.author != "user":
                    text = "".join(
                        p.text or "" for p in event.content.parts if p.text and not getattr(p, "thought", False)
                    )
                    if text:
                        yield StreamEvent(event_type="delta", data={"text": text})
            yield StreamEvent(event_type="complete", data={**usage, "finish_reason": "STOP"})
        finally:
            self._run_contexts.release(token)
            latest = await self._session_service.get_session(
                app_name=self._app_name,
                user_id=user_id,
                session_id=resolved_id,
            )
            if latest is not None:
                await self._session_service.append_event(
                    session=latest,
                    event=Event(
                        author=self._app_name,
                        actions=EventActions(state_delta={"usage.last_turn": usage}),
                    ),
                )
class AgentChatController(Controller):
    """HTTP and SSE endpoints for multi-turn agent conversations."""
    path = "/api/v1/workspaces/{workspace_id:uuid}/agent"
    tags = ["Agent"]
    @post(path="/chat", status_code=HTTP_200_OK)
    async def chat(
        self,
        workspace_id: UUID,
        data: ChatPromptRequest,
        runtime_service: Annotated[AgentRuntimeService, Dependency(skip_validation=True)],
    ) -> ServerSentEvent | ChatResponse:
        """Stream an SSE conversation turn or return a consolidated JSON response."""
        if not data.message.strip():
            raise ServiceUnavailableException(detail="Prompt message cannot be empty.")
        user_id = f"{workspace_id}:default-user"
        if data.stream:
            async def sse_stream() -> AsyncIterator[ServerSentEventMessage]:
                events = runtime_service.stream_turn(
                    workspace_id=workspace_id,
                    user_id=user_id,
                    message=data.message,
                    session_id=data.session_id,
                )
                try:
                    async for item in events:
                        payload = msgspec.json.encode(item.data).decode("utf-8")
                        yield ServerSentEventMessage(event=item.event_type, data=payload)
                except asyncio.CancelledError:
                    raise
                finally:
                    await events.aclose()
            return ServerSentEvent(sse_stream(), headers={"Content-Type": "text/event-stream; charset=utf-8"})
        chunks: list[str] = []
        resolved_session = data.session_id or ""
        usage_totals: dict[str, int] = {}
        events = runtime_service.stream_turn(
            workspace_id=workspace_id,
            user_id=user_id,
            message=data.message,
            session_id=data.session_id,
        )
        try:
            async for item in events:
                if item.event_type == "session":
                    resolved_session = str(item.data.get("session_id", resolved_session))
                elif item.event_type == "delta":
                    chunks.append(str(item.data.get("text", "")))
                elif item.event_type == "complete":
                    usage_totals = {k: int(v) for k, v in item.data.items() if isinstance(v, int)}
        finally:
            await events.aclose()
        return ChatResponse(
            session_id=resolved_session,
            message="".join(chunks),
            finish_reason="STOP",
            usage=usage_totals,
        )
sql_runtime = SQLSpec()
db_config = AsyncpgConfig(
    connection_config=AsyncpgPoolConfig(dsn="postgresql://app:app@localhost:5432/app"),
    extension_config={"adk": {"session_table": "adk_session", "events_table": "adk_event"}},
)
sql_runtime.add_config(db_config)
run_context_registry = RunContextRegistry()
session_store = SQLSpecSessionService(AsyncpgADKStore(db_config))
memory_store = SQLSpecMemoryService(AsyncpgADKMemoryStore(db_config))
specialist_agent = LlmAgent(
    name="metrics_specialist",
    model="gemini-3.8-flash",
    description="Analyzes workload health and capacity metrics.",
    static_instruction="You are the metrics specialist. Answer using workspace telemetry tools.",
    instruction=_dynamic_instruction,
    tools=build_bound_tools(run_context_registry),
)
coordinator_agent = LlmAgent(
    name="workspace_coordinator",
    model="gemini-3.8-flash",
    description="Coordinates workspace analysis and delegates to specialists.",
    static_instruction="You coordinate workspace diagnostics and delegate metrics questions.",
    instruction=_dynamic_instruction,
    sub_agents=[specialist_agent],
)
adk_app = App(
    name="litestar_assistant",
    root_agent=coordinator_agent,
    context_cache_config=ContextCacheConfig(min_tokens=2048, ttl_seconds=1800, cache_intervals=10),
    events_compaction_config=EventsCompactionConfig(compaction_interval=10, overlap_size=2),
)
adk_runner = Runner(
    app=adk_app,
    session_service=session_store,
    memory_service=memory_store,
)
async def provide_runtime_service() -> AgentRuntimeService:
    """Provide the singleton AgentRuntimeService."""
    return AgentRuntimeService(
        runner=adk_runner,
        session_service=session_store,
        run_contexts=run_context_registry,
    )
app = Litestar(
    route_handlers=[AgentChatController],
    plugins=[SQLSpecPlugin(sql_runtime)],
    dependencies={"runtime_service": Provide(provide_runtime_service)},
)