from __future__ import annotations

import argparse
import hashlib
import json
import sys
from dataclasses import replace
from pathlib import Path
from typing import Any

from .compare import run_compare
from .config_signature import (
    materialized_effective_config_signature,
    tested_config_signature_index,
)
from .harness import (
    build_harness_context,
    build_harness_guided_proposal,
    build_harness_stop_proposal,
)
from .interaction_matrix import build_interaction_screening_matrix
from .job import append_job, build_trial_job
from .lca import (
    build_study_workload_profile,
    build_workload_profile,
    resolve_length_mode,
    similarity_report,
)
from .llm import (
    build_initial_config_review_prompt,
    build_prompt,
    call_llm_for_initial_config_review,
    call_llm_for_proposal,
    load_capability_profile,
    parse_initial_config_review_text,
    parse_proposal_text,
)
from .spec import (
    ConfigPatch,
    Proposal,
    SpecError,
    StudySpec,
    StudyState,
    load_structured_file,
    load_study_spec,
    to_jsonable,
)
from .store import StudyStore
from .trace import load_trace_requests, summarize_window
from .worker import run_trial


def _is_empty_config_patch(proposal: Proposal) -> bool:
    return not proposal.config_patch.env_patch and not proposal.config_patch.flag_patch


def _reject_repeated_effective_config(
    *,
    study: StudySpec,
    state: StudyState,
    proposal: Proposal,
    proposal_name: str,
) -> None:
    if proposal.should_stop:
        return
    tested = tested_config_signature_index(study, state)
    signature = materialized_effective_config_signature(
        study=study,
        state=state,
        proposal=proposal,
    )
    matching_trials = tested.get(signature)
    if not matching_trials:
        return
    raise SpecError(
        f"Proposal {proposal_name} repeats an already tested effective full config "
        "after materialization. "
        f"matching_trial_ids={matching_trials}. "
        "Choose a different eligible candidate or return should_stop=true."
    )


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


def _proposal_effective_config_fingerprint(
    *,
    study: StudySpec,
    state: StudyState,
    proposal: Proposal,
) -> str | None:
    if proposal.should_stop:
        return None
    signature = materialized_effective_config_signature(
        study=study,
        state=state,
        proposal=proposal,
    )
    return _effective_config_fingerprint(signature)


def _visible_harness_candidates(
    harness_context: dict[str, object] | None,
) -> list[dict[str, Any]]:
    if not isinstance(harness_context, dict):
        return []
    experiment_plan = harness_context.get("experiment_plan")
    if not isinstance(experiment_plan, dict):
        return []
    candidate_set = experiment_plan.get("candidate_set")
    if not isinstance(candidate_set, dict):
        return []
    candidates = candidate_set.get("eligible_candidates")
    if not isinstance(candidates, list):
        return []
    return [item for item in candidates if isinstance(item, dict)]


def _candidate_set_hash(harness_context: dict[str, object] | None) -> object:
    if not isinstance(harness_context, dict):
        return None
    experiment_plan = harness_context.get("experiment_plan")
    if not isinstance(experiment_plan, dict):
        return None
    candidate_set = experiment_plan.get("candidate_set")
    if not isinstance(candidate_set, dict):
        return None
    return candidate_set.get("candidate_set_hash")


def _match_visible_harness_candidate(
    *,
    proposal_fingerprint: str | None,
    harness_context: dict[str, object] | None,
) -> dict[str, Any] | None:
    if proposal_fingerprint is None:
        return None
    for candidate in _visible_harness_candidates(harness_context):
        if candidate.get("effective_config_fingerprint") == proposal_fingerprint:
            return candidate
    return None


def _proposal_source_label(
    *,
    proposal_name: str,
    proposal_source: Path | None,
) -> str:
    if proposal_name.startswith("baseline-"):
        return "baseline"
    if proposal_name.startswith("harness-stop-") or proposal_name.startswith("harness-proposal-"):
        return "harness"
    return str(proposal_source) if proposal_source else "llm"


def _classify_proposal_attribution(
    *,
    study: StudySpec,
    state: StudyState,
    proposal: Proposal,
    proposal_name: str,
    proposal_source: Path | None,
    harness_context: dict[str, object] | None,
) -> dict[str, Any]:
    source_label = _proposal_source_label(
        proposal_name=proposal_name,
        proposal_source=proposal_source,
    )
    fingerprint = _proposal_effective_config_fingerprint(
        study=study,
        state=state,
        proposal=proposal,
    )
    matched = _match_visible_harness_candidate(
        proposal_fingerprint=fingerprint,
        harness_context=harness_context,
    )
    matched_candidate_id = matched.get("candidate_id") if matched else None
    if proposal.should_stop:
        origin = (
            "harness_stop"
            if proposal_name.startswith("harness-stop-")
            else "proposal_file_stop"
            if proposal_source is not None
            else "llm_stop"
        )
    elif proposal_name.startswith("baseline-"):
        origin = "baseline"
    elif proposal_name.startswith("harness-proposal-"):
        origin = "harness_top1"
    elif proposal_source is not None:
        origin = "proposal_file"
    elif study.llm.use_harness and harness_context is not None:
        origin = "llm_selected_harness_candidate" if matched else "llm_out_of_set"
    else:
        origin = "llm_no_harness"
    return {
        "schema_version": 1,
        "proposal_name": proposal_name,
        "proposal_source": source_label,
        "proposal_origin": origin,
        "harness_candidate_policy": study.llm.harness_candidate_policy,
        "candidate_set_hash": _candidate_set_hash(harness_context),
        "proposal_effective_config_fingerprint": fingerprint,
        "matched_effective_config_signature": matched is not None,
        "matched_candidate_id": matched_candidate_id,
        "matched_candidate": _compact_candidate_for_attribution(matched),
    }


