#!/usr/bin/env python3
"""Reconstruct tuning cost and cost-to-oracle curves from existing runs.

The historical engine logs do not contain controller setup/cleanup timestamps.
Consequently the reported GPU cost is an engine-lifetime lower bound:
parallel_size * (last engine timestamp - first engine timestamp).  It must not
be presented as all-in method cost.  New experiments should record allocation
start/end directly.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import re
import subprocess
from datetime import datetime
from pathlib import Path
from typing import Any
from urllib.parse import urlsplit


SCHEMA = "aituner-tuning-cost-analysis-v1"
TIMESTAMP = re.compile(r"\b(\d{2}-\d{2} \d{2}:\d{2}:\d{2})(?:\.\d+)?\b")


def numeric_summary(values: list[float]) -> dict[str, Any]:
    values = [float(value) for value in values]
    return {
        "n": len(values),
        "min": min(values) if values else None,
        "max": max(values) if values else None,
        "distinct_n": len(set(values)),
    }


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


class Reader:
    def __init__(self, repo_root: Path):
        self.repo_root = repo_root

    def _split(self, locator: str) -> tuple[str | None, str]:
        if locator.startswith("ssh://"):
            parsed = urlsplit(locator)
            if not parsed.hostname or not parsed.path.startswith("/"):
                raise ValueError(f"invalid SSH locator: {locator}")
            return parsed.hostname, parsed.path
        path = Path(locator)
        if not path.is_absolute():
            path = self.repo_root / path
        return None, str(path)

    def read_text(self, locator: str) -> str:
        host, path = self._split(locator)
        if host is None:
            return Path(path).read_text(encoding="utf-8", errors="replace")
        completed = subprocess.run(
            ["ssh", host, "cat", "--", path],
            check=True,
            capture_output=True,
            text=True,
        )
        return completed.stdout

def join_locator(root: str, *parts: str) -> str:
    return "/".join([root.rstrip("/"), *(part.strip("/") for part in parts)])


def timestamp_span(log_text: str, year: int) -> tuple[float | None, bool, int]:
    parsed = [
        datetime.strptime(f"{year}-{match}", "%Y-%m-%d %H:%M:%S")
        for match in TIMESTAMP.findall(log_text)
    ]
    if not parsed:
        return None, True, 0
    monotonic = all(right >= left for left, right in zip(parsed, parsed[1:]))
    return (max(parsed) - min(parsed)).total_seconds(), monotonic, len(parsed)


def load_campaign(
    reader: Reader,
    root: str,
    year: int,
    missing_duration_s: dict[str, float] | None = None,
    missing_reason: dict[str, str] | None = None,
) -> dict[str, Any]:
    missing_duration_s = missing_duration_s or {}
    missing_reason = missing_reason or {}
    state_text = reader.read_text(join_locator(root, "state.json"))
    state = json.loads(state_text)
    trials = []
    for trial in state["trials"]:
        trial_id = str(trial["trial_id"])
        log_text = reader.read_text(join_locator(root, "trials", trial_id, "engine.log"))
        duration_s, monotonic, timestamp_n = timestamp_span(log_text, year)
        duration_source = "engine_log_span"
        duration_note = ""
        if duration_s is None:
            duration_s = float(missing_duration_s.get(trial_id, 0.0))
            duration_source = (
                "conservative_fallback" if trial_id in missing_duration_s else "no_engine_timestamp"
            )
            duration_note = missing_reason.get(trial_id, "")
        parallel_size = int(trial["parallel_size"])
        score = trial.get("best_request_rate_per_gpu")
        trials.append(
            {
                "trial_id": trial_id,
                "status": trial["status"],
                "failure_stage": trial.get("failure_stage", ""),
                "parallel_size": parallel_size,
                "score_req_s_per_gpu": None if score is None else float(score),
                "config_patch": trial["config_patch"],
                "duration_s": duration_s,
                "duration_source": duration_source,
                "duration_note": duration_note,
                "engine_timestamp_n": timestamp_n,
                "engine_timestamps_monotonic": monotonic,
                "engine_h20_hours_lower_bound": duration_s * parallel_size / 3600.0,
                "engine_log_sha256": sha256_text(log_text),
            }
        )
    return {
        "root": root,
        "state_sha256": sha256_text(state_text),
        "trials": trials,
    }


def fixed_task_context(spec: dict[str, Any]) -> dict[str, Any]:
    flags = dict(spec["engine"]["base_flags"])
    flags.pop("port", None)
    return {
        "model": spec["model"],
        "hardware": spec["hardware"],
        "engine_version": spec["engine"]["engine_version"],
        "launch_args": spec["engine"]["launch_args"],
        "base_flags_without_port": flags,
        "search": spec["search"],
        "slo": spec["slo"],
        "trace": spec["trace"],
    }


def regret(score: float, reference: float) -> float:
    if reference <= 0:
        raise ValueError("reference score must be positive")
    return max(0.0, 1.0 - float(score) / float(reference))


def sequential_curve(
    trials: list[dict[str, Any]], reference: float, thresholds: list[float]
) -> dict[str, Any]:
    cumulative_cost = 0.0
    best_score: float | None = None
    points = []
    for trial in trials:
        cumulative_cost += float(trial["engine_h20_hours_lower_bound"])
        score = trial["score_req_s_per_gpu"]
        if score is not None:
            best_score = score if best_score is None else max(best_score, score)
        points.append(
            {
                "trial_id": trial["trial_id"],
                "cumulative_engine_h20_hours_lower_bound": cumulative_cost,
                "best_score_req_s_per_gpu": best_score,
                "regret": None if best_score is None else regret(best_score, reference),
            }
        )
    hits = {}
    for threshold in thresholds:
        hit = next(
            (
                point
                for point in points
                if point["regret"] is not None
                and float(point["regret"]) <= float(threshold) + 1e-12
            ),
            None,
        )
        hits[f"regret_le_{threshold:g}"] = hit
    return {
        "reference_score_req_s_per_gpu": reference,
        "points": points,
        "cost_to_threshold": hits,
        "total_engine_h20_hours_lower_bound": cumulative_cost,
    }


def surface_cell(trial: dict[str, Any]) -> str:
    flags = trial["config_patch"]["flag_patch"]
    return f"tp{int(flags['tensor-parallel-size'])}_mns{int(flags['max-num-seqs'])}"


def tie_expanded_candidates(scores: dict[str, float], nominal_k: int) -> list[str]:
    ordered = sorted(scores, key=lambda cell: (-float(scores[cell]), cell))
    if nominal_k <= 0 or nominal_k > len(ordered):
        raise ValueError("nominal k outside score surface")
    cutoff = float(scores[ordered[nominal_k - 1]])
    tolerance = max(1e-12, abs(cutoff) * 1e-12)
    return [cell for cell in ordered if float(scores[cell]) >= cutoff - tolerance]


def real_final_policy(
    candidates: list[str], real_scores: dict[str, float], cell_costs: dict[str, float]
) -> dict[str, Any]:
    oracle = max(real_scores.values())
    selected = max(candidates, key=lambda cell: (real_scores[cell], cell))
    return {
        "candidate_cells": candidates,
        "real_evaluations": len(candidates),
        "selected_cell": selected,
        "selected_real_score_req_s_per_gpu": real_scores[selected],
        "real_regret": regret(real_scores[selected], oracle),
        "engine_h20_hours_lower_bound": sum(cell_costs[cell] for cell in candidates),
    }


def percentage_saving(new: float, old: float) -> float:
    if old <= 0:
        raise ValueError("baseline cost must be positive")
    return 1.0 - new / old


def build_analysis(manifest_path: Path) -> dict[str, Any]:
    repo_root = manifest_path.resolve().parents[2]
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    reader = Reader(repo_root)
    year = int(manifest["year"])

    sequential = {}
    task_contexts = {}
    for name, run in manifest["sequential_runs"].items():
        sequential[name] = load_campaign(
            reader,
            run["root"],
            year,
            run.get("missing_log_duration_s"),
            run.get("missing_log_reason"),
        )
        spec_text = reader.read_text(join_locator(run["root"], "study_spec.snapshot.json"))
        sequential[name]["study_spec_sha256"] = sha256_text(spec_text)
        task_contexts[name] = fixed_task_context(json.loads(spec_text))

    all_sequential_scores = [
        trial["score_req_s_per_gpu"]
        for campaign in sequential.values()
        for trial in campaign["trials"]
        if trial["score_req_s_per_gpu"] is not None
    ]
    empirical_reference = max(all_sequential_scores)
    thresholds = [float(value) for value in manifest["threshold_regrets"]]
    for campaign in sequential.values():
        campaign["curve"] = sequential_curve(
            campaign["trials"], empirical_reference, thresholds
        )

    surface_manifest = manifest["real_surface"]
    primary = load_campaign(reader, surface_manifest["primary_root"], year)
    companion = load_campaign(reader, surface_manifest["tp4_companion_root"], year)
    completed_surface_trials = [
        trial
        for trial in primary["trials"] + companion["trials"]
        if trial["status"] == "completed" and trial["score_req_s_per_gpu"] is not None
    ]
    cells: dict[str, dict[str, Any]] = {}
    for trial in completed_surface_trials:
        flags = trial["config_patch"]["flag_patch"]
        if int(flags["max-num-batched-tokens"]) != int(
            surface_manifest["fixed_max_num_batched_tokens"]
        ):
            raise ValueError("surface MBT invariant failed")
        cell = surface_cell(trial)
        if cell in cells:
            raise ValueError(f"duplicate completed surface cell: {cell}")
        cells[cell] = trial
    real_scores = {cell: float(trial["score_req_s_per_gpu"]) for cell, trial in cells.items()}
    cell_costs = {
        cell: float(trial["engine_h20_hours_lower_bound"]) for cell, trial in cells.items()
    }
    surface_oracle_score = max(real_scores.values())
    surface_oracle_cells = [
        cell for cell, score in real_scores.items() if math.isclose(score, surface_oracle_score)
    ]

    simulator_text = reader.read_text(manifest["simulator_metrics"])
    simulator = json.loads(simulator_text)
    simulator_real_scores = {
        cell: float(score) for cell, score in simulator["real_scores"].items()
    }
    surface_matches_simulator = set(real_scores) == set(simulator_real_scores) and all(
        math.isclose(real_scores[cell], simulator_real_scores[cell], abs_tol=1e-12)
        for cell in real_scores
    )
    throughput = simulator["analyses"]["frozen-calibrated/throughput-proxy"]
    throughput_scores = {
        cell: float(score) for cell, score in throughput["simulated_scores"].items()
    }
    throughput_real_final = {
        f"nominal_k_{k}": real_final_policy(
            tie_expanded_candidates(throughput_scores, k), real_scores, cell_costs
        )
        for k in (1, 2, 3)
    }
    throughput_top1 = tie_expanded_candidates(throughput_scores, 1)
    simulator_only_cell = throughput_top1[0]
    slo = simulator["analyses"]["frozen-calibrated/SLO-gated"]
    slo_top_bucket = list(slo["metrics"]["top1"]["candidate_cells"])
    slo_diagnostic = real_final_policy(slo_top_bucket, real_scores, cell_costs)

    pure_hits = sequential["pure_llm"]["curve"]["cost_to_threshold"]
    guided_hits = sequential["guided_harness"]["curve"]["cost_to_threshold"]
    direct_comparison = {}
    for threshold in (0.05, 0.02):
        key = f"regret_le_{threshold:g}"
        pure_hit = pure_hits[key]
        guided_hit = guided_hits[key]
        if pure_hit is None or guided_hit is None:
            continue
        pure_cost = float(pure_hit["cumulative_engine_h20_hours_lower_bound"])
        guided_cost = float(guided_hit["cumulative_engine_h20_hours_lower_bound"])
        direct_comparison[key] = {
            "pure_llm_h20_hours_lower_bound": pure_cost,
            "guided_harness_h20_hours_lower_bound": guided_cost,
            "guided_saving_vs_pure_llm": percentage_saving(guided_cost, pure_cost),
        }

    five_cost = direct_comparison["regret_le_0.05"][
        "guided_harness_h20_hours_lower_bound"
    ]
    two_cost = direct_comparison["regret_le_0.02"][
        "guided_harness_h20_hours_lower_bound"
    ]
    sim_real_cost = float(slo_diagnostic["engine_h20_hours_lower_bound"])
    target_bars = {
        "five_percent_regret": {
            "twenty_percent_below_current_guided": 0.8 * five_cost,
            "thirty_percent_below_posthoc_sim_slo_real_final": 0.7 * sim_real_cost,
            "development_target_h20_hours_lower_bound": min(
                0.8 * five_cost, 0.7 * sim_real_cost
            ),
        },
        "two_percent_regret": {
            "twenty_percent_below_current_guided": 0.8 * two_cost,
            "thirty_percent_below_posthoc_sim_slo_real_final": 0.7 * sim_real_cost,
            "development_target_h20_hours_lower_bound": min(
                0.8 * two_cost, 0.7 * sim_real_cost
            ),
        },
    }

    all_trial_costs = [
        float(trial["engine_h20_hours_lower_bound"])
        for campaign in sequential.values()
        for trial in campaign["trials"]
    ] + [
        float(trial["engine_h20_hours_lower_bound"])
        for campaign in (primary, companion)
        for trial in campaign["trials"]
    ]
    all_regrets = [
        float(point["regret"])
        for campaign in sequential.values()
        for point in campaign["curve"]["points"]
        if point["regret"] is not None
    ]
    monotonic_logs = all(
        trial["engine_timestamps_monotonic"]
        for campaign in [*sequential.values(), primary, companion]
        for trial in campaign["trials"]
    )
    invariants = {
        "dash0_task_contexts_equal_except_method_and_port": len(
            {json.dumps(value, sort_keys=True) for value in task_contexts.values()}
        )
        == 1,
        "surface_has_expected_cell_count": len(cells)
        == int(surface_manifest["expected_cells"]),
        "surface_matches_simulator_real_scores": surface_matches_simulator,
        "all_costs_non_negative": all(value >= 0 for value in all_trial_costs),
        "all_regrets_in_0_1": all(0 <= value <= 1 for value in all_regrets),
        "surface_scores_not_all_identical": len(set(real_scores.values())) > 1,
        "sequential_scores_not_all_identical": len(set(all_sequential_scores)) > 1,
        "engine_log_timestamps_monotonic": monotonic_logs,
        "sequential_trial_counts_match_manifest": all(
            len(sequential[name]["trials"]) == int(run["expected_trials"])
            for name, run in manifest["sequential_runs"].items()
        ),
        "simulator_suite_has_no_failed_runs": int(simulator["execution"]["failed_runs"])
        == 0,
        "simulator_scores_not_all_identical": len(set(throughput_scores.values())) > 1,
    }
    failed_invariants = [name for name, passed in invariants.items() if not passed]
    if failed_invariants:
        raise RuntimeError(f"data sanity invariant failed: {failed_invariants}")

    failed_primary_attempts = [
        trial for trial in primary["trials"] if trial["status"] != "completed"
    ]
    return {
        "schema": SCHEMA,
        "cost_definition": {
            "reported_metric": "engine H20-hours lower bound",
            "formula": "parallel_size * (last_engine_log_timestamp - first_engine_log_timestamp) / 3600",
            "included": ["engine startup after first timestamp", "warm-up/probes until last timestamp"],
            "not_reconstructable": [
                "GPU allocation before first engine timestamp",
                "controller/LLM latency",
                "cleanup after last engine timestamp",
                "one-time simulator operator profiling GPU-hours",
            ],
            "future_all_in_metric": "allocation_start_to_GPU_idle * allocated_GPU_count, including failures",
        },
        "comparison_scope": {
            "pure_llm_vs_guided_harness": "direct: same dash0 fixed task context",
            "simulator_vs_surface": "direct: simulator predictions and exact dash1 12-cell real surface",
            "dash0_methods_vs_dash1_simulator": "indicative only: matched model/engine/workload/GPU type, different host and campaign",
        },
        "empirical_reference": {
            "score_req_s_per_gpu": empirical_reference,
            "meaning": "best observed across the two dash0 sequential runs; not a global oracle",
        },
        "sequential_runs": sequential,
        "direct_dash0_comparison": direct_comparison,
        "real_surface": {
            "cells": {
                cell: {
                    "score_req_s_per_gpu": real_scores[cell],
                    "engine_h20_hours_lower_bound": cell_costs[cell],
                }
                for cell in sorted(cells)
            },
            "oracle_score_req_s_per_gpu": surface_oracle_score,
            "oracle_cells": surface_oracle_cells,
            "completed_annotation_engine_h20_hours_lower_bound": sum(cell_costs.values()),
            "failed_primary_attempt_n": len(failed_primary_attempts),
            "failed_primary_attempt_engine_h20_hours_lower_bound": sum(
                float(trial["engine_h20_hours_lower_bound"])
                for trial in failed_primary_attempts
            ),
            "primary_state_sha256": primary["state_sha256"],
            "tp4_companion_state_sha256": companion["state_sha256"],
        },
        "simulator": {
            "marginal_gpu_hours_without_real_verification": 0.0,
            "observed_fidelity_suite_cpu_hours": float(
                simulator["execution"]["suite_elapsed_seconds"]
            )
            / 3600.0,
            "observed_fidelity_suite_runs": int(simulator["execution"]["attempted_runs"]),
            "one_time_profile_gpu_hours": None,
            "one_time_profile_cost_status": "not recorded; total cold-start cost is unknown",
            "decision_bearing_throughput_proxy_sim_only_top1": {
                "selected_cell": simulator_only_cell,
                "selected_real_score_req_s_per_gpu": real_scores[simulator_only_cell],
                "real_regret": regret(real_scores[simulator_only_cell], surface_oracle_score),
                "gpu_hours": 0.0,
            },
            "decision_bearing_throughput_proxy_plus_real_final": throughput_real_final,
            "posthoc_slo_gated_plus_real_final": {
                **slo_diagnostic,
                "status": "diagnostic/post-hoc, not a preregistered prospective policy",
                "false_feasible": int(slo["false_feasibility"]["overall"]["false_feasible"]),
                "false_infeasible": int(slo["false_feasibility"]["overall"]["false_infeasible"]),
            },
            "metrics_sha256": sha256_text(simulator_text),
        },
        "provisional_development_targets": {
            "status": "lower-bound, single-task targets; require same-host prospective validation",
            **target_bars,
        },
        "data_sanity": {
            "invariants": invariants,
            "sequential_score_summary": numeric_summary(all_sequential_scores),
            "surface_score_summary": numeric_summary(list(real_scores.values())),
            "simulator_throughput_score_summary": numeric_summary(
                list(throughput_scores.values())
            ),
            "trial_cost_summary": numeric_summary(all_trial_costs),
            "regret_summary": numeric_summary(all_regrets),
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, default=Path(__file__).with_name("manifest.json"))
    parser.add_argument("--output", type=Path, default=Path(__file__).with_name("metrics.json"))
    args = parser.parse_args()
    analysis = build_analysis(args.manifest)
    args.output.write_text(json.dumps(analysis, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({
        "status": "ok",
        "output": str(args.output),
        "empirical_reference": analysis["empirical_reference"],
        "surface_oracle": analysis["real_surface"]["oracle_score_req_s_per_gpu"],
        "sanity": analysis["data_sanity"],
    }, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()
