"""Runtime backend — Backend protocol, TrialResult, and BenchflowBackend.

Called by: executor.py, eval.py, optimize.py

BenchflowBackend is the sole backend implementation. It uses benchflow's
ACP-native SDK for agent execution, skills deployment, trajectory capture,
and agent installation.
"""

from __future__ import annotations

import asyncio
import datetime
import hashlib
import json
import os
import shutil
import tempfile
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Protocol, runtime_checkable

import click

from clawbench._paths import _PROJECT_ROOT


# ---------------------------------------------------------------------------
# Protocol & data types
# ---------------------------------------------------------------------------


@dataclass
class TrialResult:
    """What every backend must return from a trial run.

    The backend is responsible for extracting the score from verifier rewards.
    Use scoring.extract_score() for the standard normalization.
    """

    score: float
    side_info: dict = field(default_factory=dict)
    trial_dir: Path | None = None


@runtime_checkable
class Backend(Protocol):
    """Interface for trial execution backends.

    Lifecycle: setup() → run_trial / run_trial_async (many) → cleanup()

    Both sync and async are required. The executor picks which to call
    based on --env and --workers:
      - Docker (single worker): sync path — run_trial()
      - Cloud envs / multi-worker: async path — run_trial_async()
    """

    def setup(self) -> None:
        """One-time initialization before trials (monkey-patches, SDK init, etc.)."""
        ...

    def cleanup(self) -> None:
        """Resource teardown after all trials complete."""
        ...

    def run_trial(
        self,
        candidate: dict[str, str],
        task: dict,
        *,
        skills_dir: Path | None = None,
        run_id: str | None = None,
    ) -> TrialResult:
        """Sync trial execution. Default: asyncio.run(run_trial_async(...))."""
        ...

    async def run_trial_async(
        self,
        candidate: dict[str, str],
        task: dict,
        *,
        skills_dir: Path | None = None,
        run_id: str | None = None,
    ) -> TrialResult:
        """Async trial execution. Backends must implement this."""
        ...


# ---------------------------------------------------------------------------
# Agent environment helpers
# ---------------------------------------------------------------------------


def build_agent_env(agent_name: str | None = None) -> dict[str, str]:
    """Forward provider credentials that installed agents expect.

    Passthrough strategy: forward every key that any supported agent
    might need. Keys absent from the host environment are silently
    skipped, so this is safe for all agents.
    """
    agent_env: dict[str, str] = {}

    passthrough_keys = (
        "ANTHROPIC_API_KEY",
        "GEMINI_API_KEY",
        "GOOGLE_API_KEY",
        "GOOGLE_APPLICATION_CREDENTIALS",
        "GOOGLE_CLOUD_PROJECT",
        "GOOGLE_CLOUD_LOCATION",
        "GOOGLE_GENAI_USE_VERTEXAI",
        "OPENAI_API_KEY",
        "OPENAI_BASE_URL",
        "OPENAI_MODEL",
        "OPENAI_ORG_ID",
        "OPENAI_PROJECT",
        "AZURE_OPENAI_API_KEY",
        "AZURE_OPENAI_ENDPOINT",
        "AZURE_API_KEY",
        "AZURE_RESOURCE_NAME",
    )

    for key in passthrough_keys:
        val = os.environ.get(key)
        if val:
            agent_env[key] = val

    # Azure fallback: many agents expect OPENAI_API_KEY even for Azure
    if "OPENAI_API_KEY" not in agent_env and agent_env.get("AZURE_OPENAI_API_KEY"):
        agent_env["OPENAI_API_KEY"] = agent_env["AZURE_OPENAI_API_KEY"]

    return agent_env


# ---------------------------------------------------------------------------
# BenchflowBackend
# ---------------------------------------------------------------------------


