from __future__ import annotations

import json
from dataclasses import replace
from pathlib import Path
from typing import Any, Mapping

from .spec import (
    CompareCandidateSpec,
    CompareSpec,
    ConfigPatch,
    SpecError,
    StudySpec,
    TrialSpec,
    load_structured_file,
    load_study_spec,
)
from .store import StudyStore
from .worker import run_trial


def load_compare_spec(path: Path) -> CompareSpec:
    return CompareSpec.from_dict(load_structured_file(path))


def run_compare(compare_spec_path: Path, *, output_root: Path | None = None) -> dict[str, Any]:
    compare_spec_path = compare_spec_path.resolve()
    compare = load_compare_spec(compare_spec_path)
    base_study_path = _resolve_path(compare.study_spec_path, base_dir=compare_spec_path.parent)
    study = load_study_spec(base_study_path)
    compare_root = (output_root or (Path(".aituner-compare") / compare.compare_id)).resolve()
    compare_root.mkdir(parents=True, exist_ok=True)
    StudyStore.write_json(compare_root / "compare_spec.snapshot.json", _compare_snapshot(compare))

    windows = _select_windows(compare, study=study, study_spec_path=base_study_path)
    baseline_patch, baseline_source = _resolve_candidate_config(
        compare.baseline,
        compare_spec_path=compare_spec_path,
    )
    tuned_patch, tuned_source = _resolve_candidate_config(
        compare.tuned,
        compare_spec_path=compare_spec_path,
    )

    per_window: list[dict[str, Any]] = []
    for window in windows:
        window_id = str(window["window_id"])
        row = {
            "window_id": window_id,
            "trace_type": window.get("trace_type"),
            "date": window.get("date"),
            "slot_token": window.get("slot_token"),
            "slot_label": window.get("slot_label"),
            "window_start": window.get("window_start"),
            "window_end": window.get("window_end"),
            "baseline": _run_compare_candidate(
                compare_root=compare_root,
                compare_id=compare.compare_id,
                study=study,
                study_spec_path=base_study_path,
                window_id=window_id,
                candidate_name="baseline",
                config_patch=baseline_patch,
                source=baseline_source,
            ),
            "tuned": _run_compare_candidate(
                compare_root=compare_root,
                compare_id=compare.compare_id,
                study=study,
                study_spec_path=base_study_path,
                window_id=window_id,
                candidate_name="tuned",
                config_patch=tuned_patch,
                source=tuned_source,
            ),
        }
        row["delta"] = _window_delta(row["baseline"], row["tuned"])
        per_window.append(row)

    summary = {
        "compare_id": compare.compare_id,
        "study_spec_path": str(base_study_path),
        "compare_root": str(compare_root),
        "baseline_source": baseline_source,
        "tuned_source": tuned_source,
        "windows": per_window,
        "aggregate": _aggregate_summary(per_window),
    }
    StudyStore.write_json(compare_root / "summary.json", summary)
    (compare_root / "report.md").write_text(_render_report(summary), encoding="utf-8")
    return summary


def _compare_snapshot(compare: CompareSpec) -> dict[str, Any]:
    return {
        "compare_id": compare.compare_id,
        "study_spec_path": compare.study_spec_path,
        "window_ids": compare.window_ids,
        "window_selector": (
            {
                "trace_type": compare.window_selector.trace_type,
                "date_prefix": compare.window_selector.date_prefix,
                "date_from": compare.window_selector.date_from,
                "date_to": compare.window_selector.date_to,
                "slot_token": compare.window_selector.slot_token,
            }
            if compare.window_selector is not None
            else None
        ),
        "baseline": _candidate_snapshot(compare.baseline),
        "tuned": _candidate_snapshot(compare.tuned),
    }


def _candidate_snapshot(candidate: CompareCandidateSpec) -> dict[str, Any]:
    if candidate.config_patch is not None:
        return {
            "kind": "config_patch",
            "config_patch": {
                "env_patch": dict(candidate.config_patch.env_patch),
                "flag_patch": dict(candidate.config_patch.flag_patch),
            },
        }
    assert candidate.trial_ref is not None
    return {
        "kind": "trial_ref",
        "trial_ref": {
            "study_root": candidate.trial_ref.study_root,
            "trial_id": candidate.trial_ref.trial_id,
        },
    }


def _resolve_path(raw_path: str, *, base_dir: Path) -> Path:
    path = Path(raw_path)
    if not path.is_absolute():
        path = (base_dir / path).resolve()
    return path


