"""Optimization dispatcher — runs GEPA optimize_anything() loop.

Called by: cli.py
Calls: backend.BenchflowBackend, proposer.py, GEPA engine

Usage:
    clawbench optimize \
        --skill skills/gmail-agent-gws/SKILL.md \
        --tasks-dir tasks \
        --max-iter 20 --workers 10 -e daytona
"""

from __future__ import annotations

import hashlib
import json
import threading
from pathlib import Path

import click

from clawbench._env import DEFAULT_OPTIMIZER_MODEL, PASS_THRESHOLD
from clawbench.backend import BenchflowBackend
from clawbench.task_utils import discover_tasks
from clawbench.proposer import Proposer
from clawbench.task_contract import validate_tasks


def _create_fitness_fn(
    agent_name: str = "claude-code",
    model_name: str = "anthropic/claude-sonnet-4-6",
    trials_dir: Path | None = None,
    environment_type: str = "docker",
    trajectory_store: list | None = None,
    trajectory_lock: threading.Lock | None = None,
    agent_import_path: str | None = None,
    keep_trials: bool = False,
    skills_dir: Path | None = None,
    run_id: str | None = None,
) -> callable:
    """Return a GEPA-compatible evaluator wrapping a Backend.

    Returns a sync fitness function ``(candidate, example) → (score, side_info)``
    with an ``.async_fn`` attribute for the async variant.
    """
    backend_impl = BenchflowBackend(
        agent_name=agent_name,
        model_name=model_name,
        environment_type=environment_type,
        trials_dir=trials_dir,
        trajectory_store=trajectory_store,
        trajectory_lock=trajectory_lock,
        agent_import_path=agent_import_path,
        keep_trials=keep_trials,
    )
    backend_impl.setup()

    def fitness_fn(
        candidate: dict[str, str],
        example: dict | None = None,
    ) -> tuple[float, dict]:
        """Score a skill candidate by running a trial (sync)."""
        if example is None:
            return 0.0, {"Feedback": {"Status": "no_example"}}
        result = backend_impl.run_trial(
            candidate, example, skills_dir=skills_dir, run_id=run_id,
        )
        return result.score, result.side_info

    async def fitness_fn_async(
        candidate: dict[str, str],
        example: dict | None = None,
    ) -> tuple[float, dict]:
        """Score a skill candidate by running a trial (async)."""
        if example is None:
            return 0.0, {"Feedback": {"Status": "no_example"}}
        result = await backend_impl.run_trial_async(
            candidate, example, skills_dir=skills_dir, run_id=run_id,
        )
        return result.score, result.side_info

    fitness_fn.async_fn = fitness_fn_async
    return fitness_fn


def _make_split_block(task_list: list[dict], subscores: dict, split_name: str) -> dict:
    """Build a train/val block with summary for the dashboard templates."""
    results = []
    for task_idx, task in enumerate(task_list):
        score = subscores.get(task_idx, 0.0)
        results.append({
            "task_name": task["task_name"],
            "avg_reward": score,
            "passed": score > PASS_THRESHOLD,
            "rewards": [score],
        })

    n_passed = sum(1 for r in results if r["passed"])
    avg = sum(r["avg_reward"] for r in results) / len(results) if results else 0

    return {
        "avg_reward": avg,
        "pass_rate": n_passed / len(results) if results else 0,
        "summary": {
            "passed": n_passed,
            "total": len(results),
            "avg_reward": avg,
        },
        "results": results,
    }