def _compact_candidate_for_attribution(candidate: dict[str, Any] | None) -> dict[str, Any] | None:
    if not candidate:
        return None
    return {
        "candidate_id": candidate.get("candidate_id"),
        "action_id": candidate.get("action_id"),
        "knob_family": candidate.get("knob_family"),
        "score": candidate.get("score"),
        "config_patch": candidate.get("config_patch"),
    }


def _validate_harness_candidate_policy(attribution: dict[str, Any]) -> None:
    if (
        attribution.get("harness_candidate_policy") == "strict"
        and attribution.get("proposal_origin") == "llm_out_of_set"
    ):
        raise SpecError(
            "LLM proposal is outside the visible harness eligible candidate set while "
            "llm.harness_candidate_policy=strict. Use an eligible harness candidate, "
            "switch to advisory mode, or disable harness context."
        )


def _config_patch_knob_keys(config_patch: object) -> set[str]:
    if not isinstance(config_patch, dict):
        return set()
    keys: set[str] = set()
    env_patch = config_patch.get("env_patch")
    if isinstance(env_patch, dict):
        keys.update(f"env:{key}" for key in env_patch)
    flag_patch = config_patch.get("flag_patch")
    if isinstance(flag_patch, dict):
        keys.update(f"flag:{key}" for key in flag_patch)
    return keys


def _proposal_knob_keys(proposal: Proposal) -> set[str]:
    return {
        *{f"env:{key}" for key in proposal.config_patch.env_patch},
        *{f"flag:{key}" for key in proposal.config_patch.flag_patch},
    }


def _nearest_visible_harness_candidates(
    *,
    proposal: Proposal,
    harness_context: dict[str, object] | None,
    limit: int = 3,
) -> list[dict[str, Any]]:
    proposal_keys = _proposal_knob_keys(proposal)
    scored: list[tuple[int, float, dict[str, Any]]] = []
    for candidate in _visible_harness_candidates(harness_context):
        candidate_keys = _config_patch_knob_keys(candidate.get("config_patch"))
        overlap = len(proposal_keys & candidate_keys)
        score = candidate.get("score")
        candidate_score = float(score) if isinstance(score, (int, float)) else 0.0
        scored.append((overlap, candidate_score, candidate))
    scored.sort(key=lambda item: (item[0], item[1]), reverse=True)
    return [_compact_candidate_for_attribution(item[2]) or {} for item in scored[:limit]]


def _candidate_family_gap_payload(
    *,
    study: StudySpec,
    trial_id: str,
    proposal_name: str,
    proposal: Proposal,
    attribution: dict[str, Any],
    harness_context: dict[str, object] | None,
    incumbent_rate_per_gpu: float,
    result_rate_per_gpu: float,
) -> dict[str, Any]:
    nearest = _nearest_visible_harness_candidates(
        proposal=proposal,
        harness_context=harness_context,
    )
    proposal_keys = _proposal_knob_keys(proposal)
    has_same_knob_candidate = any(
        proposal_keys
        & _config_patch_knob_keys(candidate.get("config_patch") if isinstance(candidate, dict) else None)
        for candidate in nearest
    )
    return {
        "schema_version": 1,
        "study_id": study.study_id,
        "trial_id": trial_id,
        "proposal_name": proposal_name,
        "proposal_origin": attribution.get("proposal_origin"),
        "gap_type": "same_operator_new_step" if has_same_knob_candidate else "missing_operator",
        "review_status": "pending",
        "incumbent_request_rate_per_gpu": incumbent_rate_per_gpu,
        "result_request_rate_per_gpu": result_rate_per_gpu,
        "absolute_gain": result_rate_per_gpu - incumbent_rate_per_gpu,
        "relative_gain": (
            (result_rate_per_gpu - incumbent_rate_per_gpu) / incumbent_rate_per_gpu
            if incumbent_rate_per_gpu > 0
            else None
        ),
        "proposal_patch": {
            "env_patch": dict(proposal.config_patch.env_patch),
            "flag_patch": dict(proposal.config_patch.flag_patch),
        },
        "changed_knobs": sorted(proposal_keys),
        "candidate_set_hash": attribution.get("candidate_set_hash"),
        "nearest_harness_candidates": nearest,
        "interpretation": (
            "LLM changed a knob already present in the visible harness candidate set; "
            "treat this as a step-size/acquisition gap until offline review accepts "
            "a descriptor or operator change."
            if has_same_knob_candidate
            else "LLM changed knobs not represented by the visible harness candidates; "
            "offline review should decide whether this is a missing operator, "
            "descriptor, or mechanism."
        ),
    }


def _result_request_rate_per_gpu(result: dict[str, object], parallel_size: int | None) -> float | None:
    best_request_rate = result.get("best_request_rate")
    if not isinstance(best_request_rate, (int, float)):
        return None
    if not isinstance(parallel_size, int) or parallel_size <= 0:
        return None
    return float(best_request_rate) / float(parallel_size)


def _parse_parallel_int(value: object, *, default: int = 1) -> int:
    if value is None:
        return default
    if isinstance(value, bool):
        raise SpecError("Boolean values are not valid topology settings.")
    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 topology setting from {value!r}.")


def _parallel_size_for_config_patch(study: StudySpec, config_patch: object) -> int | None:
    if not isinstance(config_patch, ConfigPatch):
        return None
    flags: dict[str, object] = dict(study.engine.base_flags)
    flags.update(config_patch.flag_patch)
    tp = _parse_parallel_int(flags.get("tensor-parallel-size"), default=1)
    dp = _parse_parallel_int(flags.get("data-parallel-size"), default=1)
    return tp * dp