def _resolve_windows_path(study: StudySpec, *, study_spec_path: Path) -> Path:
    path = Path(study.trace.windows_path)
    if not path.is_absolute():
        path = (study_spec_path.parent / path).resolve()
    return path


def _load_windows_payload(study: StudySpec, *, study_spec_path: Path) -> list[dict[str, Any]]:
    windows_path = _resolve_windows_path(study, study_spec_path=study_spec_path)
    payload = json.loads(windows_path.read_text(encoding="utf-8"))
    raw_windows = payload.get("windows") if isinstance(payload, Mapping) else payload
    if not isinstance(raw_windows, list):
        raise SpecError(f"windows payload must contain a list: {windows_path}")
    return [{str(key): value for key, value in item.items()} for item in raw_windows if isinstance(item, Mapping)]


def _select_windows(compare: CompareSpec, *, study: StudySpec, study_spec_path: Path) -> list[dict[str, Any]]:
    windows = _load_windows_payload(study, study_spec_path=study_spec_path)
    indexed = {str(item.get("window_id") or "").strip(): item for item in windows}
    selected: dict[str, dict[str, Any]] = {}

    for window_id in compare.window_ids:
        item = indexed.get(window_id)
        if item is None:
            raise SpecError(f"window_id not found in windows payload: {window_id}")
        selected[window_id] = item

    selector = compare.window_selector
    if selector is not None:
        for item in windows:
            window_id = str(item.get("window_id") or "").strip()
            if not window_id:
                continue
            if selector.trace_type and str(item.get("trace_type") or "").strip() != selector.trace_type:
                continue
            date_value = str(item.get("date") or "").strip()
            if selector.date_prefix and not date_value.startswith(selector.date_prefix):
                continue
            if selector.date_from and date_value and date_value < selector.date_from:
                continue
            if selector.date_to and date_value and date_value > selector.date_to:
                continue
            if selector.slot_token and str(item.get("slot_token") or "").strip() != selector.slot_token:
                continue
            selected[window_id] = item

    ordered = sorted(
        selected.values(),
        key=lambda item: (
            str(item.get("date") or ""),
            str(item.get("slot_token") or ""),
            str(item.get("window_id") or ""),
        ),
    )
    if not ordered:
        raise SpecError("Compare spec selected zero windows.")
    return ordered


def _resolve_candidate_config(
    candidate: CompareCandidateSpec,
    *,
    compare_spec_path: Path,
) -> tuple[ConfigPatch, dict[str, Any]]:
    if candidate.config_patch is not None:
        return candidate.config_patch, {
            "kind": "config_patch",
            "config_patch": {
                "env_patch": dict(candidate.config_patch.env_patch),
                "flag_patch": dict(candidate.config_patch.flag_patch),
            },
        }
    assert candidate.trial_ref is not None
    study_root = _resolve_path(candidate.trial_ref.study_root, base_dir=compare_spec_path.parent)
    trial_spec_path = study_root / "trials" / candidate.trial_ref.trial_id / "trial_spec.json"
    if not trial_spec_path.exists():
        raise SpecError(f"trial_ref target not found: {trial_spec_path}")
    payload = json.loads(trial_spec_path.read_text(encoding="utf-8"))
    config_patch = ConfigPatch.from_dict(payload.get("config_patch") or {})
    return config_patch, {
        "kind": "trial_ref",
        "study_root": str(study_root),
        "trial_id": candidate.trial_ref.trial_id,
        "config_patch": {
            "env_patch": dict(config_patch.env_patch),
            "flag_patch": dict(config_patch.flag_patch),
        },
    }