class BenchflowBackend:
    """Benchflow SDK backend — runs trials via ACP-native agent communication.

    Lifecycle:
        setup() — validate benchflow is importable
        run_trial() / run_trial_async() — execute a single trial
        cleanup() — no-op
    """

    def __init__(
        self,
        agent_name: str = "claude-agent-acp",
        model_name: str = "claude-haiku-4-5-20251001",
        environment_type: str = "docker",
        trials_dir: Path | None = None,
        trajectory_store: list | None = None,
        trajectory_lock=None,
        agent_import_path: str | None = None,
        keep_trials: bool = False,
    ):
        self.agent_name = agent_name
        self.model_name = model_name
        self.environment_type = environment_type
        self.trials_dir = trials_dir or (_PROJECT_ROOT / ".local" / "trials")
        self.trajectory_store = trajectory_store
        self.trajectory_lock = trajectory_lock
        self.agent_import_path = agent_import_path
        self.keep_trials = keep_trials
        self._sdk = None

    def setup(self) -> None:
        """Import and initialize benchflow SDK (idempotent)."""
        if self._sdk is not None:
            return
        try:
            from benchflow import SDK
            self._sdk = SDK()
        except ImportError:
            raise RuntimeError(
                "benchflow SDK not installed. Install with: "
                "pip install 'benchflow @ git+https://github.com/benchflow-ai/benchflow.git'"
            )

    def cleanup(self) -> None:
        """No-op — benchflow handles cleanup per-trial."""
        pass

    def run_trial(
        self,
        candidate: dict[str, str],
        task: dict,
        *,
        skills_dir: Path | None = None,
        run_id: str | None = None,
    ) -> TrialResult:
        """Sync trial execution — wraps async via asyncio.run()."""
        self.setup()
        return asyncio.run(
            self.run_trial_async(
                candidate, task, skills_dir=skills_dir, run_id=run_id,
            )
        )

    async def run_trial_async(
        self,
        candidate: dict[str, str],
        task: dict,
        *,
        skills_dir: Path | None = None,
        run_id: str | None = None,
    ) -> TrialResult:
        """Async trial execution via benchflow SDK.

        Maps the clawbench (candidate, task) → (score, side_info) contract
        onto benchflow's SDK.run() → RunResult.
        """
        self.setup()
        if task is None:
            return TrialResult(score=0.0, side_info={"Feedback": {"Status": "no_example"}})

        task_name = task.get("task_name", "?")
        task_path = task["task_path"]
        skills_content = candidate.get("skills", "")

        # Always copy task dir — SDK's context_root staging mutates in-place
        tmp_task = _copy_task(task_path, skills_content, run_id=run_id)

        try:
            # Resolve model name — benchflow uses bare model IDs,
            # clawbench/Harbor uses "provider/model" format
            model = self.model_name
            if model and "/" in model:
                model = model.split("/", 1)[1]

            # Build agent env — auto-inherit API keys
            agent_env = build_agent_env(self.agent_name)

            # Generate unique trial/job names
            trial_name = f"{task_name}__{uuid.uuid4().hex[:8]}"
            _ts = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            job_name = run_id or f"cb-{_ts}-{uuid.uuid4().hex[:8]}"
            self.trials_dir.mkdir(parents=True, exist_ok=True)
            jobs_dir = str(self.trials_dir)

            # Build pre-agent hooks for claw-* services
            pre_agent_hooks = _build_service_hooks(tmp_task)

            result = await self._sdk.run(
                task_path=tmp_task,
                agent=self.agent_name,
                model=model,
                agent_env=agent_env,
                job_name=job_name,
                trial_name=trial_name,
                jobs_dir=jobs_dir,
                environment=self.environment_type,
                skills_dir=str(skills_dir) if skills_dir else None,
                sandbox_user="agent" if self.agent_name != "oracle" else None,
                pre_agent_hooks=pre_agent_hooks or None,
                context_root=str(_PROJECT_ROOT),
            )
        finally:
            shutil.rmtree(tmp_task, ignore_errors=True)

        # Write metadata files for export.py compatibility
        trial_dir = self.trials_dir / job_name / trial_name
        if trial_dir.is_dir():
            (trial_dir / "task_name.txt").write_text(task_name + "\n")
            if run_id:
                (trial_dir / "run_id.txt").write_text(run_id + "\n")

            # Write agent/trajectory.json in ATIF format for export.py
            trajectory = result.trajectory or []
            if trajectory:
                agent_dir = trial_dir / "agent"
                agent_dir.mkdir(exist_ok=True)
                atif = _acp_to_atif(trajectory, self.agent_name, model or "")
                (agent_dir / "trajectory.json").write_text(
                    json.dumps(atif, indent=2) + "\n"
                )

        # Map RunResult → TrialResult with clawbench side_info contract
        return self._to_trial_result(result, task_name, skills_content)

    def _to_trial_result(self, result, task_name: str, skills_content: str) -> TrialResult:
        """Convert benchflow RunResult to clawbench TrialResult."""
        score = 0.0
        side_info: dict = {"Input": {"task_name": task_name}}

        # Honor verifier result even if agent timed out (verifier runs after timeout)
        if result.rewards:
            rewards = result.rewards
            raw = rewards.get("score", rewards.get("reward", 0.0))
            score = max(-1.0, min(1.0, float(raw)))
            side_info["Feedback"] = {"Status": "completed", "Rewards": rewards}
            side_info["scores"] = rewards
            if result.error:
                click.echo(f"  [{task_name}] WARN: agent error ({result.error}) but verifier returned score={score}", err=True)
        elif result.error:
            click.echo(f"  [{task_name}] ERROR: {result.error}", err=True)
            side_info["Feedback"] = {
                "Status": "trial_error",
                "Error": result.error,
            }
        else:
            click.echo(f"  [{task_name}] WARNING: no verifier result", err=True)
            side_info["Feedback"] = {"Status": "no_verifier_result"}

        # Build trajectory summary from ACP trajectory
        trajectory = result.trajectory or []
        trajectory_text = _format_acp_trajectory(trajectory)
        side_info["Generated Outputs"] = {
            "Agent Trace": trajectory_text,
        }

        # Timing
        timing = {}
        if result.started_at and result.finished_at:
            timing["total"] = round(
                (result.finished_at - result.started_at).total_seconds(), 1
            )
        side_info["Input"]["timing"] = timing
        if timing.get("total"):
            side_info["Input"]["duration_sec"] = timing["total"]

        if not result.error or score != 0.0:
            duration = f" ({timing['total']:.0f}s)" if timing.get("total") else ""
            click.echo(f"  [{task_name}] score={score}{duration}")

        # Save to trajectory store (for dashboard/proposer)
        if self.trajectory_store is not None:
            candidate_hash = hashlib.md5(skills_content.encode()).hexdigest()[:12]
            raw_steps = []
            for event in trajectory:
                etype = event.get("type", "")
                if etype == "tool_call":
                    raw_steps.append({
                        "type": "tool",
                        "tool": event.get("kind", ""),
                        "command": event.get("title", ""),
                        "content": json.dumps(event.get("content", []))[:300],
                    })
                elif etype == "agent_message":
                    raw_steps.append({
                        "type": "thought",
                        "content": event.get("text", "")[:500],
                    })

            entry = {
                "candidate_hash": candidate_hash,
                "task_name": task_name,
                "score": score,
                "steps": raw_steps,
                "summary": trajectory_text[:500] if trajectory_text else "",
                "duration_sec": timing.get("total", 0),
            }
            if self.trajectory_lock:
                with self.trajectory_lock:
                    self.trajectory_store.append(entry)
            else:
                self.trajectory_store.append(entry)

        return TrialResult(score=score, side_info=side_info)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _copy_task(task_path: str, skills_content: str = "", run_id: str | None = None) -> str:
    """Copy task dir to a temp location, optionally injecting skill content.

    Always copies so that SDK operations (context_root staging, skills_dir
    injection) don't mutate the original task directory.
    """
    tmpdir_name = f"bf{uuid.uuid4().hex[:12]}"
    tmpdir = os.path.join(tempfile.gettempdir(), tmpdir_name)
    os.makedirs(tmpdir)
    src = Path(task_path).resolve()
    dst = Path(tmpdir)

    for item in src.iterdir():
        dest = dst / item.name
        if item.is_dir():
            shutil.copytree(item, dest)
        else:
            shutil.copy2(item, dest)

    # Inject a per-eval-run cache-bust ARG so Daytona rebuilds images between
    # runs but reuses the same image for all repeats within a single run.
    dockerfile = dst / "environment" / "Dockerfile"
    if dockerfile.exists() and run_id:
        content = dockerfile.read_text()
        content = content.replace("FROM ", f"ARG EVAL_RUN={run_id}\nFROM ", 1)
        dockerfile.write_text(content)

    if skills_content:
        instruction_path = dst / "instruction.md"
        original = instruction_path.read_text() if instruction_path.exists() else ""
        instruction_path.write_text(
            f"# Agent Skill\n\n{skills_content}\n\n---\n\n# Task\n\n{original}"
        )

    return tmpdir


