"""Job execution utilities for the unified CLI."""

from __future__ import annotations

import asyncio
import contextlib
import logging
import shutil
from datetime import UTC, datetime
from pathlib import Path
from time import perf_counter, sleep
from typing import Any, Literal, Mapping, Sequence
from pydantic import BaseModel, Field, field_validator

from verifiers.types import GenerateOutputs
from verifiers.utils.eval_utils import run_evaluation

from medarc_verifiers.cli._constants import DEFAULT_ENDPOINTS_PATH
from medarc_verifiers.cli._eval_builder import build_client_config, build_eval_config
from medarc_verifiers.cli._job_builder import ResolvedJob
from medarc_verifiers.cli._manifest import RunManifest
from medarc_verifiers.cli._schemas import ModelConfigSchema
from medarc_verifiers.cli.utils.endpoint_utils import (
    EndpointRegistry,
    EndpointRegistryCache,
    EnvMetadataCache,
    load_endpoint_registry,
    load_env_metadata,
)
from medarc_verifiers.cli.utils.resume import (
    format_resume_mismatch_lines,
    is_resume_metadata_mismatch_error,
    is_valid_resume_results_path,
    load_resume_metadata_values,
)
from medarc_verifiers.cli.utils.shared import DEFAULT_BATCH_MAX_CONCURRENT, ensure_root_logging, resolve_env_identifier

try:
    from rich import print as rich_print  # type: ignore
except ImportError:  # pragma: no cover - rich is optional
    rich_print = None

logger = logging.getLogger(__name__)


class ExecutorSettings(BaseModel):
    """Run-level options controlling how jobs are executed."""

    run_id: str
    output_dir: Path
    env_dir: Path
    endpoints_path: Path | None = None
    endpoints_path_explicit: bool = False
    default_api_key_var: str
    default_api_key_var_explicit: bool = False
    default_api_base_url: str
    api_base_url_override: str | None = None
    log_level: str = "INFO"
    verbose: bool = False
    save_results: bool = True
    save_to_hf_hub: bool = False
    hf_hub_dataset_name: str | None = None
    max_concurrent_generation: int | None = None  # Deprecated; accepted for compatibility and ignored.
    max_concurrent_scoring: int | None = None  # Deprecated; accepted for compatibility and ignored.
    max_concurrent: int | None = None  # CLI override for max_concurrent
    http_max_retries: int | None = None  # CLI override for ClientConfig.max_retries
    rollout_max_retries: int = 0  # CLI override for EvalConfig.max_retries
    timeout: float | None = None
    sleep: float = 0.0
    dry_run: bool = False
    cli_env_args: dict[str, Any] | None = None
    cli_sampling_args: dict[str, Any] | None = None
    forced_job_ids: set[str] = Field(default_factory=set)

    @field_validator("output_dir", "env_dir", mode="before")
    @classmethod
    def _expand_path(cls, value: Path | str) -> Path:
        return Path(value).expanduser()

    @field_validator("endpoints_path", mode="before")
    @classmethod
    def _expand_optional_path(cls, value: Path | str | None) -> Path | None:
        if value is None:
            return None
        return Path(value).expanduser()


class JobExecutionResult(BaseModel):
    """Outcome emitted for each executed job."""

    job_id: str
    status: Literal["succeeded", "failed", "skipped"]
    error: str | None = None
    duration_seconds: float | None = None
    output_path: Path | None = None
    result: Any | None = None