def _run_compare_candidate(
    *,
    compare_root: Path,
    compare_id: str,
    study: StudySpec,
    study_spec_path: Path,
    window_id: str,
    candidate_name: str,
    config_patch: ConfigPatch,
    source: dict[str, Any],
) -> dict[str, Any]:
    run_root = compare_root / "runs" / window_id / candidate_name
    run_root.mkdir(parents=True, exist_ok=True)
    window_study = replace(study, trace=replace(study.trace, window_id=window_id))
    actual_study_path = run_root / "study_spec.json"
    source_path = run_root / "study_spec.source"
    actual_study_path.write_text(json.dumps(_study_to_snapshot(window_study), ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    source_path.write_text(str(actual_study_path) + "\n", encoding="utf-8")
    trial = TrialSpec(
        study_id=compare_id,
        trial_id=candidate_name,
        config_patch=config_patch,
        search=window_study.search,
        study_spec_path=str(source_path),
        artifact_dir=str(run_root),
        probe_log_path=str(run_root / "probe_history.json"),
        engine_log_path=str(run_root / "engine.log"),
        result_path=str(run_root / "result.json"),
    )
    trial_spec_path = run_root / "trial_spec.json"
    StudyStore.write_json(trial_spec_path, _trial_snapshot(trial))
    result = run_trial(trial_spec_path)
    parallel_size = _parallel_size_for_candidate(study=window_study, patch=config_patch)
    best_rate = result.get("best_request_rate")
    best_rate_per_gpu = (
        float(best_rate) / float(parallel_size)
        if isinstance(best_rate, (int, float)) and parallel_size > 0
        else None
    )
    return {
        "candidate": candidate_name,
        "source": source,
        "parallel_size": parallel_size,
        "config_patch": {
            "env_patch": dict(config_patch.env_patch),
            "flag_patch": dict(config_patch.flag_patch),
        },
        "status": result.get("status"),
        "best_sampling_u": result.get("best_sampling_u"),
        "best_request_rate": best_rate,
        "best_request_rate_per_gpu": best_rate_per_gpu,
        "best_pass_rate": result.get("best_pass_rate"),
        "best_request_count": result.get("best_request_count"),
        "failure_stage": result.get("failure_stage"),
        "failure_reason": result.get("failure_reason"),
        "artifact_dir": str(run_root),
        "result_path": str(run_root / "result.json"),
        "probe_log_path": str(run_root / "probe_history.json"),
        "engine_log_path": str(run_root / "engine.log"),
    }


def _study_to_snapshot(study: StudySpec) -> dict[str, Any]:
    from .spec import to_jsonable

    return to_jsonable(study)


def _trial_snapshot(trial: TrialSpec) -> dict[str, Any]:
    from .spec import to_jsonable

    return to_jsonable(trial)


def _parse_int_like(value: Any, *, default: int = 1) -> int:
    if value is None:
        return default
    if isinstance(value, bool):
        raise SpecError("Topology values must be integers.")
    if isinstance(value, int):
        return value
    if isinstance(value, float) and value.is_integer():
        return int(value)
    if isinstance(value, str) and value.strip():
        return int(value.strip())
    raise SpecError(f"Unable to parse integer topology value: {value!r}")


def _parallel_size_for_candidate(*, study: StudySpec, patch: ConfigPatch) -> int:
    flags = dict(study.engine.base_flags)
    flags.update(patch.flag_patch)
    tp = _parse_int_like(flags.get("tensor-parallel-size"), default=1)
    dp = _parse_int_like(flags.get("data-parallel-size"), default=1)
    return tp * dp


def _window_delta(baseline: dict[str, Any], tuned: dict[str, Any]) -> dict[str, Any]:
    return {
        "request_rate_delta": _numeric_delta(
            baseline.get("best_request_rate"), tuned.get("best_request_rate")
        ),
        "request_rate_per_gpu_delta": _numeric_delta(
            baseline.get("best_request_rate_per_gpu"),
            tuned.get("best_request_rate_per_gpu"),
        ),
        "pass_rate_delta": _numeric_delta(
            baseline.get("best_pass_rate"), tuned.get("best_pass_rate")
        ),
        "winner": _winner(
            baseline.get("best_request_rate_per_gpu"), tuned.get("best_request_rate_per_gpu")
        ),
    }


def _numeric_delta(lhs: Any, rhs: Any) -> float | None:
    if not isinstance(lhs, (int, float)) or not isinstance(rhs, (int, float)):
        return None
    return float(rhs) - float(lhs)


def _winner(baseline_per_gpu: Any, tuned_per_gpu: Any) -> str:
    if isinstance(baseline_per_gpu, (int, float)) and isinstance(tuned_per_gpu, (int, float)):
        if float(tuned_per_gpu) > float(baseline_per_gpu):
            return "tuned"
        if float(tuned_per_gpu) < float(baseline_per_gpu):
            return "baseline"
        return "tie"
    return "incomparable"


def _aggregate_summary(rows: list[dict[str, Any]]) -> dict[str, Any]:
    baseline_rates = [
        float(row["baseline"]["best_request_rate"])
        for row in rows
        if isinstance(row["baseline"].get("best_request_rate"), (int, float))
    ]
    tuned_rates = [
        float(row["tuned"]["best_request_rate"])
        for row in rows
        if isinstance(row["tuned"].get("best_request_rate"), (int, float))
    ]
    baseline_per_gpu = [
        float(row["baseline"]["best_request_rate_per_gpu"])
        for row in rows
        if isinstance(row["baseline"].get("best_request_rate_per_gpu"), (int, float))
    ]
    tuned_per_gpu = [
        float(row["tuned"]["best_request_rate_per_gpu"])
        for row in rows
        if isinstance(row["tuned"].get("best_request_rate_per_gpu"), (int, float))
    ]
    wins = {"baseline": 0, "tuned": 0, "tie": 0, "incomparable": 0}
    for row in rows:
        wins[row["delta"]["winner"]] += 1
    baseline_counts = _candidate_result_counts(rows, "baseline")
    tuned_counts = _candidate_result_counts(rows, "tuned")
    return {
        "window_count": len(rows),
        "wins": wins,
        "baseline_mean_request_rate": _mean_or_none(baseline_rates),
        "tuned_mean_request_rate": _mean_or_none(tuned_rates),
        "baseline_mean_request_rate_per_gpu": _mean_or_none(baseline_per_gpu),
        "tuned_mean_request_rate_per_gpu": _mean_or_none(tuned_per_gpu),
        "baseline_completed_window_count": baseline_counts["completed"],
        "baseline_failed_window_count": baseline_counts["failed"],
        "baseline_no_feasible_window_count": baseline_counts["no_feasible"],
        "tuned_completed_window_count": tuned_counts["completed"],
        "tuned_failed_window_count": tuned_counts["failed"],
        "tuned_no_feasible_window_count": tuned_counts["no_feasible"],
    }


def _candidate_result_counts(rows: list[dict[str, Any]], name: str) -> dict[str, int]:
    counts = {"completed": 0, "failed": 0, "no_feasible": 0}
    for row in rows:
        result = row.get(name)
        if not isinstance(result, dict):
            continue
        status = str(result.get("status") or "")
        if status == "completed":
            counts["completed"] += 1
        elif status == "failed":
            counts["failed"] += 1
        if not isinstance(result.get("best_request_rate_per_gpu"), (int, float)):
            counts["no_feasible"] += 1
    return counts


def _mean_or_none(values: list[float]) -> float | None:
    if not values:
        return None
    return sum(values) / len(values)


def _render_report(summary: dict[str, Any]) -> str:
    lines = [
        f"# {summary['compare_id']}",
        "",
        "## Setup",
        "",
        f"- Study spec: `{summary['study_spec_path']}`",
        f"- Compare root: `{summary['compare_root']}`",
        f"- Baseline source: `{summary['baseline_source']['kind']}`",
        f"- Tuned source: `{summary['tuned_source']['kind']}`",
        "",
        "## Aggregate",
        "",
        f"- Window count: `{summary['aggregate']['window_count']}`",
        f"- Wins: `{json.dumps(summary['aggregate']['wins'], ensure_ascii=False)}`",
        f"- Baseline mean request rate: `{summary['aggregate']['baseline_mean_request_rate']}`",
        f"- Tuned mean request rate: `{summary['aggregate']['tuned_mean_request_rate']}`",
        f"- Baseline mean request rate per GPU: `{summary['aggregate']['baseline_mean_request_rate_per_gpu']}`",
        f"- Tuned mean request rate per GPU: `{summary['aggregate']['tuned_mean_request_rate_per_gpu']}`",
        f"- Baseline completed/failed/no-feasible windows: `{summary['aggregate']['baseline_completed_window_count']}`/`{summary['aggregate']['baseline_failed_window_count']}`/`{summary['aggregate']['baseline_no_feasible_window_count']}`",
        f"- Tuned completed/failed/no-feasible windows: `{summary['aggregate']['tuned_completed_window_count']}`/`{summary['aggregate']['tuned_failed_window_count']}`/`{summary['aggregate']['tuned_no_feasible_window_count']}`",
        "",
        "## Per Window",
        "",
        "| Window | Date | Baseline req/s | Baseline req/s/gpu | Tuned req/s | Tuned req/s/gpu | Winner |",
        "| --- | --- | ---: | ---: | ---: | ---: | --- |",
    ]
    for row in summary["windows"]:
        baseline = row["baseline"]
        tuned = row["tuned"]
        lines.append(
            f"| `{row['window_id']}` | `{row.get('date') or ''}` | "
            f"`{baseline.get('best_request_rate')}` | `{baseline.get('best_request_rate_per_gpu')}` | "
            f"`{tuned.get('best_request_rate')}` | `{tuned.get('best_request_rate_per_gpu')}` | "
            f"`{row['delta']['winner']}` |"
        )
    lines.append("")
    return "\n".join(lines)
