"""Trial dispatch — execute trials via a Backend, handle sync/async routing.

Called by: eval.py
Calls: backend.run_trial(), scoring.extract_score(), scoring.validate_side_info()

Dispatch strategy (chosen by environment_type + n_workers):
  - Docker, single worker:  sequential sync — backend.run_trial()
  - Daytona / multi-worker: asyncio.gather() with semaphore — backend.run_trial_async()
"""

from __future__ import annotations

import asyncio
import logging
import time
from pathlib import Path

from daytona.common.errors import DaytonaError

from clawbench._env import PASS_THRESHOLD
from clawbench.backend import Backend, TrialResult
from clawbench.scoring import InvalidEvalOutput, extract_score, validate_side_info

logger = logging.getLogger(__name__)


def _process_result(
    result: TrialResult,
    task_name: str,
    duration: float,
    verbose: bool = False,
) -> dict:
    """Extract score/status/rewards from a TrialResult into a flat dict.

    Returns dict with keys: task_name, score, passed, duration_sec, rewards,
    status, error.
    """
    side_info = result.side_info
    warnings = validate_side_info(side_info, task_name)
    for w in warnings:
        logger.warning(f"Malformed side_info: {w}")

    feedback = side_info.get("Feedback", {}) if isinstance(side_info, dict) else {}
    status = feedback.get("Status", "unknown")
    rewards = feedback.get("Rewards", {})

    score = result.score
    if status == "completed" and rewards:
        score = extract_score(rewards)

    if verbose:
        print(f"    Status: {status}, score: {score}, duration: {duration:.1f}s", flush=True)
        if status != "completed":
            error_msg = feedback.get("Error", "")
            if error_msg:
                print(f"    Error: {error_msg}", flush=True)

    return {
        "task_name": task_name,
        "score": score,
        "passed": score > PASS_THRESHOLD,
        "duration_sec": round(duration, 1),
        "rewards": rewards if isinstance(rewards, dict) else {},
        "status": status,
        "error": None,
    }


def _run_one_sync(
    backend: Backend,
    candidate: dict,
    task: dict,
    *,
    verbose: bool = False,
    n_repeats: int = 1,
    skills_dir: Path | None = None,
    run_id: str | None = None,
) -> dict:
    """Execute a single trial synchronously and process the result."""
    task_name = task["task_name"]
    repeat = task.get("_repeat")

    if verbose:
        label = (
            f"{task_name} (repeat {repeat + 1}/{n_repeats})"
            if repeat is not None
            else task_name
        )
        print(f"  Evaluating: {label} ...", flush=True)

    t0 = time.time()
    try:
        result = backend.run_trial(candidate, task, skills_dir=skills_dir, run_id=run_id)
        duration = time.time() - t0
        return _process_result(result, task_name, duration, verbose=verbose)

    except InvalidEvalOutput:
        raise

    except DaytonaError as exc:
        duration = time.time() - t0
        if verbose:
            print(
                f"    Status: sandbox_error, score: nan, duration: {duration:.1f}s",
                flush=True,
            )
            print(f"    Error: {exc}", flush=True)
        logger.error(f"Sandbox error for {task_name}: {exc}")
        return {
            "task_name": task_name,
            "score": float("nan"),
            "passed": False,
            "duration_sec": round(duration, 1),
            "rewards": {},
            "status": "sandbox_error",
            "error": str(exc),
        }

    except Exception as exc:
        duration = time.time() - t0
        if verbose:
            print(
                f"    Status: trial_error, score: 0.0, duration: {duration:.1f}s",
                flush=True,
            )
            print(f"    Error: {exc}", flush=True)
        logger.error(f"Trial error for {task_name}: {exc}")
        return {
            "task_name": task_name,
            "score": 0.0,
            "passed": False,
            "duration_sec": round(duration, 1),
            "rewards": {},
            "status": "trial_error",
            "error": str(exc),
        }