def convert_gepa_to_dashboard(
    result, output_dir: Path, train: list[dict], val: list[dict],
    trajectory_store: list[dict] | None = None,
    iter_map: dict[int, int] | None = None,
    skipped_history: list[dict] | None = None,
):
    """Convert GEPA result to the dashboard's iterative format.

    Generates results.json, iter_N_results.json, and skill_iter_N/SKILL.md
    so the clawbench dashboard can display GEPA optimization runs.

    Args:
        iter_map: Optional mapping from candidate index to real GEPA iteration
            number. When provided, files use the real iteration numbers (with
            gaps for skipped iterations). When ``None``, falls back to
            sequential numbering (candidate 0 → iter 1, etc.).
        skipped_history: Optional list of history entries for iterations that
            were skipped at the subsample gate. Merged into the final history.
    """
    output_dir = Path(output_dir)
    n_candidates = len(result.candidates)
    best_idx = result.best_idx

    # Save baseline (seed) skill as skill_iter_0/SKILL.md (enables iter-1 diff)
    if n_candidates > 0:
        seed = result.candidates[0]
        seed_content = seed.get("skills", "") if isinstance(seed, dict) else str(seed)
        baseline_dir = output_dir / "skill_iter_0"
        baseline_dir.mkdir(parents=True, exist_ok=True)
        (baseline_dir / "SKILL.md").write_text(seed_content)

    history = []
    best_val_score = "0/0"
    best_iter_num = 0

    for i, candidate in enumerate(result.candidates):
        subscores = result.val_subscores[i] if i < len(result.val_subscores) else {}
        iter_num = iter_map[i] if iter_map and i in iter_map else i + 1
        # Skip seed (iteration 0) — already saved as skill_iter_0 above
        if iter_num == 0:
            continue

        # GEPA only exposes val subscores; train subscores are not tracked in
        # GEPAResult. Emit train as null so the dashboard skips it instead of
        # showing a duplicated/zeroed column.
        val_block = _make_split_block(val, subscores, "val")

        # Save iter_N_results.json
        (output_dir / f"iter_{iter_num}_results.json").write_text(
            json.dumps({"train": None, "val": val_block}, indent=2)
        )

        # Save skill_iter_N/SKILL.md
        skill_dir = output_dir / f"skill_iter_{iter_num}"
        skill_dir.mkdir(parents=True, exist_ok=True)
        skill_content = candidate.get("skills", "") if isinstance(candidate, dict) else str(candidate)
        (skill_dir / "SKILL.md").write_text(skill_content)

        # Save trajectory data for this iteration
        if trajectory_store:
            skill_content = candidate.get("skills", "") if isinstance(candidate, dict) else str(candidate)
            candidate_hash = hashlib.md5(skill_content.encode()).hexdigest()[:12]
            # Group trajectory entries by task name for this candidate.
            # GEPA evaluates each candidate on both train and val — the fitness
            # fn can't distinguish them, so all runs (train+val) are included.
            # Consume matched entries to prevent double-counting across iterations.
            iter_trajectories: dict[str, list] = {}
            matching = [t for t in trajectory_store if t["candidate_hash"] == candidate_hash]
            trajectory_store[:] = [t for t in trajectory_store if t["candidate_hash"] != candidate_hash]
            for t in matching:
                task_name = t["task_name"]
                if task_name not in iter_trajectories:
                    iter_trajectories[task_name] = []
                iter_trajectories[task_name].append({
                    "run": len(iter_trajectories[task_name]),
                    "score": t["score"],
                    "steps": t["steps"],
                    "summary": t["summary"],
                    "duration_sec": t["duration_sec"],
                })
            if iter_trajectories:
                (output_dir / f"iter_{iter_num}_trajectories.json").write_text(
                    json.dumps(iter_trajectories, indent=2)
                )

        # History entry
        val_score_str = f"{val_block['summary']['passed']}/{val_block['summary']['total']}"

        history.append({
            "iteration": iter_num,
            "train_score": None,
            "val_score": val_score_str,
            "train_avg_reward": None,
            "val_avg_reward": val_block["avg_reward"],
        })

        if i == best_idx:
            best_val_score = val_score_str
            best_iter_num = iter_num

    # Merge skipped iterations into history
    if skipped_history:
        history.extend(skipped_history)
    history.sort(key=lambda h: h["iteration"])

    # Total iterations = max iteration number seen
    total_iters = max((h["iteration"] for h in history), default=0)

    # Save results.json
    results = {
        "engine": "gepa",
        "iterations_run": total_iters,
        "iterations_evaluated": len(history),
        "best_iteration": best_iter_num if best_iter_num != 0 else (best_idx + 1),
        "best_avg_reward": result.val_aggregate_scores[best_idx] if result.val_aggregate_scores else 0,
        "best_train_score": None,
        "best_val_score": best_val_score,
        "total_metric_calls": result.total_metric_calls,
        "history": history,
    }
    (output_dir / "results.json").write_text(json.dumps(results, indent=2))