def _is_candidate_family_gap(
    *,
    attribution: dict[str, Any],
    incumbent_rate_per_gpu: float | None,
    result_rate_per_gpu: float | None,
) -> bool:
    if attribution.get("proposal_origin") != "llm_out_of_set":
        return False
    if not isinstance(incumbent_rate_per_gpu, (int, float)):
        return False
    if not isinstance(result_rate_per_gpu, (int, float)):
        return False
    min_gain = max(1e-6, float(incumbent_rate_per_gpu) * 0.01)
    return float(result_rate_per_gpu) > float(incumbent_rate_per_gpu) + min_gain


def _harness_snapshot_payload(
    *,
    study: StudySpec,
    state: StudyState,
    harness_context: dict[str, object],
) -> dict[str, object]:
    experiment_plan = harness_context.get("experiment_plan")
    if not isinstance(experiment_plan, dict):
        experiment_plan = {}
    candidate_set = experiment_plan.get("candidate_set")
    if not isinstance(candidate_set, dict):
        candidate_set = {}
    return {
        "schema_version": 1,
        "study_id": study.study_id,
        "iteration": state.next_trial_index,
        "planner_version": experiment_plan.get("planner_version"),
        "candidate_set_hash": candidate_set.get("candidate_set_hash"),
        "state_ref": {
            "best_trial_id": state.best_trial_id,
            "best_parallel_size": state.best_parallel_size,
            "best_request_rate": state.best_request_rate,
            "best_request_rate_per_gpu": state.best_request_rate_per_gpu,
            "next_trial_index": state.next_trial_index,
            "trial_count": len(state.trials),
        },
        "candidate_set": candidate_set,
        "decisions": {
            "next_action": experiment_plan.get("next_action"),
            "harness_proposal": harness_context.get("harness_proposal"),
            "harness_stop": harness_context.get("harness_stop"),
            "stop_authority": harness_context.get("stop_authority"),
        },
    }


def _maybe_run_initial_config_review(
    *,
    study: StudySpec,
    spec_path: Path,
    store: StudyStore,
    capability_profile: dict[str, Any] | None,
) -> dict[str, Any] | None:
    mode = study.llm.initial_config_review.mode
    if mode == "off":
        return None
    state = store.load_state(study.study_id)
    if state.trials or state.next_trial_index != 1:
        return None

    audit_name = "initial-config-0001"
    audit_root = store.study_root(study.study_id) / "preflight_audits"
    audit_path = audit_root / f"{audit_name}.json"
    if audit_path.exists():
        return json.loads(audit_path.read_text(encoding="utf-8"))

    base_payload: dict[str, Any] = {
        "schema_version": 1,
        "audit_type": "initial_config_review",
        "study_id": study.study_id,
        "mode": mode,
        "repair_applied": False,
    }
    if study.llm.endpoint is None:
        payload = {
            **base_payload,
            "status": "skipped",
            "reason": "llm.endpoint_not_configured",
        }
        store.write_preflight_audit(study.study_id, audit_name, payload)
        return payload

    window, requests = load_trace_requests(study, study_spec_path=spec_path)
    window_summary = summarize_window(requests, window)
    workload_profile = build_study_workload_profile(study, requests, window)
    prompt = build_initial_config_review_prompt(
        study=study,
        window_summary=window_summary,
        capability_profile=capability_profile,
        workload_profile=workload_profile,
    )
    prompt_path = audit_root / f"{audit_name}.prompt.txt"
    raw_path = audit_root / f"{audit_name}.raw.txt"
    prompt_path.write_text(prompt, encoding="utf-8")
    try:
        raw_text = call_llm_for_initial_config_review(policy=study.llm, prompt=prompt)
        raw_path.write_text(raw_text, encoding="utf-8")
        review = parse_initial_config_review_text(raw_text, study)
        payload = {
            **base_payload,
            "status": "completed",
            "prompt_path": str(prompt_path),
            "raw_response_path": str(raw_path),
            "review": review,
        }
    except Exception as exc:  # noqa: BLE001
        payload = {
            **base_payload,
            "status": "failed",
            "prompt_path": str(prompt_path),
            "raw_response_path": str(raw_path) if raw_path.exists() else None,
            "error": str(exc),
        }
    store.write_preflight_audit(study.study_id, audit_name, payload)
    return payload


def _latency_percentiles(summary: object, metric: str) -> dict[str, float]:
    if not isinstance(summary, dict):
        return {}
    payload = summary.get(metric)
    if not isinstance(payload, dict):
        return {}
    selected: dict[str, float] = {}
    for key in ("mean", "p50", "p95", "p99"):
        value = payload.get(key)
        if isinstance(value, (int, float)):
            selected[key] = float(value)
    return selected


def _format_latency_percentiles(metric: str, values: dict[str, float]) -> str:
    if not values:
        return ""
    ordered = ", ".join(
        f"{key}={values[key]:.3f}"
        for key in ("mean", "p50", "p95", "p99")
        if key in values
    )
    return f"{metric}({ordered})"


