# platform/deskclaw/experiment/template_adapter.py
# ═══════════════════════════════════════════════════════════════
# TemplateAdapter — Parametric adapter for controlled experiments
# ═══════════════════════════════════════════════════════════════

from __future__ import annotations

from collections.abc import AsyncIterator
from typing import Any

from deskclaw.adapters.protocol import (
    AgentAdapter,
    AgentCapability,
    AgentMessage,
    AgentResponse,
)
from deskclaw.experiment.models import AgentBlueprint


class TemplateAdapter(AgentAdapter):
    """Wraps an inner AgentAdapter, applying controlled variables from an AgentBlueprint."""

    def __init__(self, blueprint: AgentBlueprint, inner: AgentAdapter) -> None:
        self._blueprint = blueprint
        self._inner = inner

    @property
    def blueprint(self) -> AgentBlueprint:
        return self._blueprint

    def name(self) -> str:
        model_suffix = (
            f":{self._blueprint.model_id}" if self._blueprint.model_id else ""
        )
        return f"{self._blueprint.agent_name}{model_suffix}"

    def capabilities(self) -> set[AgentCapability]:
        if self._blueprint.capabilities:
            return {
                AgentCapability(c)
                for c in self._blueprint.capabilities
                if c in {cap.value for cap in AgentCapability}
            }
        return self._inner.capabilities()

    async def execute(
        self,
        messages: list[AgentMessage],
        *,
        max_tokens: int = 4096,
        temperature: float = 0.0,
        tools: list[dict] | None = None,
        timeout_s: int = 300,
    ) -> AgentResponse:
        controlled_messages = list(messages)
        if self._blueprint.system_prompt:
            has_system = any(m.role == "system" for m in controlled_messages)
            if has_system:
                controlled_messages = [
                    AgentMessage(
                        role=m.role,
                        content=f"{self._blueprint.system_prompt}\n\n{m.content}"
                        if m.role == "system"
                        else m.content,
                        tool_calls=m.tool_calls,
                        metadata=m.metadata,
                    )
                    for m in controlled_messages
                ]
            else:
                controlled_messages.insert(
                    0,
                    AgentMessage(role="system", content=self._blueprint.system_prompt),
                )

        if self._blueprint.metadata and hasattr(self._inner, "set_run_context"):
            self._inner.set_run_context(dict(self._blueprint.metadata))

        try:
            return await self._inner.execute(
                controlled_messages,
                max_tokens=self._blueprint.max_tokens,
                temperature=self._blueprint.temperature,
                tools=tools,
                timeout_s=timeout_s,
            )
        finally:
            if hasattr(self._inner, "clear_run_context"):
                self._inner.clear_run_context()

    async def stream(
        self,
        messages: list[AgentMessage],
        **kwargs: Any,
    ) -> AsyncIterator[AgentMessage]:
        async for msg in self._inner.stream(messages, **kwargs):
            yield msg

    def max_context_window(self) -> int:
        return self._inner.max_context_window()

    def supports_parallel_tool_calls(self) -> bool:
        return self._inner.supports_parallel_tool_calls()

    def cost_per_1k_tokens(self) -> dict[str, float]:
        return self._inner.cost_per_1k_tokens()