def _build_service_hooks(task_path: str) -> list:
    """Auto-detect claw-* services from Dockerfile and build pre-agent hooks.

    Reads config.toml for the service registry, scans the Dockerfile for
    service names, and returns async hook functions that start each matched
    service and wait for its health check.
    """
    import tomllib

    env_dir = Path(task_path) / "environment"
    dockerfile_path = env_dir / "Dockerfile"
    if not dockerfile_path.exists():
        return []
    # docker-compose tasks handle their own services
    if (env_dir / "docker-compose.yaml").exists():
        return []

    toml_path = _PROJECT_ROOT / "config.toml"
    if not toml_path.exists():
        return []

    with open(toml_path, "rb") as f:
        env_cfg = tomllib.load(f)

    service_registry: dict[str, tuple[str, int]] = {
        name: (cfg["db_path"], cfg["port"])
        for name, cfg in env_cfg.items()
        if name.startswith("claw-")
    }

    dockerfile_text = dockerfile_path.read_text()
    services = [
        (name, db_path, port)
        for name, (db_path, port) in service_registry.items()
        if name in dockerfile_text
    ]

    if not services:
        return []

    async def start_services(env):
        for cli_name, db_path, port in services:
            await env.exec(
                f"{cli_name} --db {db_path} serve --host 0.0.0.0 --port {port} --no-mcp &"
            )
        for cli_name, _db_path, port in services:
            await env.exec(
                f"for i in $(seq 1 30); do curl -sf http://localhost:{port}/health > /dev/null && break; sleep 1; done"
            )

    return [start_services]