def _baseline_all_infeasible_stop(result: dict[str, object]) -> tuple[str, dict[str, object]] | None:
    if result.get("status") != "completed":
        return None
    if isinstance(result.get("best_request_rate"), (int, float)):
        return None
    probes = result.get("probes")
    if not isinstance(probes, list) or not probes:
        return None
    if any(isinstance(probe, dict) and probe.get("feasible") for probe in probes):
        return None

    diagnostics = result.get("all_infeasible_diagnostics")
    if not isinstance(diagnostics, dict):
        diagnostics = {}
    lowest_rate = diagnostics.get("request_rate")
    lowest_threshold = diagnostics.get("threshold")
    pass_rate = diagnostics.get("pass_rate")
    early_stop_reason = str(diagnostics.get("early_stop_reason") or "").strip()
    latency_summary = diagnostics.get("latency_summary")
    ttft = _latency_percentiles(latency_summary, "ttft_ms")
    tpot = _latency_percentiles(latency_summary, "tpot_ms")
    details: dict[str, object] = {
        "lowest_sampled_request_rate": lowest_rate,
        "lowest_sampling_u": lowest_threshold,
        "lowest_probe_pass_rate": pass_rate,
        "early_stop_reason": early_stop_reason,
        "lowest_probe_latency_ms": {
            "ttft": ttft,
            "tpot": tpot,
        },
        "lowest_probe_latency_summary": latency_summary if isinstance(latency_summary, dict) else {},
    }
    pieces = [
        "Baseline configuration has no feasible probe under the current SLO.",
        "Stopping tuning because even the lowest sampled request rate did not meet the target pass rate.",
    ]
    if isinstance(lowest_rate, (int, float)):
        pieces.append(f"lowest_sampled_request_rate={float(lowest_rate):.6g}")
    if isinstance(lowest_threshold, (int, float)):
        pieces.append(f"lowest_sampling_u={float(lowest_threshold):.6g}")
    if isinstance(pass_rate, (int, float)):
        pieces.append(f"lowest_probe_pass_rate={float(pass_rate):.6g}")
    if early_stop_reason:
        pieces.append(f"early_stop_reason={early_stop_reason}")
    for item in (
        _format_latency_percentiles("lowest_probe_ttft_ms", ttft),
        _format_latency_percentiles("lowest_probe_tpot_ms", tpot),
    ):
        if item:
            pieces.append(item)
    return " ".join(pieces), details


def _study_source_path(study_root: Path) -> Path:
    return Path((study_root / "study_spec.source").read_text(encoding="utf-8").strip())


def cmd_study_init(args: argparse.Namespace) -> int:
    spec_path = Path(args.spec).resolve()
    study = load_study_spec(spec_path)
    store = StudyStore(Path(args.store_root) if args.store_root else None)
    root = store.init_study(spec_path=spec_path, study=study)
    print(root)
    return 0


def cmd_study_prompt(args: argparse.Namespace) -> int:
    store = StudyStore(Path(args.store_root) if args.store_root else None)
    study_root = Path(args.study_root).resolve()
    source_path = _study_source_path(study_root)
    study = load_study_spec(source_path)
    state = store.load_state(study.study_id)
    capability_profile = load_capability_profile(study, study_spec_path=source_path)
    window, requests = load_trace_requests(study, study_spec_path=source_path)
    prompt = build_prompt(
        study=study,
        window_summary=summarize_window(requests, window),
        state=state,
        capability_profile=capability_profile,
        workload_profile=build_study_workload_profile(study, requests, window),
    )
    prompt_name = args.prompt_name or f"prompt-{state.next_trial_index:04d}"
    path = store.write_prompt(study.study_id, prompt_name, prompt)
    print(path)
    return 0


def cmd_study_llm_propose(args: argparse.Namespace) -> int:
    store = StudyStore(Path(args.store_root) if args.store_root else None)
    study_root = Path(args.study_root).resolve()
    source_path = _study_source_path(study_root)
    study = load_study_spec(source_path)
    state = store.load_state(study.study_id)
    capability_profile = load_capability_profile(study, study_spec_path=source_path)
    window, requests = load_trace_requests(study, study_spec_path=source_path)
    prompt = build_prompt(
        study=study,
        window_summary=summarize_window(requests, window),
        state=state,
        capability_profile=capability_profile,
        workload_profile=build_study_workload_profile(study, requests, window),
    )
    proposal_text = call_llm_for_proposal(
        policy=study.llm,
        prompt=prompt,
        use_harness=study.llm.use_harness,
    )
    proposal = parse_proposal_text(proposal_text, study)
    name = args.proposal_name or f"proposal-{state.next_trial_index:04d}"
    path = store.write_proposal(study.study_id, name, proposal)
    print(path)
    return 0


def cmd_study_register_proposal(args: argparse.Namespace) -> int:
    store = StudyStore(Path(args.store_root) if args.store_root else None)
    study_root = Path(args.study_root).resolve()
    source_path = _study_source_path(study_root)
    study = load_study_spec(source_path)
    proposal = parse_proposal_text(Path(args.proposal_file).read_text(encoding="utf-8"), study)
    name = args.proposal_name or Path(args.proposal_file).stem
    path = store.write_proposal(study.study_id, name, proposal)
    print(path)
    return 0


def cmd_study_emit_job(args: argparse.Namespace) -> int:
    store = StudyStore(Path(args.store_root) if args.store_root else None)
    study_root = Path(args.study_root).resolve()
    source_path = _study_source_path(study_root)
    study = load_study_spec(source_path)
    state = store.load_state(study.study_id)
    proposal_text = Path(args.proposal_file).read_text(encoding="utf-8")
    proposal = parse_proposal_text(proposal_text, study)
    trial, _ = store.materialize_trial(study=study, state=state, proposal=proposal)
    repo_root = Path(__file__).resolve().parents[2]
    job = build_trial_job(study=study, trial=trial, repo_root=repo_root)
    append_job(Path(args.jobs_file).resolve(), job)
    print(trial.trial_id)
    return 0


