# platform/deskclaw/experiment/pipeline.py
# ═══════════════════════════════════════════════════════════════
# ExperimentPipeline — Matrix expansion and controlled execution
# ═══════════════════════════════════════════════════════════════

from __future__ import annotations

import asyncio
import json
import logging
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from deskclaw.adapters.protocol import AgentMessage
from deskclaw.adapters.registry import get_adapter
from deskclaw.experiment.models import (
    AgentBlueprint,
    ExperimentConfig,
    ExperimentError,
    ExperimentResult,
    RunResult,
)
from deskclaw.experiment.template_adapter import TemplateAdapter

logger = logging.getLogger("deskclaw.experiment.pipeline")


class PipelineError(ExperimentError):
    """ExperimentPipeline execution failed."""


class ExperimentPipeline:
    def __init__(
        self,
        config: ExperimentConfig,
        *,
        test_messages: list[AgentMessage] | None = None,
    ) -> None:
        self._config = config
        self._test_messages = test_messages or [
            AgentMessage(role="user", content="Hello, please respond.")
        ]

    @property
    def config(self) -> ExperimentConfig:
        return self._config

    def run(self) -> ExperimentResult:
        started_at = datetime.now(UTC).isoformat()
        experiment_id = (
            datetime.now(UTC).strftime("%Y%m%dT%H%M%S")
            + f"_{self._config.name}"
        )

        run_plan = self._expand_runs()
        logger.info(
            "Experiment '%s': %d total runs (%d blueprints × %d combos × %d reps)",
            self._config.name,
            len(run_plan),
            len(self._config.blueprints),
            len(self._config.matrix.expand()),
            self._config.repetitions,
        )

        adapters_to_close: list[Any] = []
        adapter_names = {bp.adapter_name for bp in self._config.blueprints}
        for name in adapter_names:
            try:
                adapter = get_adapter(name)
                if hasattr(adapter, "connect"):
                    asyncio.run(adapter.connect())
                    adapters_to_close.append(adapter)
                    logger.info("Connected adapter '%s' for experiment", name)
            except Exception as exc:
                logger.warning("Failed to pre-connect adapter '%s': %s", name, exc)

        try:
            runs: list[RunResult] = []
            for i, (blueprint, variables, rep) in enumerate(run_plan):
                logger.info(
                    "Run %d/%d: blueprint=%s, variables=%s, rep=%d",
                    i + 1,
                    len(run_plan),
                    blueprint.agent_name,
                    variables,
                    rep,
                )
                result = self._execute_single_run(blueprint, variables, rep)
                runs.append(result)
        finally:
            for adapter in adapters_to_close:
                try:
                    asyncio.run(adapter.close())
                    logger.info("Closed adapter '%s'", adapter.name())
                except Exception as exc:
                    logger.warning("Failed to close adapter: %s", exc)

        finished_at = datetime.now(UTC).isoformat()

        experiment_result = ExperimentResult(
            config=self._config,
            runs=runs,
            started_at=started_at,
            finished_at=finished_at,
            experiment_id=experiment_id,
        )

        self._save_results(experiment_result)
        return experiment_result

    def _expand_runs(self) -> list[tuple[AgentBlueprint, dict[str, Any], int]]:
        combos = self._config.matrix.expand()
        plan: list[tuple[AgentBlueprint, dict[str, Any], int]] = []
        for blueprint in self._config.blueprints:
            for variables in combos:
                for rep in range(self._config.repetitions):
                    plan.append((blueprint, variables, rep))
        return plan

    def _execute_single_run(
        self,
        blueprint: AgentBlueprint,
        variables: dict[str, Any],
        repetition: int,
    ) -> RunResult:
        try:
            inner_adapter = get_adapter(blueprint.adapter_name)
        except KeyError as exc:
            return RunResult(
                blueprint_name=blueprint.agent_name,
                variables=tuple(sorted(variables.items())),
                repetition=repetition,
                status="error",
                error=f"Adapter not found: {exc}",
            )

        effective_blueprint = self._apply_variable_overrides(blueprint, variables)
        template = TemplateAdapter(effective_blueprint, inner_adapter)

        start = time.monotonic()
        try:
            response = asyncio.run(
                asyncio.wait_for(
                    template.execute(
                        self._test_messages, timeout_s=self._config.timeout_s
                    ),
                    timeout=self._config.timeout_s,
                )
            )
            elapsed_ms = int((time.monotonic() - start) * 1000)

            usage = response.usage
            cost_info = template.cost_per_1k_tokens()
            input_cost = (usage.get("input_tokens", 0) / 1000) * cost_info.get(
                "input", 0.0
            )
            output_cost = (usage.get("output_tokens", 0) / 1000) * cost_info.get(
                "output", 0.0
            )

            return RunResult(
                blueprint_name=blueprint.agent_name,
                variables=tuple(sorted(variables.items())),
                repetition=repetition,
                status="passed",
                elapsed_ms=elapsed_ms,
                input_tokens=usage.get("input_tokens", 0),
                output_tokens=usage.get("output_tokens", 0),
                api_calls=usage.get("api_calls", 0),
                cost_usd=input_cost + output_cost,
                final_answer=response.final_answer,
                trace=tuple(sorted(response.trace.items())),
            )

        except TimeoutError:
            elapsed_ms = int((time.monotonic() - start) * 1000)
            return RunResult(
                blueprint_name=blueprint.agent_name,
                variables=tuple(sorted(variables.items())),
                repetition=repetition,
                status="timeout",
                elapsed_ms=elapsed_ms,
                error=f"Timeout after {self._config.timeout_s}s",
            )

        except Exception as exc:
            elapsed_ms = int((time.monotonic() - start) * 1000)
            return RunResult(
                blueprint_name=blueprint.agent_name,
                variables=tuple(sorted(variables.items())),
                repetition=repetition,
                status="error",
                elapsed_ms=elapsed_ms,
                error=f"{type(exc).__name__}: {exc}",
            )

    @staticmethod
    def _apply_variable_overrides(
        blueprint: AgentBlueprint, variables: dict[str, Any]
    ) -> AgentBlueprint:
        overrides: dict[str, Any] = {}
        extra_metadata: dict[str, Any] = {}
        bp_fields = {
            "temperature",
            "max_tokens",
            "system_prompt",
            "model_id",
        }
        for key, value in variables.items():
            if key in bp_fields:
                overrides[key] = value
            else:
                extra_metadata[key] = value

        if not overrides and not extra_metadata:
            return blueprint

        merged_metadata = dict(blueprint.metadata)
        merged_metadata.update(extra_metadata)

        return AgentBlueprint(
            agent_name=blueprint.agent_name,
            adapter_name=blueprint.adapter_name,
            model_id=overrides.get("model_id", blueprint.model_id),
            temperature=overrides.get("temperature", blueprint.temperature),
            max_tokens=overrides.get("max_tokens", blueprint.max_tokens),
            system_prompt=overrides.get("system_prompt", blueprint.system_prompt),
            capabilities=blueprint.capabilities,
            metadata=tuple(sorted(merged_metadata.items())),
        )

    def _save_results(self, result: ExperimentResult) -> None:
        output_dir = Path(self._config.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        results_file = output_dir / f"{result.experiment_id}.json"
        results_file.write_text(
            json.dumps(result.to_dict(), indent=2, ensure_ascii=False, default=str),
            encoding="utf-8",
        )
        logger.info("Experiment results saved to %s", results_file)