def execute_jobs(
    jobs: Sequence[ResolvedJob],
    settings: ExecutorSettings,
    *,
    endpoints_cache: EndpointRegistryCache | None = None,
    env_metadata_cache: EnvMetadataCache | None = None,
    manifest: RunManifest | None = None,
) -> list[JobExecutionResult]:
    """Execute a sequence of resolved jobs."""
    ensure_root_logging(settings.log_level)
    logger.info("Starting run '%s' with %d job(s).", settings.run_id, len(jobs))

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

    job_statuses: dict[str, str] = {job.job_id: "pending" for job in jobs}
    results: list[JobExecutionResult] = []
    interrupted = False

    for index, job in enumerate(jobs):
        is_last_job = index == len(jobs) - 1
        env_identifier = resolve_env_identifier(job.env)
        model_identifier = job.model.id or job.model.model or job.job_id
        job_label = f"{job.job_id} (env={env_identifier}, model={model_identifier})"
        logger.info("Job %d/%d starting: %s", index + 1, len(jobs), job_label)
        job_dir = (run_dir / job.job_id).resolve()
        job_dir.mkdir(parents=True, exist_ok=True)
        job_statuses[job.job_id] = "running"
        forced_clean = job.job_id in settings.forced_job_ids

        if settings.dry_run:
            logger.info("Dry run enabled; skipping execution for job '%s'.", job.job_id)
            results.append(
                JobExecutionResult(
                    job_id=job.job_id,
                    status="skipped",
                    output_path=job_dir,
                )
            )
            job_statuses[job.job_id] = "skipped"
            _log_job_progress_window(jobs, index, job_statuses, event="dry-run skip")
            continue

        if manifest is not None:
            manifest.record_job_start(job.job_id)

        try:
            _prepare_job_dir_for_resume(job_id=job.job_id, job_dir=job_dir, forced_clean=forced_clean)
        except Exception as exc:  # noqa: BLE001
            error_message = f"{job_label} preflight failed: {exc}"
            logger.exception("%s", error_message)
            _record_job_failure(
                results=results,
                job_statuses=job_statuses,
                jobs=jobs,
                center_index=index,
                manifest=manifest,
                job_id=job.job_id,
                output_path=job_dir,
                error_message=error_message,
                manifest_error=str(exc),
                duration_seconds=None,
                status_label="failed",
                event="failure",
                note="during preflight",
            )
            _maybe_sleep_between_jobs(job, settings, is_last=is_last_job)
            continue

        try:
            endpoints = _load_endpoints_for_model(job.model, settings, cache=endpoints_cache)
            resolved_model, client_config, prime_sampling_overrides = build_client_config(
                job.model,
                endpoints=endpoints,
                default_api_key_var=settings.default_api_key_var,
                default_api_key_var_explicit=settings.default_api_key_var_explicit,
                default_api_base_url=settings.default_api_base_url,
                api_base_url_override=settings.api_base_url_override,
                http_max_retries_override=settings.http_max_retries,
                timeout_override=settings.timeout,
                headers=job.model.headers,
            )
            # Merge Prime Inference overrides with job sampling args (job args take precedence)
            merged_sampling_args = {**prime_sampling_overrides, **job.sampling_args}
            eval_config = build_eval_config(
                job_label=job.job_id,
                model_cfg=job.model,
                env_cfg=job.env,
                env_args=job.env_args,
                sampling_args=merged_sampling_args,
                cli_env_args=settings.cli_env_args,
                cli_sampling_args=settings.cli_sampling_args,
                resolved_model=resolved_model,
                client_config=client_config,
                env_dir=settings.env_dir,
                max_concurrent_override=settings.max_concurrent,
                max_concurrent_generation=settings.max_concurrent_generation,
                max_concurrent_scoring=settings.max_concurrent_scoring,
                rollout_max_retries=settings.rollout_max_retries,
                resume_path=job_dir,
                default_max_concurrent=DEFAULT_BATCH_MAX_CONCURRENT,
                save_results=settings.save_results,
                save_to_hf_hub=settings.save_to_hf_hub,
                hf_hub_dataset_name=settings.hf_hub_dataset_name,
                verbose=settings.verbose,
                env_metadata_cache=env_metadata_cache,
                env_metadata_loader=load_env_metadata,
                enforce_required_env_args=True,
            )
        except KeyboardInterrupt:
            logger.warning("Interrupted while preparing job %s.", job_label)
            interruption_message = f"{job_label} interrupted by user"
            _record_job_failure(
                results=results,
                job_statuses=job_statuses,
                jobs=jobs,
                center_index=index,
                manifest=manifest,
                job_id=job.job_id,
                output_path=job_dir,
                error_message=interruption_message,
                manifest_error="interrupted by user",
                duration_seconds=None,
                status_label="interrupted",
                event="interruption",
                note="during preparation",
            )
            interrupted = True
            break
        except Exception as exc:  # noqa: BLE001
            error_message = f"{job_label} preparation failed: {exc}"
            logger.exception("%s", error_message)
            _record_job_failure(
                results=results,
                job_statuses=job_statuses,
                jobs=jobs,
                center_index=index,
                manifest=manifest,
                job_id=job.job_id,
                output_path=job_dir,
                error_message=error_message,
                manifest_error=str(exc),
                duration_seconds=None,
                status_label="failed",
                event="failure",
                note="during preparation",
            )
            _maybe_sleep_between_jobs(job, settings, is_last=is_last_job)
            continue

        start = perf_counter()
        try:
            eval_result = asyncio.run(run_evaluation(eval_config))
        except KeyboardInterrupt:
            duration = perf_counter() - start
            logger.warning("Job %s interrupted by user after %.2fs.", job_label, duration)
            interruption_message = f"{job_label} interrupted by user"
            _record_job_failure(
                results=results,
                job_statuses=job_statuses,
                jobs=jobs,
                center_index=index,
                manifest=manifest,
                job_id=job.job_id,
                output_path=job_dir,
                error_message=interruption_message,
                manifest_error="interrupted by user",
                duration_seconds=duration,
                status_label="interrupted",
                event="interruption",
            )
            interrupted = True
            break
        except Exception as exc:  # noqa: BLE001
            duration = perf_counter() - start
            if is_resume_metadata_mismatch_error(exc):
                _log_resume_mismatch_diagnostics(job_id=job.job_id, resume_path=job_dir, eval_config=eval_config)
                prescriptive = (
                    "Job output dir contains incompatible prior results; "
                    "use --force to rerun cleanly or start a new run_id."
                )
                error_message = f"{job_label} failed after {duration:.2f}s: {prescriptive}"
                logger.error("%s", error_message)
                manifest_error = prescriptive
            else:
                error_message = f"{job_label} failed after {duration:.2f}s: {exc}"
                logger.exception("%s", error_message)
                manifest_error = str(exc)
            _record_job_failure(
                results=results,
                job_statuses=job_statuses,
                jobs=jobs,
                center_index=index,
                manifest=manifest,
                job_id=job.job_id,
                output_path=job_dir,
                error_message=error_message,
                manifest_error=manifest_error,
                duration_seconds=duration,
                status_label="failed",
                event="failure",
            )
            _maybe_sleep_between_jobs(job, settings, is_last=is_last_job)
            continue

        duration = perf_counter() - start
        logger.info("Job '%s' completed in %.2fs.", job.job_id, duration)

        _materialize_results(job_dir, eval_result)
        avg_reward = _extract_avg_reward(eval_result)
        metrics_avg = _extract_avg_metrics(eval_result)
        metadata = _safe_get(eval_result, "metadata", None)
        num_examples = _safe_get(metadata, "num_examples", None)
        rollouts_per_example = _safe_get(metadata, "rollouts_per_example", None)

        if manifest is not None:
            manifest.record_job_completion(
                job.job_id,
                duration_seconds=duration,
                results_dir=job_dir,
                avg_reward=avg_reward,
                metrics=metrics_avg,
                num_examples=num_examples,
                rollouts_per_example=rollouts_per_example,
            )

        results.append(
            JobExecutionResult(
                job_id=job.job_id,
                status="succeeded",
                duration_seconds=duration,
                output_path=job_dir,
                result=eval_result,
            )
        )
        job_statuses[job.job_id] = "completed"
        _log_job_progress_window(jobs, index, job_statuses, event="completion")
        _maybe_sleep_between_jobs(job, settings, is_last=is_last_job)

    if interrupted:
        logger.warning("Execution interrupted by user; %d job(s) left pending.", len(jobs) - len(results))

    return results