def cmd_study_ingest(args: argparse.Namespace) -> int:
    store = StudyStore(Path(args.store_root) if args.store_root else None)
    study_root = Path(args.study_root).resolve()
    study_id = study_root.name
    state = store.ingest_trial_results(study_id)
    print(json.dumps({"best_trial_id": state.best_trial_id, "best_request_rate": state.best_request_rate}))
    return 0


def cmd_study_tune(args: argparse.Namespace) -> int:
    spec_path = Path(args.spec).resolve()
    study = load_study_spec(spec_path)
    store = StudyStore(Path(args.store_root) if args.store_root else None)
    study_root = store.init_study(spec_path=spec_path, study=study)
    capability_profile = load_capability_profile(study, study_spec_path=spec_path)
    preflight_audit = _maybe_run_initial_config_review(
        study=study,
        spec_path=spec_path,
        store=store,
        capability_profile=capability_profile,
    )
    proposal_files = [Path(item).resolve() for item in (args.proposal_file or [])]
    max_trials = args.max_trials or (len(proposal_files) if proposal_files else 2)
    proposal_policy = args.proposal_policy
    if max_trials <= 0:
        raise SpecError("max_trials must be positive")
    if proposal_files and max_trials > len(proposal_files):
        max_trials = len(proposal_files)
    executed: list[dict[str, object]] = []
    stop_vetoes = 0
    max_llm_stop_vetoes = 1
    for idx in range(max_trials):
        state = store.load_state(study.study_id)
        if state.tuning_stop_reason:
            executed.append(
                {
                    "trial_id": None,
                    "stopped": True,
                    "reason": state.tuning_stop_reason,
                    "diagnosis": state.tuning_stop_diagnosis,
                    "details": state.tuning_stop_details,
                    "state_best_trial_id": state.best_trial_id,
                    "state_best_request_rate": state.best_request_rate,
                }
            )
            break
        if state.next_trial_index > max_trials:
            break
        window, requests = load_trace_requests(study, study_spec_path=spec_path)
        window_summary = summarize_window(requests, window)
        workload_profile = build_study_workload_profile(study, requests, window)
        harness_context = (
            build_harness_context(
                study=study,
                window_summary=window_summary,
                state=state,
                workload_profile=workload_profile,
            )
            if study.llm.use_harness
            else None
        )
        if harness_context is not None:
            store.write_harness_snapshot(
                study.study_id,
                f"candidate-set-{state.next_trial_index:04d}",
                _harness_snapshot_payload(
                    study=study,
                    state=state,
                    harness_context=harness_context,
                ),
            )
        prompt = build_prompt(
            study=study,
            window_summary=window_summary,
            state=state,
            capability_profile=capability_profile,
            workload_profile=workload_profile,
        )
        prompt_name = f"prompt-{state.next_trial_index:04d}"
        store.write_prompt(study.study_id, prompt_name, prompt)

        if (
            not proposal_files
            and not args.skip_baseline
            and state.next_trial_index == 1
            and not state.trials
        ):
            proposal_source = None
            proposal_name = "baseline-0001"
            proposal_text = json.dumps(
                {
                    "observation": "Evaluate the study's initial engine configuration before LLM-guided edits.",
                    "diagnosis": "Baseline trial aligned with the AITuner evaluate-then-search loop.",
                    "config_patch": {"env_patch": {}, "flag_patch": {}},
                    "expected_effects": [
                        "establish incumbent performance",
                        "provide bottleneck evidence for harness-guided proposals",
                    ],
                    "why_not_previous_failures": "No config changes are applied.",
                    "should_stop": False,
                },
                ensure_ascii=False,
            )
        elif proposal_files:
            proposal_index = state.next_trial_index - 1
            if proposal_index >= len(proposal_files):
                break
            proposal_source = proposal_files[proposal_index]
            proposal_text = proposal_source.read_text(encoding="utf-8")
            proposal_name = proposal_source.stem
        else:
            proposal_source = None
            stop_proposal = (
                build_harness_stop_proposal(harness_context)
                if harness_context is not None
                else None
            )
            if stop_proposal is not None:
                proposal_text = json.dumps(to_jsonable(stop_proposal), ensure_ascii=False)
                proposal_name = f"harness-stop-{state.next_trial_index:04d}"
            else:
                guided_proposal = (
                    build_harness_guided_proposal(harness_context)
                    if harness_context is not None and proposal_policy == "harness-first"
                    else None
                )
                if guided_proposal is not None:
                    proposal_text = json.dumps(
                        to_jsonable(guided_proposal),
                        ensure_ascii=False,
                    )
                    proposal_name = f"harness-proposal-{state.next_trial_index:04d}"
                else:
                    if study.llm.endpoint is None:
                        raise SpecError(
                            "No proposal files provided, study.llm.endpoint is not configured, "
                            "and the harness stop guard did not fire."
                        )
                    proposal_text = call_llm_for_proposal(
                        policy=study.llm,
                        prompt=prompt,
                        use_harness=study.llm.use_harness,
                    )
                    proposal_name = f"proposal-{state.next_trial_index:04d}"
        raw_proposal_path = store.study_root(study.study_id) / "proposals" / f"{proposal_name}.raw.txt"
        raw_proposal_path.write_text(proposal_text, encoding="utf-8")
        proposal = parse_proposal_text(proposal_text, study)
        _reject_repeated_effective_config(
            study=study,
            state=state,
            proposal=proposal,
            proposal_name=proposal_name,
        )
        proposal_attribution = _classify_proposal_attribution(
            study=study,
            state=state,
            proposal=proposal,
            proposal_name=proposal_name,
            proposal_source=proposal_source,
            harness_context=harness_context,
        )
        _validate_harness_candidate_policy(proposal_attribution)
        store.write_proposal(study.study_id, proposal_name, proposal)
        if proposal.should_stop:
            proposal_attribution["stopped"] = True
            store.write_proposal_attribution(
                study.study_id,
                proposal_name,
                proposal_attribution,
            )
            is_harness_stop = proposal_name.startswith("harness-stop-")
            is_llm_stop = not is_harness_stop and proposal_source is None
            stop_authority = (
                harness_context.get("stop_authority")
                if isinstance(harness_context, dict)
                else None
            )
            authorized = stop_authority is None or bool(stop_authority.get("authorized"))
            # Stop-B authority: the deterministic validator overrides an
            # LLM-originated stop. Veto an unauthorized stop (bounded) so the
            # loop does not converge prematurely on the agent's say-so alone.
            if is_llm_stop and not authorized and stop_vetoes < max_llm_stop_vetoes:
                stop_vetoes += 1
                executed.append(
                    {
                        "trial_id": None,
                        "proposal_name": proposal_name,
                        "proposal_source": "llm",
                        "stop_vetoed": True,
                        "reason": "validator_did_not_authorize_stop",
                        "validator_reason": (
                            stop_authority.get("reason") if stop_authority else None
                        ),
                        "diagnosis": proposal.diagnosis,
                    }
                )
                continue
            if is_harness_stop:
                proposal_source_label = "harness"
            else:
                proposal_source_label = str(proposal_source) if proposal_source else "llm"
            stop_authorized_by = (
                "validator"
                if (is_harness_stop or authorized)
                else "file_proposal"
                if proposal_source is not None
                else "llm_after_veto_budget"
            )
            stop_reason = (
                "harness_stop"
                if is_harness_stop
                else "proposal_file_stop"
                if proposal_source is not None
                else "llm_stop"
            )
            stop_details = {
                "proposal_name": proposal_name,
                "proposal_source": proposal_source_label,
                "stop_authorized_by": stop_authorized_by,
            }
            if stop_authority:
                stop_details["validator_reason"] = stop_authority.get("reason")
            state.tuning_stop_reason = stop_reason
            state.tuning_stop_diagnosis = proposal.diagnosis
            state.tuning_stop_details = stop_details
            store.save_state(state)
            executed.append(
                {
                    "trial_id": None,
                    "proposal_name": proposal_name,
                    "proposal_source": proposal_source_label,
                    "stopped": True,
                    "reason": state.tuning_stop_reason,
                    "stop_authorized_by": stop_authorized_by,
                    "diagnosis": proposal.diagnosis,
                    "details": stop_details,
                    "state_best_trial_id": state.best_trial_id,
                    "state_best_request_rate": state.best_request_rate,
                }
            )
            break
        is_auto_baseline = (
            not proposal_files
            and not args.skip_baseline
            and state.next_trial_index == 1
            and not state.trials
            and _is_empty_config_patch(proposal)
        )
        pre_trial_best_rate_per_gpu = state.best_request_rate_per_gpu
        trial, _ = store.materialize_trial(study=study, state=state, proposal=proposal)
        proposal_attribution["trial_id"] = trial.trial_id
        store.write_proposal_attribution(
            study.study_id,
            proposal_name,
            proposal_attribution,
        )
        trial_spec_path = Path(trial.artifact_dir) / "trial_spec.json"
        result = run_trial(trial_spec_path)
        state = store.ingest_trial_results(study.study_id)
        trial_parallel_size = _parallel_size_for_config_patch(study, trial.config_patch)
        result_rate_per_gpu = _result_request_rate_per_gpu(result, trial_parallel_size)
        gap_path: Path | None = None
        if _is_candidate_family_gap(
            attribution=proposal_attribution,
            incumbent_rate_per_gpu=pre_trial_best_rate_per_gpu,
            result_rate_per_gpu=result_rate_per_gpu,
        ):
            gap_payload = _candidate_family_gap_payload(
                study=study,
                trial_id=trial.trial_id,
                proposal_name=proposal_name,
                proposal=proposal,
                attribution=proposal_attribution,
                harness_context=harness_context,
                incumbent_rate_per_gpu=float(pre_trial_best_rate_per_gpu),
                result_rate_per_gpu=float(result_rate_per_gpu),
            )
            gap_path = store.write_candidate_family_gap(
                study.study_id,
                trial.trial_id,
                gap_payload,
            )
        executed.append(
            {
                "trial_id": trial.trial_id,
                "proposal_name": proposal_name,
                "proposal_source": (
                    "harness"
                    if proposal_name.startswith("harness-proposal-")
                    else str(proposal_source) if proposal_source else "llm"
                ),
                "proposal_origin": proposal_attribution.get("proposal_origin"),
                "matched_candidate_id": proposal_attribution.get("matched_candidate_id"),
                "candidate_family_gap_path": str(gap_path) if gap_path else None,
                "best_sampling_u": result.get("best_sampling_u"),
                "best_request_rate": result.get("best_request_rate"),
                "best_request_rate_per_gpu": result_rate_per_gpu,
                "best_pass_rate": result.get("best_pass_rate"),
                "state_best_trial_id": state.best_trial_id,
                "state_best_request_rate": state.best_request_rate,
            }
        )
        if is_auto_baseline:
            stop = _baseline_all_infeasible_stop(result)
            if stop is not None:
                diagnosis, details = stop
                state.tuning_stop_reason = "baseline_all_infeasible"
                state.tuning_stop_diagnosis = diagnosis
                state.tuning_stop_details = details
                store.save_state(state)
                executed.append(
                    {
                        "trial_id": None,
                        "stopped": True,
                        "reason": state.tuning_stop_reason,
                        "diagnosis": diagnosis,
                        "details": details,
                        "state_best_trial_id": state.best_trial_id,
                        "state_best_request_rate": state.best_request_rate,
                    }
                )
                break

    final_state = store.load_state(study.study_id)
    print(
        json.dumps(
            {
                "study_root": str(study_root),
                "preflight_audit": preflight_audit,
                "executed_trials": executed,
                "best_trial_id": final_state.best_trial_id,
                "best_request_rate": final_state.best_request_rate,
                "tuning_stop_reason": final_state.tuning_stop_reason,
                "tuning_stop_diagnosis": final_state.tuning_stop_diagnosis,
                "tuning_stop_details": final_state.tuning_stop_details,
            },
            ensure_ascii=False,
        )
    )
    return 0