def _acp_to_atif(events: list[dict], agent_name: str, model_name: str) -> dict:
    """Convert ACP trajectory events → ATIF-v1.6 format for export.py."""
    steps = []
    for i, event in enumerate(events, start=1):
        etype = event.get("type", "")
        step: dict = {"step_id": i, "source": "agent"}
        if etype == "tool_call":
            step["message"] = ""
            step["tool_calls"] = [{
                "tool_call_id": event.get("tool_call_id", ""),
                "function_name": event.get("kind", ""),
                "arguments": {"command": event.get("title", "")},
            }]
            content = event.get("content")
            if content:
                step["observation"] = {"results": [{"content": str(content)[:2000]}]}
        elif etype == "agent_message":
            step["message"] = event.get("text", "")
        elif etype == "agent_thought":
            step["source"] = "agent_thought"
            step["message"] = event.get("text", "")
        else:
            continue
        steps.append(step)

    return {
        "schema_version": "ATIF-v1.6",
        "agent": {"name": agent_name, "model_name": model_name},
        "steps": steps,
    }


def _format_acp_trajectory(trajectory: list[dict], max_chars: int = 8000) -> str:
    """Format ACP trajectory events into a human-readable summary."""
    if not trajectory:
        return "(no trajectory)"

    lines = []
    for event in trajectory:
        etype = event.get("type", "")
        if etype == "tool_call":
            kind = event.get("kind", "")
            title = event.get("title", "")
            status = event.get("status", "")
            lines.append(f"[{kind}] {title} → {status}")
        elif etype == "agent_message":
            text = event.get("text", "")
            lines.append(f"[Message]: {text[:300]}")
        elif etype == "agent_thought":
            text = event.get("text", "")
            lines.append(f"[Thought]: {text[:200]}")

    full = "\n".join(lines)
    if len(full) > max_chars:
        return full[:max_chars] + f"\n... ({len(lines)} events total, truncated)"
    return full + f"\n[Summary: {len(trajectory)} events]"