def run_optimization(
    skill_path: Path,
    train_tasks: list[dict] | None = None,
    val_tasks: list[dict] | None = None,
    test_tasks: list[dict] | None = None,
    tasks_dir: Path | None = None,
    agent_name: str = "claude-code",
    model_name: str = "anthropic/claude-sonnet-4-6",
    environment_type: str = "docker",
    max_iter: int = 20,
    n_workers: int = 1,
    output_dir: Path = Path(".local/clawbench-output"),
    seed: int = 42,
    agent_import_path: str | None = None,
    engine: str = "gepa",
    skills_dir: Path | None = None,
):
    """Run optimization loop.

    Args:
        skill_path: Path to the seed SKILL.md file.
        train_tasks: Training task dicts.
        val_tasks: Validation task dicts (optional, skip val when empty/None).
        test_tasks: Held-out test task dicts (optional).
        tasks_dir: Tasks directory to auto-discover from.
        agent_name: Agent to use (e.g. "claude-code", "openhands").
        model_name: Model for agent execution.
        environment_type: "docker" or "daytona".
        max_iter: Max optimization iterations (proposals).
        n_workers: Evaluation workers (must be 1 in v1).
        output_dir: Where to save results.
        seed: Random seed.
        agent_import_path: Import path for unregistered agents.
        engine: "gepa" (sole option, kept for future extensibility).

    Returns:
        GEPAResult.
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Resolve task lists
    if train_tasks is not None:
        val_tasks = val_tasks or []
        test_tasks = test_tasks or []
    elif tasks_dir is not None:
        train_tasks = discover_tasks(tasks_dir)
        val_tasks = val_tasks or []
        test_tasks = test_tasks or []
    else:
        raise ValueError("Provide train_tasks or tasks_dir")

    if not train_tasks:
        raise ValueError("No train tasks provided")

    # Pre-loop validation gate (both engines)
    all_tasks = train_tasks + val_tasks + test_tasks
    # Deduplicate for validation (tasks may appear in multiple splits)
    seen_names = set()
    unique_tasks = []
    for t in all_tasks:
        if t["task_name"] not in seen_names:
            seen_names.add(t["task_name"])
            unique_tasks.append(t)
    validate_tasks(unique_tasks, check_callable=False)

    return _run_gepa_engine(
        skill_path=skill_path,
        train_tasks=train_tasks,
        val_tasks=val_tasks,
        test_tasks=test_tasks,
        agent_name=agent_name,
        model_name=model_name,
        environment_type=environment_type,
        max_iter=max_iter,
        n_workers=n_workers,
        output_dir=output_dir,
        seed=seed,
        agent_import_path=agent_import_path,
        skills_dir=skills_dir,
    )


def _run_gepa_engine(
    skill_path: Path,
    train_tasks: list[dict],
    val_tasks: list[dict],
    test_tasks: list[dict],
    agent_name: str,
    model_name: str,
    environment_type: str,
    max_iter: int = 20,
    n_workers: int = 1,
    output_dir: Path = Path(".local/clawbench-output"),
    seed: int = 42,
    agent_import_path: str | None = None,
    skills_dir: Path | None = None,
):
    """Run GEPA optimize_anything() loop (existing path)."""
    try:
        from gepa.optimize_anything import optimize_anything, GEPAConfig, EngineConfig, ReflectionConfig, TrackingConfig
    except ImportError:
        raise click.ClickException(
            "--engine gepa requires the 'gepa' package (pip install gepa)."
        )

    # 2. Load seed skill
    skill_path = Path(skill_path)
    seed_skills = skill_path.read_text()
    seed_candidate = {"skills": seed_skills}

    # 3. Use pre-split task lists; replicate for single-task variance reduction
    train = list(train_tasks)
    test = test_tasks
    is_single_task = len(train) == 1

    if is_single_task:
        # Replicate the single task so GEPA evaluates it N times per minibatch.
        # Each replica is a fresh stochastic run → different trajectory and score.
        train = train * n_workers
        minibatch_size = n_workers
    else:
        if n_workers > len(train):
            raise click.BadParameter(
                f"--workers ({n_workers}) exceeds task count ({len(train)}). "
                f"Use --workers <= {len(train)} for multi-task optimization."
            )
        minibatch_size = n_workers

    val = val_tasks if val_tasks else train  # GEPA requires valset

    # Safety cap: max_metric_calls prevents runaway costs.
    safety_metric_cap = max_iter * minibatch_size * 4

    # 4. Create fitness function (with trajectory store for dashboard export)
    #
    # Daytona's aiohttp SDK uses asyncio locks that are bound to a single event
    # loop. GEPA dispatches evaluations via ThreadPoolExecutor, where each thread
    # calls asyncio.run() creating a *new* loop — causing "Lock bound to a
    # different event loop" errors.
    #
    # Fix: spin up one persistent event loop in a background thread and have the
    # sync wrapper submit coroutines to it via run_coroutine_threadsafe().
    trajectory_store: list[dict] = []
    fitness_fn = _create_fitness_fn(
        agent_name=agent_name,
        model_name=model_name,
        environment_type=environment_type,
        trajectory_store=trajectory_store,
        agent_import_path=agent_import_path,
        skills_dir=skills_dir,
    )

    _daytona_loop = None
    if environment_type != "docker" and n_workers > 1:
        import asyncio
        import threading

        _daytona_loop = asyncio.new_event_loop()
        _loop_thread = threading.Thread(target=_daytona_loop.run_forever, daemon=True)
        _loop_thread.start()

        _raw_async_fn = fitness_fn.async_fn

        def _daytona_safe_evaluator(candidate, example):
            """Submit async trial to the shared event loop (thread-safe)."""
            future = asyncio.run_coroutine_threadsafe(
                _raw_async_fn(candidate, example), _daytona_loop
            )
            return future.result()  # blocks this GEPA worker thread until done

        # Preserve the function signature GEPA expects
        fitness_fn = _daytona_safe_evaluator

    # 5. Create proposer (uses DEFAULT_OPTIMIZER_MODEL, not the agent model)
    proposer = Proposer(
        run_dir=output_dir / "proposals",
    )

    # 6. Configure GEPA
    config = GEPAConfig(
        engine=EngineConfig(
            run_dir=str(output_dir),
            max_candidate_proposals=max_iter,
            max_metric_calls=safety_metric_cap,
            seed=seed,
            display_progress_bar=True,
            candidate_selection_strategy="pareto",
            parallel=n_workers > 1,
            max_workers=n_workers,
        ),
        reflection=ReflectionConfig(
            reflection_lm=DEFAULT_OPTIMIZER_MODEL,
            reflection_minibatch_size=minibatch_size,
            custom_candidate_proposer=proposer,
            skip_perfect_score=True,
            perfect_score=1.0,
        ),
        tracking=TrackingConfig(use_wandb=False),
    )

    # 6b. Incremental dashboard callback — writes iter files as GEPA runs.
    #
    # optimize_anything() doesn't expose GEPAEngine's callbacks param, so we
    # monkey-patch GEPAEngine.__init__ to inject ours. The patch is scoped to
    # this call only and restored immediately after.
    class _IncrementalDashboardCallback:
        """Write dashboard artifacts after each GEPA iteration.

        Uses GEPA's real iteration numbers so skipped iterations appear as
        gaps (e.g. iter 1,2,3,5,8,10) and are recorded with a ``skipped``
        flag in history/results.json.
        """

        def __init__(self):
            self._gepa_iter: int | None = None  # current GEPA iteration (set by on_iteration_start)
            self._seed_saved = False
            self._history: list[dict] = []
            self._best_iter = 0
            self._best_avg = 0.0
            self._total_iters = 0  # total iterations attempted (incl. skipped)
            self._candidate_count = 0  # sequential count of evaluated candidates
            self.iter_map: dict[int, int] = {}  # candidate_index → GEPA iteration
            self.skipped_history: list[dict] = []  # skipped entries for post-hoc merge

        def on_iteration_start(self, event) -> None:
            self._gepa_iter = event["iteration"]

        def on_candidate_rejected(self, event) -> None:
            """Record skipped iterations so the dashboard shows them."""
            iter_num = event["iteration"]
            self._total_iters = max(self._total_iters, iter_num)
            entry = {
                "iteration": iter_num,
                "train_score": None,
                "val_score": None,
                "train_avg_reward": None,
                "val_avg_reward": None,
                "skipped": True,
                "reason": event.get("reason", "subsample gate"),
            }
            self._history.append(entry)
            self.skipped_history.append(entry)

            # Save the rejected candidate's skill text so the dashboard can
            # show what was proposed even for skipped iterations.
            proposals_dir = output_dir / "proposals"
            result_path = proposals_dir / f"proposal_{iter_num}_result_skills.md"
            if result_path.exists():
                skill_dir = output_dir / f"skill_iter_{iter_num}"
                skill_dir.mkdir(parents=True, exist_ok=True)
                (skill_dir / "SKILL.md").write_text(result_path.read_text())

            self._write_results_json()

        def on_valset_evaluated(self, event) -> None:
            # GEPA fires iteration=0 for the seed, iteration=1+ for real
            # iterations.  Use the event's iteration directly; fall back to
            # _gepa_iter (set by on_iteration_start) or sequential numbering.
            gepa_iter = event.get("iteration")
            if gepa_iter is not None:
                iter_num = gepa_iter
            elif self._gepa_iter is not None:
                iter_num = self._gepa_iter
            else:
                iter_num = len(self._history) + 1

            candidate = event["candidate"]
            scores = event["scores_by_val_id"]
            avg_score = event["average_score"]

            # Seed evaluation (iteration 0): save baseline files and add to
            # history so the dashboard always has at least one data point.
            if iter_num == 0:
                if not self._seed_saved:
                    self._seed_saved = True
                    seed_content = candidate.get("skills", "") if isinstance(candidate, dict) else str(candidate)
                    baseline_dir = output_dir / "skill_iter_0"
                    baseline_dir.mkdir(parents=True, exist_ok=True)
                    (baseline_dir / "SKILL.md").write_text(seed_content)
                # Save seed results
                val_block = _make_split_block(val, scores, "val")
                (output_dir / "iter_0_results.json").write_text(
                    json.dumps({"train": None, "val": val_block}, indent=2)
                )
                # Track as candidate 0 in iter_map for post-hoc conversion
                self.iter_map[self._candidate_count] = 0
                self._candidate_count += 1
                # Track best (seed can be best if no iteration beats it)
                if avg_score >= self._best_avg:
                    self._best_avg = avg_score
                    self._best_iter = 0
                # Add seed to history so chart always has a baseline
                passed = val_block["summary"]["passed"]
                total = val_block["summary"]["total"]
                self._history.append({
                    "iteration": 0,
                    "train_score": None,
                    "val_score": f"{passed}/{total}",
                    "train_avg_reward": None,
                    "val_avg_reward": avg_score,
                    "seed": True,
                })
                self._write_results_json()
                return

            self._total_iters = max(self._total_iters, iter_num)
            # Track candidate_index → GEPA iteration for post-hoc conversion
            self.iter_map[self._candidate_count] = iter_num
            self._candidate_count += 1

            # Save seed as skill_iter_0 on first valset evaluation (fallback
            # in case the seed event was missed)
            if not self._seed_saved:
                self._seed_saved = True
                seed_content = candidate.get("skills", "") if isinstance(candidate, dict) else str(candidate)
                baseline_dir = output_dir / "skill_iter_0"
                baseline_dir.mkdir(parents=True, exist_ok=True)
                (baseline_dir / "SKILL.md").write_text(seed_content)

            # Write iter_N_results.json
            val_block = _make_split_block(val, scores, "val")
            (output_dir / f"iter_{iter_num}_results.json").write_text(
                json.dumps({"train": None, "val": val_block}, indent=2)
            )

            # Write skill_iter_N/SKILL.md
            skill_content = candidate.get("skills", "") if isinstance(candidate, dict) else str(candidate)
            skill_dir = output_dir / f"skill_iter_{iter_num}"
            skill_dir.mkdir(parents=True, exist_ok=True)
            (skill_dir / "SKILL.md").write_text(skill_content)

            # Track best
            if avg_score >= self._best_avg:
                self._best_avg = avg_score
                self._best_iter = iter_num

            # History entry
            passed = val_block["summary"]["passed"]
            total = val_block["summary"]["total"]
            self._history.append({
                "iteration": iter_num,
                "train_score": None,
                "val_score": f"{passed}/{total}",
                "train_avg_reward": None,
                "val_avg_reward": avg_score,
            })

            self._write_results_json()

        def _write_results_json(self) -> None:
            """Persist current state to results.json."""
            # Sort history by iteration number for consistent display
            sorted_history = sorted(self._history, key=lambda h: h["iteration"])
            last = sorted_history[-1] if sorted_history else {}
            best_val_str = last.get("val_score", "0/0")
            # Find the best non-skipped val_score string
            for h in sorted_history:
                if h.get("iteration") == self._best_iter and h.get("val_score"):
                    best_val_str = h["val_score"]
                    break
            (output_dir / "results.json").write_text(json.dumps({
                "engine": "gepa",
                "status": "running",
                "iterations_run": self._total_iters,
                "iterations_evaluated": sum(1 for h in sorted_history if not h.get("skipped")),
                "best_iteration": self._best_iter,
                "best_avg_reward": self._best_avg,
                "best_train_score": None,
                "best_val_score": best_val_str,
                "total_metric_calls": 0,  # not available until GEPA finishes
                "history": sorted_history,
            }, indent=2))

    dashboard_cb = _IncrementalDashboardCallback()

    try:
        from gepa.core.engine import GEPAEngine
        _orig_init = GEPAEngine.__init__

        def _patched_init(self, *args, **kwargs):
            cbs = kwargs.get("callbacks") or []
            cbs.append(dashboard_cb)
            kwargs["callbacks"] = cbs
            _orig_init(self, *args, **kwargs)

        GEPAEngine.__init__ = _patched_init
        _patched = True
    except (ImportError, AttributeError):
        _patched = False

    # 7. Run optimization
    try:
        result = optimize_anything(
            seed_candidate=seed_candidate,
            evaluator=fitness_fn,
            dataset=train,
            valset=val,
            config=config,
        )
    finally:
        if _patched:
            GEPAEngine.__init__ = _orig_init

    # 8. Extract and save results
    best_skills = result.best_candidate
    if isinstance(best_skills, dict):
        best_skills = best_skills.get("skills", seed_skills)
    elif not isinstance(best_skills, str):
        best_skills = seed_skills
    (output_dir / "best_SKILL.md").write_text(best_skills)

    # Save optimization summary
    summary = {
        "best_idx": result.best_idx,
        "total_metric_calls": result.total_metric_calls,
        "num_candidates": len(result.candidates),
        "val_scores": result.val_aggregate_scores,
        "best_val_score": result.val_aggregate_scores[result.best_idx]
        if result.val_aggregate_scores
        else None,
        "train_tasks": [t["task_name"] for t in train],
        "val_tasks": [t["task_name"] for t in val],
        "test_tasks": [t["task_name"] for t in test],
        "config": {
            "agent_name": agent_name,
            "model_name": model_name,
            "environment_type": environment_type,
            "max_iter": max_iter,
            "minibatch_size": minibatch_size,
            "safety_metric_cap": safety_metric_cap,
            "n_workers": n_workers,
            "is_single_task": is_single_task,
        },
    }
    (output_dir / "optimization_summary.json").write_text(
        json.dumps(summary, indent=2)
    )

    # 9. Convert to dashboard format (including trajectories and skipped iters)
    convert_gepa_to_dashboard(
        result, output_dir, train, val, trajectory_store,
        iter_map=dashboard_cb.iter_map if dashboard_cb.iter_map else None,
        skipped_history=dashboard_cb.skipped_history if dashboard_cb.skipped_history else None,
    )

    # 10. Test evaluation on held-out set
    if test:
        test_scores = []
        for task in test:
            score, _ = fitness_fn(result.best_candidate, task)
            test_scores.append({"task_name": task["task_name"], "score": score})
        summary["test_scores"] = test_scores
        summary["test_avg_score"] = (
            sum(s["score"] for s in test_scores) / len(test_scores) if test_scores else 0
        )
        (output_dir / "optimization_summary.json").write_text(
            json.dumps(summary, indent=2)
        )

    return result