def cmd_worker_run_trial(args: argparse.Namespace) -> int:
    result = run_trial(Path(args.trial_spec).resolve())
    print(json.dumps(result))
    return 0


def cmd_compare_run(args: argparse.Namespace) -> int:
    summary = run_compare(
        Path(args.spec).resolve(),
        output_root=Path(args.output_root).resolve() if args.output_root else None,
    )
    print(
        json.dumps(
            {
                "compare_id": summary["compare_id"],
                "compare_root": summary["compare_root"],
                "window_count": summary["aggregate"]["window_count"],
                "wins": summary["aggregate"]["wins"],
            },
            ensure_ascii=False,
        )
    )
    return 0


def _resolve_profile_gpu_count(args: argparse.Namespace, study: StudySpec) -> int:
    gpu_count = args.gpu_count
    if gpu_count is None:
        gpu_count = study.hardware.gpu_count
    if gpu_count <= 0:
        raise SpecError("--gpu-count must be > 0.")
    return int(gpu_count)


def _load_profile_study_spec(spec_path: Path) -> StudySpec:
    payload = dict(load_structured_file(spec_path))
    llm_payload = dict(payload.get("llm") or {})
    llm_payload.pop("endpoint", None)
    payload["llm"] = llm_payload
    return StudySpec.from_dict(payload)