def _load_endpoints_for_model(
    model_cfg: ModelConfigSchema,
    settings: ExecutorSettings,
    *,
    cache: EndpointRegistryCache | None,
) -> EndpointRegistry:
    """Load the endpoint registry to use for a model."""
    registry_path = model_cfg.endpoints_path or settings.endpoints_path
    if registry_path is None:
        return {}

    registry_path_obj = Path(registry_path).expanduser()
    default_registry_path = Path(DEFAULT_ENDPOINTS_PATH).expanduser()
    explicit_path = bool(model_cfg.endpoints_path) or settings.endpoints_path_explicit

    if not registry_path_obj.exists():
        if explicit_path:
            raise FileNotFoundError(f"Endpoint registry not found at {registry_path_obj}")
        if _same_path(registry_path_obj, default_registry_path):
            logger.warning(
                "Default endpoints registry '%s' not found; continuing without endpoint aliases.",
                registry_path_obj,
            )
            return {}

    endpoints = load_endpoint_registry(registry_path_obj, cache=cache)
    if explicit_path and not endpoints:
        raise ValueError(f"Failed to load endpoint registry from explicit path '{registry_path_obj}'")
    return endpoints


def _record_job_failure(
    *,
    results: list[JobExecutionResult],
    job_statuses: dict[str, str],
    jobs: Sequence[ResolvedJob],
    center_index: int,
    manifest: RunManifest | None,
    job_id: str,
    output_path: Path,
    error_message: str,
    manifest_error: str,
    duration_seconds: float | None,
    status_label: str,
    event: str,
    note: str | None = None,
) -> None:
    if manifest is not None:
        manifest.record_job_failure(job_id, error=manifest_error, duration_seconds=duration_seconds)
    results.append(
        JobExecutionResult(
            job_id=job_id,
            status="failed",
            error=error_message,
            duration_seconds=duration_seconds,
            output_path=output_path,
        )
    )
    job_statuses[job_id] = status_label
    _log_job_progress_window(jobs, center_index, job_statuses, event=event, note=note)