async def _run_one_async(
    backend: Backend,
    candidate: dict,
    task: dict,
    sem: asyncio.Semaphore,
    *,
    verbose: bool = False,
    n_repeats: int = 1,
    skills_dir: Path | None = None,
    run_id: str | None = None,
) -> dict:
    """Execute a single trial asynchronously (with semaphore) and process the result."""
    async with sem:
        task_name = task["task_name"]
        repeat = task.get("_repeat")

        if verbose:
            label = (
                f"{task_name} (repeat {repeat + 1}/{n_repeats})"
                if repeat is not None
                else task_name
            )
            print(f"  Evaluating: {label} ...", flush=True)

        t0 = time.time()
        try:
            result = await backend.run_trial_async(
                candidate, task, skills_dir=skills_dir, run_id=run_id
            )
            duration = time.time() - t0
            return _process_result(result, task_name, duration, verbose=verbose)

        except InvalidEvalOutput:
            raise

        except DaytonaError as exc:
            duration = time.time() - t0
            if verbose:
                print(
                    f"    Status: sandbox_error, score: nan, duration: {duration:.1f}s",
                    flush=True,
                )
                print(f"    Error: {exc}", flush=True)
            logger.error(f"Sandbox error for {task_name}: {exc}")
            return {
                "task_name": task_name,
                "score": float("nan"),
                "passed": False,
                "duration_sec": round(duration, 1),
                "rewards": {},
                "status": "sandbox_error",
                "error": str(exc),
            }

        except Exception as exc:
            duration = time.time() - t0
            if verbose:
                print(
                    f"    Status: trial_error, score: 0.0, duration: {duration:.1f}s",
                    flush=True,
                )
                print(f"    Error: {exc}", flush=True)
            logger.error(f"Trial error for {task_name}: {exc}")
            return {
                "task_name": task_name,
                "score": 0.0,
                "passed": False,
                "duration_sec": round(duration, 1),
                "rewards": {},
                "status": "trial_error",
                "error": str(exc),
            }


def execute_trials(
    backend: Backend,
    candidate: dict,
    work_items: list[dict],
    *,
    environment_type: str = "docker",
    n_workers: int = 1,
    n_repeats: int = 1,
    verbose: bool = False,
    skills_dir: Path | None = None,
    run_id: str | None = None,
) -> list[dict]:
    """Execute a batch of trials, routing to sync or async dispatch.

    Args:
        backend: Initialized Backend instance (setup() already called).
        candidate: Candidate dict (e.g. {"skills": "..."}).
        work_items: Task dicts to evaluate (already expanded for repeats).
        environment_type: "docker" or "daytona".
        n_workers: Concurrency level for async dispatch.
        n_repeats: Total repeats (for verbose labeling only).
        verbose: Print per-task progress.
        skills_dir: Passed through to backend.run_trial().
        run_id: Passed through to backend.run_trial().

    Returns:
        List of result dicts, one per work item. Each dict has keys:
        task_name, score, passed, duration_sec, rewards, status, error.
    """
    common = dict(
        verbose=verbose,
        n_repeats=n_repeats,
        skills_dir=skills_dir,
        run_id=run_id,
    )

    if environment_type == "daytona":
        # Daytona requires a single persistent event loop — asyncio.run()
        # closes the loop after each call, breaking the SDK's internal
        # aiohttp session on subsequent trials. Use asyncio.gather() with
        # a semaphore to keep one loop alive across all trials.
        effective_workers = max(n_workers, 1)
        sem = asyncio.Semaphore(effective_workers)

        async def _run_all():
            coros = [
                _run_one_async(backend, candidate, t, sem, **common)
                for t in work_items
            ]
            return await asyncio.gather(*coros, return_exceptions=True)

        raw = asyncio.run(_run_all())
        results = []
        for r in raw:
            if isinstance(r, BaseException):
                logger.error(f"Async eval error: {r}")
            else:
                results.append(r)
        return results

    else:
        # Docker / local: sequential sync dispatch
        return [
            _run_one_sync(backend, candidate, t, **common)
            for t in work_items
        ]