def _profile_current_study_window(args: argparse.Namespace) -> dict[str, object]:
    spec_path = Path(args.spec).resolve()
    study = _load_profile_study_spec(spec_path)
    mode = resolve_length_mode(
        request_mode=study.trace.request_mode,
        length_mode=args.length_mode,
    )
    window, requests = load_trace_requests(study, study_spec_path=spec_path)
    profile = build_workload_profile(
        requests,
        window,
        gpu_count=_resolve_profile_gpu_count(args, study),
        length_mode=mode,
    )
    return {
        "profile": profile.to_dict(),
        "source": {
            "study_spec_path": str(spec_path),
            "window_id": study.trace.window_id,
        },
    }


def _resolve_windows_path_for_profile(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_profile_windows(
    study: StudySpec,
    *,
    study_spec_path: Path,
) -> list[dict[str, object]]:
    windows_path = _resolve_windows_path_for_profile(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, dict) 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, dict)
    ]


def _selected_profile_windows(
    args: argparse.Namespace,
    study: StudySpec,
    *,
    study_spec_path: Path,
) -> list[dict[str, object]]:
    windows = _load_profile_windows(study, study_spec_path=study_spec_path)
    window_ids = set(args.window_id or [])
    selected: list[dict[str, object]] = []
    for item in windows:
        window_id = str(item.get("window_id") or "").strip()
        if not window_id:
            continue
        if window_ids and window_id not in window_ids:
            continue
        if not window_ids and not args.all:
            if window_id != study.trace.window_id:
                continue
        trace_type = str(item.get("trace_type") or "").strip()
        if args.trace_type and trace_type != args.trace_type:
            continue
        date_value = str(item.get("date") or "").strip()
        if args.date_from and date_value and date_value < args.date_from:
            continue
        if args.date_to and date_value and date_value > args.date_to:
            continue
        if args.slot_token and str(item.get("slot_token") or "").strip() != args.slot_token:
            continue
        selected.append(item)
    selected.sort(
        key=lambda item: (
            str(item.get("date") or ""),
            str(item.get("slot_token") or ""),
            str(item.get("window_id") or ""),
        )
    )
    if args.limit is not None:
        selected = selected[: args.limit]
    if not selected:
        raise SpecError("No trace windows selected for profile similarity.")
    return selected


def cmd_profile_window(args: argparse.Namespace) -> int:
    print(json.dumps(_profile_current_study_window(args), ensure_ascii=False, indent=2))
    return 0


def cmd_profile_similarity(args: argparse.Namespace) -> int:
    spec_path = Path(args.spec).resolve()
    study = _load_profile_study_spec(spec_path)
    mode = resolve_length_mode(
        request_mode=study.trace.request_mode,
        length_mode=args.length_mode,
    )
    gpu_count = _resolve_profile_gpu_count(args, study)
    profiles = []
    selected = _selected_profile_windows(args, study, study_spec_path=spec_path)
    for item in selected:
        window_id = str(item["window_id"])
        window_study = replace(study, trace=replace(study.trace, window_id=window_id))
        window, requests = load_trace_requests(window_study, study_spec_path=spec_path)
        profiles.append(
            build_workload_profile(
                requests,
                window,
                gpu_count=gpu_count,
                length_mode=mode,
            )
        )
    print(
        json.dumps(
            {
                "source": {
                    "study_spec_path": str(spec_path),
                    "selected_window_count": len(profiles),
                    "length_mode": mode,
                    "gpu_count": gpu_count,
                },
                "profiles": [profile.to_dict() for profile in profiles],
                "similarity": similarity_report(profiles),
            },
            ensure_ascii=False,
            indent=2,
        )
    )
    return 0


