"""Multi-model parallel runner."""

from __future__ import annotations

import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path
from typing import Callable

from .docker import run_container
from .eject import EjectPolicy
from .models import RunResult
from .preflight import preflight

RUNS_DIR = Path("runs")


def run_single(
    model_id: str,
    *,
    image_name: str,
    admin_key: str,
    or_key: str,
    budget_usd: float = 0.10,
    alias: str | None = None,
    extra_models: list[str] | None = None,
    runs_dir: Path = RUNS_DIR,
    run_id: str | None = None,
    **kwargs,
) -> RunResult:
    """Run a single model (convenience wrapper around run_multi).

    Useful for validation runs, quick tests, or model comparisons.

    Args:
        model_id: OpenRouter model ID.
        image_name: Docker image to run.
        admin_key: OpenRouter admin key.
        or_key: OpenRouter API key.
        budget_usd: Budget cap in USD (default 0.10).
        alias: Short name (defaults to last part of model_id).
        extra_models: Additional model IDs to validate during preflight.
        runs_dir: Root directory for artifacts.
        run_id: Override run ID.
        **kwargs: Forwarded to run_multi (env, mounts, on_event, etc.).

    Returns:
        :class:`~agent_preflight.models.RunResult` for the single model.
    """
    if alias is None:
        alias = model_id.split("/")[-1].replace(".", "-")[:16]

    preflight_kwargs = dict(kwargs.pop("preflight_kwargs", {}) or {})
    merged_extra_models = []
    if extra_models:
        merged_extra_models.extend(extra_models)

    pf_models = preflight_kwargs.pop("models", [])
    if isinstance(pf_models, str):
        pf_models = [pf_models]
    if pf_models:
        merged_extra_models.extend(pf_models)

    if merged_extra_models:
        preflight_kwargs["models"] = list(dict.fromkeys(merged_extra_models))

    results = run_multi(
        models=[(model_id, alias)],
        image_name=image_name,
        admin_key=admin_key,
        or_key=or_key,
        budget_usd=budget_usd,
        runs_dir=runs_dir,
        run_id=run_id,
        parallel=False,
        preflight_kwargs=preflight_kwargs or None,
        **kwargs,
    )
    return results[0]

# Thread-local log function so each container thread writes to its own file
_thread_log = threading.local()


def _tlog(*args: object, **kwargs: object) -> None:
    fn = getattr(_thread_log, "fn", print)
    fn(*args, **kwargs)


def _make_logger(log_path: Path, alias: str) -> Callable:
    lock = threading.Lock()

    def log(*args: object, **kwargs: object) -> None:
        msg = " ".join(str(a) for a in args)
        with lock:
            with open(log_path, "a") as f:
                f.write(msg + "\n")
            print(f"[{alias}] {msg}", flush=True)

    return log