def _safe_get(obj: Any, key: str, default: Any = None) -> Any:
    """Retrieve attribute or dict key, allowing newer dict-style GenerateOutputs."""
    if isinstance(obj, dict):
        return obj.get(key, default)
    return getattr(obj, key, default)


def _same_path(left: Path, right: Path) -> bool:
    try:
        return left.resolve(strict=False) == right.resolve(strict=False)
    except OSError:
        return left == right


def _materialize_results(job_dir: Path, results: GenerateOutputs) -> None:
    """Move evaluation artifacts into the job directory."""
    metadata = _safe_get(results, "metadata", None)
    raw_path = _safe_get(metadata, "path_to_save", None)
    src_path = Path(raw_path) if raw_path else job_dir
    try:
        resolved_src = src_path.resolve()
    except OSError:
        resolved_src = src_path
    try:
        resolved_job_dir = job_dir.resolve()
    except OSError:
        resolved_job_dir = job_dir

    if resolved_src == resolved_job_dir:
        logger.debug("Results already in job_dir; _materialize_results no-op for %s.", job_dir)
        return

    if src_path.exists() and resolved_src != resolved_job_dir:
        logger.warning(
            "Unexpected results source path for job '%s': src=%s, job_dir=%s. Materializing as a safety net.",
            job_dir.name,
            src_path,
            job_dir,
        )
        for item in src_path.iterdir():
            target = job_dir / item.name
            if target.exists():
                if target.is_dir():
                    shutil.rmtree(target)
                else:
                    target.unlink()
            shutil.move(str(item), target)
        with contextlib.suppress(OSError):
            src_path.rmdir()


def _prepare_job_dir_for_resume(*, job_id: str, job_dir: Path, forced_clean: bool) -> None:
    if not job_dir.exists():
        job_dir.mkdir(parents=True, exist_ok=True)
        return
    if not job_dir.is_dir():
        msg = f"Job output dir '{job_dir}' is not a directory. Use --force to rerun cleanly or choose a new run_id."
        raise ValueError(msg)
    if not any(job_dir.iterdir()):
        return

    if forced_clean:
        archive_path = _archive_job_dir(job_dir)
        logger.info("Forced rerun for job '%s': archived '%s' -> '%s'.", job_id, job_dir, archive_path)
        job_dir.mkdir(parents=True, exist_ok=True)
        return

    if is_valid_resume_results_path(job_dir):
        return

    msg = (
        f"Job output dir '{job_dir}' is non-empty but not a valid evaluation results path "
        "(expected results.jsonl and metadata.json). "
        "Use --force to rerun cleanly or choose a new run_id."
    )
    raise ValueError(msg)