def cmd_profile_interaction_matrix(args: argparse.Namespace) -> int:
    spec_path = Path(args.spec).resolve()
    study = _load_profile_study_spec(spec_path)
    window, requests = load_trace_requests(study, study_spec_path=spec_path)
    matrix = build_interaction_screening_matrix(
        study=study,
        window_summary=summarize_window(requests, window),
    )
    print(json.dumps(matrix, ensure_ascii=False, indent=2))
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="AITuner CLI")
    subparsers = parser.add_subparsers(dest="command", required=True)

    study = subparsers.add_parser("study")
    study_sub = study.add_subparsers(dest="study_command", required=True)

    init = study_sub.add_parser("init")
    init.add_argument("--spec", required=True)
    init.add_argument("--store-root")
    init.set_defaults(func=cmd_study_init)

    prompt = study_sub.add_parser("prompt")
    prompt.add_argument("--study-root", required=True)
    prompt.add_argument("--store-root")
    prompt.add_argument("--prompt-name")
    prompt.set_defaults(func=cmd_study_prompt)

    llm_propose = study_sub.add_parser("llm-propose")
    llm_propose.add_argument("--study-root", required=True)
    llm_propose.add_argument("--store-root")
    llm_propose.add_argument("--proposal-name")
    llm_propose.set_defaults(func=cmd_study_llm_propose)

    register = study_sub.add_parser("register-proposal")
    register.add_argument("--study-root", required=True)
    register.add_argument("--store-root")
    register.add_argument("--proposal-file", required=True)
    register.add_argument("--proposal-name")
    register.set_defaults(func=cmd_study_register_proposal)

    emit = study_sub.add_parser("emit-job")
    emit.add_argument("--study-root", required=True)
    emit.add_argument("--store-root")
    emit.add_argument("--proposal-file", required=True)
    emit.add_argument("--jobs-file", required=True)
    emit.set_defaults(func=cmd_study_emit_job)

    ingest = study_sub.add_parser("ingest")
    ingest.add_argument("--study-root", required=True)
    ingest.add_argument("--store-root")
    ingest.set_defaults(func=cmd_study_ingest)

    tune = study_sub.add_parser("tune")
    tune.add_argument("--spec", required=True)
    tune.add_argument("--store-root")
    tune.add_argument("--proposal-file", action="append")
    tune.add_argument("--max-trials", type=int)
    tune.add_argument(
        "--proposal-policy",
        choices=("harness-first", "llm-first"),
        default="harness-first",
        help=(
            "Choose whether deterministic harness proposals are tried before the LLM "
            "or whether the LLM proposes directly from the harness prompt/context."
        ),
    )
    tune.add_argument(
        "--skip-baseline",
        action="store_true",
        help="Do not automatically evaluate the initial config before LLM proposals.",
    )
    tune.set_defaults(func=cmd_study_tune)

    worker = subparsers.add_parser("worker")
    worker_sub = worker.add_subparsers(dest="worker_command", required=True)
    run = worker_sub.add_parser("run-trial")
    run.add_argument("--trial-spec", required=True)
    run.set_defaults(func=cmd_worker_run_trial)

    compare = subparsers.add_parser("compare")
    compare_sub = compare.add_subparsers(dest="compare_command", required=True)
    compare_run = compare_sub.add_parser("run")
    compare_run.add_argument("--spec", required=True)
    compare_run.add_argument("--output-root")
    compare_run.set_defaults(func=cmd_compare_run)

    profile = subparsers.add_parser("profile")
    profile_sub = profile.add_subparsers(dest="profile_command", required=True)

    profile_window = profile_sub.add_parser("window")
    profile_window.add_argument("--spec", required=True)
    profile_window.add_argument(
        "--length-mode",
        default="auto",
        choices=["auto", "total", "input", "output"],
        help="Token length basis for the L vector. auto uses output for decode_only and total otherwise.",
    )
    profile_window.add_argument(
        "--gpu-count",
        type=int,
        help="GPU denominator for per-GPU arrival rate. Defaults to hardware.gpu_count.",
    )
    profile_window.set_defaults(func=cmd_profile_window)

    profile_similarity = profile_sub.add_parser("similarity")
    profile_similarity.add_argument("--spec", required=True)
    profile_similarity.add_argument("--window-id", action="append")
    profile_similarity.add_argument("--trace-type")
    profile_similarity.add_argument("--date-from")
    profile_similarity.add_argument("--date-to")
    profile_similarity.add_argument("--slot-token")
    profile_similarity.add_argument("--limit", type=int)
    profile_similarity.add_argument(
        "--all",
        action="store_true",
        help="Profile all windows selected by filters. Without this or --window-id, only the study window is used.",
    )
    profile_similarity.add_argument(
        "--length-mode",
        default="auto",
        choices=["auto", "total", "input", "output"],
        help="Token length basis for the L vector. auto uses output for decode_only and total otherwise.",
    )
    profile_similarity.add_argument(
        "--gpu-count",
        type=int,
        help="GPU denominator for per-GPU arrival rate. Defaults to hardware.gpu_count.",
    )
    profile_similarity.set_defaults(func=cmd_profile_similarity)

    profile_interaction = profile_sub.add_parser("interaction-matrix")
    profile_interaction.add_argument("--spec", required=True)
    profile_interaction.set_defaults(func=cmd_profile_interaction_matrix)

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        return int(args.func(args))
    except SpecError as exc:
        print(str(exc), file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())