def run_multi(
    models: list[tuple[str, str]],  # [(model_id, alias), ...]
    *,
    image_name: str,
    admin_key: str,
    or_key: str,
    budget_usd: float,
    runs_dir: Path = RUNS_DIR,
    run_id: str | None = None,
    env: dict[str, str] | None = None,
    mounts: list[tuple[str, str, str] | str] | None = None,
    initial_prompt: str | None = None,
    container_args: list[str] | None = None,
    eject_policy: EjectPolicy | None = None,
    on_event: Callable[[dict], None] | None = None,
    on_complete: Callable[[Path, RunResult], None] | None = None,
    budget_mode: str = "simple",
    idle_timeout_s: float = 300.0,
    parallel: bool = True,
    max_workers: int = 8,
    skip_preflight: bool = False,
    preflight_kwargs: dict | None = None,
) -> list[RunResult]:
    """Run multiple models in parallel Docker containers.

    Calls :func:`~agent_preflight.preflight.preflight` first (unless
    ``skip_preflight=True``), then launches one Docker container per model
    using a :class:`~concurrent.futures.ThreadPoolExecutor`.

    Args:
        models: List of ``(model_id, alias)`` pairs.
        image_name: Docker image to run.
        admin_key: OpenRouter admin key for provisioning sub-keys.
        or_key: Regular OpenRouter key for preflight balance/model checks.
        budget_usd: USD budget per model.
        runs_dir: Root directory for run artifacts.
        run_id: Override the auto-generated run ID.
        env: Extra env vars forwarded to every container.
        mounts: Volume mounts for every container as
            ``[(host, container, mode), ...]`` or Docker-style
            ``["host:container:mode", ...]`` strings.
        initial_prompt: JSON RPC initial prompt (for pi-mode agents).
        container_args: Extra args appended to ``docker run`` for every container.
        eject_policy: Shared eject policy (or None for defaults).
        on_event: Callback called with each parsed JSON event.
        on_complete: Callback ``(run_dir, result)`` after each model finishes.
        budget_mode: ``"simple"`` or ``"detailed"`` budget display.
        idle_timeout_s: Per-container idle timeout in seconds.
        parallel: Run all models concurrently (default True).
        max_workers: Thread pool size (capped at ``len(models)``).
        skip_preflight: Skip the preflight checks.
        preflight_kwargs: Extra kwargs forwarded to :func:`preflight`.

    Returns:
        List of :class:`~agent_preflight.models.RunResult`, one per model.
    """
    if run_id is None:
        run_id = f"run-{datetime.now().strftime('%Y%m%d-%H%M%S')}"

    run_dir = runs_dir / run_id
    run_dir.mkdir(parents=True, exist_ok=True)

    model_ids = [m for m, _ in models]

    if not skip_preflight:
        pf_kwargs = dict(preflight_kwargs or {})
        extra_models = pf_kwargs.pop("models", [])
        if isinstance(extra_models, str):
            extra_models = [extra_models]
        merged_models = list(dict.fromkeys(model_ids + list(extra_models)))
        preflight(
            admin_key=admin_key,
            or_key=or_key,
            models=merged_models,
            budget_per_model=budget_usd,
            docker_image=image_name,
            **pf_kwargs,
        )

    print(f"\nStarting {len(models)} model(s) in parallel.")
    print(f"Run dir : {run_dir}")
    print(f"Monitor : python -m agent_preflight status {run_id}\n")

    def _run_one(model_id: str, alias: str) -> RunResult:
        model_dir = run_dir / alias
        model_dir.mkdir(parents=True, exist_ok=True)
        log_fn = _make_logger(model_dir / "run.log", alias)
        _thread_log.fn = log_fn
        result = run_container(
            image_name=image_name,
            model_id=model_id,
            alias=alias,
            admin_key=admin_key,
            budget_usd=budget_usd,
            run_dir=run_dir,
            run_id=run_id,
            env=env,
            mounts=mounts,
            initial_prompt=initial_prompt,
            container_args=container_args,
            eject_policy=eject_policy,
            on_event=on_event,
            budget_mode=budget_mode,
            idle_timeout_s=idle_timeout_s,
            log_fn=log_fn,
        )
        if on_complete:
            on_complete(run_dir, result)
        return result

    results: list[RunResult] = []

    if parallel:
        workers = min(len(models), max_workers)
        with ThreadPoolExecutor(max_workers=workers) as pool:
            futures = {
                pool.submit(_run_one, mid, alias): alias
                for mid, alias in models
            }
            for future in as_completed(futures):
                alias = futures[future]
                try:
                    r = future.result()
                    results.append(r)
                except Exception as exc:
                    print(f"  [done] {alias}  EXCEPTION: {exc}")
    else:
        for model_id, alias in models:
            results.append(_run_one(model_id, alias))

    # Final summary table
    print(f"\n{'='*72}")
    print(f"  RUN COMPLETE: {run_id}")
    print(f"{'='*72}")
    print(f"  {'alias':<16}  {'cost':>7}  {'time':>7}  {'tools':>5}  {'cache':>6}  status")
    print(f"  {'-'*64}")
    total_cost = 0.0
    total_cached = 0
    total_prompt = 0
    for r in sorted(results, key=lambda x: x.alias):
        cost = r.cost_usd or 0.0
        total_cost += cost
        total_cached += r.cache.cached_tokens
        total_prompt += r.cache.total_prompt_tokens
        dur = f"{r.duration_s // 60}m{r.duration_s % 60:02d}s"
        cache_str = f"{r.cache.hit_rate * 100:.0f}%" if r.cache.request_count else "n/a"
        status = r.error or "ok"
        print(f"  {r.alias:<16}  ${cost:>6.3f}  {dur:>7}  {r.tool_calls:>5}  {cache_str:>6}  {status}")
    print(f"  {'TOTAL':<16}  ${total_cost:>6.3f}")
    if total_prompt > 0:
        agg_rate = total_cached / total_prompt * 100
        print(f"  Cache: {agg_rate:.1f}% aggregate hit rate "
              f"({total_cached:,} / {total_prompt:,} prompt tokens)")
    print(f"\nArtifacts: {run_dir}\n")

    return results