def _archive_job_dir(job_dir: Path) -> Path:
    timestamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%S")
    candidate = job_dir.with_name(f"{job_dir.name}__old_{timestamp}")
    suffix = 1
    while candidate.exists():
        candidate = job_dir.with_name(f"{job_dir.name}__old_{timestamp}_{suffix}")
        suffix += 1
    job_dir.rename(candidate)
    return candidate


def _log_resume_mismatch_diagnostics(*, job_id: str, resume_path: Path, eval_config: Any) -> None:
    logger.error("Resume metadata mismatch for job '%s' at %s.", job_id, resume_path)
    saved_values = load_resume_metadata_values(resume_path)
    current_values = {
        "env_id": getattr(eval_config, "env_id", "<missing>"),
        "model": getattr(eval_config, "model", "<missing>"),
        "rollouts_per_example": getattr(eval_config, "rollouts_per_example", "<missing>"),
        "num_examples": getattr(eval_config, "num_examples", "<missing>"),
    }
    for line in format_resume_mismatch_lines(saved_values=saved_values, current_values=current_values):
        logger.error("  %s", line)
    logger.error("Resume supports increasing num_examples, but not decreasing it.")


def _extract_avg_reward(results: GenerateOutputs) -> float | None:
    """Return metadata-level average reward from GenerateOutputs."""
    metadata = _safe_get(results, "metadata", None)
    metadata_avg = _safe_get(metadata, "avg_reward", None)
    if metadata_avg is not None:
        return float(metadata_avg)
    return None


def _extract_avg_metrics(results: GenerateOutputs) -> dict[str, float]:
    """Return metadata-level average metrics from GenerateOutputs."""
    metadata = _safe_get(results, "metadata", None)
    raw_metrics = _safe_get(metadata, "avg_metrics", None)
    if not isinstance(raw_metrics, Mapping):
        return {}

    metrics: dict[str, float] = {}
    for key, value in raw_metrics.items():
        if value is None:
            continue
        try:
            metrics[str(key)] = float(value)
        except (TypeError, ValueError):
            continue
    return metrics


def _log_job_progress_window(
    jobs: Sequence[ResolvedJob],
    center_index: int,
    job_statuses: Mapping[str, str],
    *,
    event: str,
    note: str | None = None,
) -> None:
    if not jobs:
        return
    start = max(0, center_index - 1)
    end = min(len(jobs), center_index + 2)
    lines: list[str] = []
    header = "Segment | Job ID | Status | Model | Env | Name"
    divider = "-" * len(header)
    lines.append(header)
    lines.append(divider)
    for idx in range(start, end):
        job = jobs[idx]
        segment = "current" if idx == center_index else ("previous" if idx < center_index else "next")
        status = job_statuses.get(job.job_id, "pending")
        model_label = job.model.id or job.model.model or "-"
        try:
            env_label = resolve_env_identifier(job.env)
        except ValueError:
            env_label = job.env.id or job.job_id
        lines.append(
            f"{segment:8} | {job.job_id:20} | {status:10} | {model_label:15} | {env_label:20} | {job.name or '-'}"
        )
    label = f"Job progress after {event}"
    if note:
        label = f"{label} ({note})"
    logger.info("%s:\n%s", label, "\n".join(lines))


def _maybe_sleep_between_jobs(job: ResolvedJob, settings: ExecutorSettings, *, is_last: bool) -> None:
    """Optionally pause between jobs to spread out environment runs."""
    if settings.dry_run or is_last:
        return
    delay = job.sleep if job.sleep is not None else settings.sleep
    if delay is None or delay <= 0:
        return
    if rich_print:
        rich_print(f"[cyan]Sleeping {delay:.2f} second(s) before next job...[/cyan]")
    logger.info("Sleeping %.2f second(s) before next job...", delay)
    sleep(delay)


__all__ = ["ExecutorSettings", "JobExecutionResult", "execute_jobs"]
