from __future__ import annotations

import json
import time
from pathlib import Path
from typing import TYPE_CHECKING, Any

from .engine_adapters.vllm import default_vllm_descriptors
from .harness import _effective_config_signature, build_harness_context, render_harness_context
from .http_client import chat_completion, stream_text_completion
from .spec import LLMPolicySpec, Proposal, SpecError, StudySpec, StudyState, to_jsonable

if TYPE_CHECKING:
    from .lca import WorkloadProfile


def _parse_bool_like(value: Any, *, context: str) -> bool:
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        normalized = value.strip().lower()
        if normalized in {"1", "true", "yes", "on"}:
            return True
        if normalized in {"0", "false", "no", "off"}:
            return False
    raise SpecError(f"{context} must be a boolean-like value.")


def _parse_int_like(value: Any, *, context: str) -> int:
    if isinstance(value, bool):
        raise SpecError(f"{context} must be an integer.")
    if isinstance(value, int):
        return value
    if isinstance(value, float) and value.is_integer():
        return int(value)
    if isinstance(value, str):
        stripped = value.strip()
        if stripped:
            try:
                return int(stripped)
            except ValueError as exc:
                raise SpecError(f"{context} must be an integer.") from exc
    raise SpecError(f"{context} must be an integer.")


def _effective_engine_envs(study: StudySpec, proposal: Proposal | None = None) -> dict[str, Any]:
    envs = dict(study.engine.base_envs)
    if proposal is not None:
        envs.update(proposal.config_patch.env_patch)
    return envs


def _effective_engine_flags(study: StudySpec, proposal: Proposal | None = None) -> dict[str, Any]:
    flags = dict(study.engine.base_flags)
    if proposal is not None:
        flags.update(proposal.config_patch.flag_patch)
    return flags


def _effective_topology(study: StudySpec, proposal: Proposal | None = None) -> dict[str, Any]:
    envs = _effective_engine_envs(study, proposal)
    flags = _effective_engine_flags(study, proposal)
    tp = _parse_int_like(flags.get("tensor-parallel-size", 1), context="tensor-parallel-size")
    dp = _parse_int_like(flags.get("data-parallel-size", 1), context="data-parallel-size")
    raw_enable_ep = flags.get("enable-expert-parallel", False)
    enable_ep = _parse_bool_like(raw_enable_ep, context="enable-expert-parallel")
    raw_ep = flags.get("expert-parallel-size")
    ep = _parse_int_like(raw_ep, context="expert-parallel-size") if raw_ep is not None else None
    effective_ep = ep if ep is not None else (tp * dp if enable_ep else 1)
    return {
        "tensor_parallel_size": tp,
        "data_parallel_size": dp,
        "tp_dp_product": tp * dp,
        "enable_expert_parallel": enable_ep,
        "expert_parallel_size": ep,
        "effective_expert_parallel_size": effective_ep,
        "vllm_tpep_reduce_scatter": _parse_bool_like(
            envs.get("VLLM_TPEP_REDUCE_SCATTER", "0"),
            context="VLLM_TPEP_REDUCE_SCATTER",
        ),
    }


def _launch_failure_history(state: StudyState) -> list[dict[str, Any]]:
    failures: list[dict[str, Any]] = []
    for trial in state.trials:
        if trial.failure_stage != "engine_launch" and "engine_process_exited_before_ready" not in (
            trial.failure_reason or ""
        ):
            continue
        config_patch = trial.config_patch or {}
        env_patch = config_patch.get("env_patch", {}) if isinstance(config_patch, dict) else {}
        flag_patch = config_patch.get("flag_patch", {}) if isinstance(config_patch, dict) else {}
        failures.append(
            {
                "trial_id": trial.trial_id,
                "failure_reason": trial.failure_reason,
                "failure_stage": trial.failure_stage or "engine_launch",
                "implicated_env_keys": sorted(env_patch),
                "implicated_flag_keys": sorted(flag_patch),
                "config_patch": config_patch,
            }
        )
    return failures


def _parallel_space_tunable(study: StudySpec) -> bool:
    tunable_flags = set(study.engine.tunable_flags)
    return {
        "tensor-parallel-size",
        "data-parallel-size",
        "expert-parallel-size",
    }.issubset(tunable_flags)


def _enumerate_parallel_candidates(study: StudySpec) -> list[dict[str, int | bool]]:
    constraints = study.engine.topology_constraints
    if constraints is None or not _parallel_space_tunable(study):
        return []
    base_flags = _effective_engine_flags(study)
    base_enable_ep = _parse_bool_like(
        base_flags.get("enable-expert-parallel", False),
        context="enable-expert-parallel",
    )
    current = _effective_topology(study)
    tp_values = constraints.allowed_tensor_parallel_sizes or [current["tensor_parallel_size"]]
    dp_values = constraints.allowed_data_parallel_sizes or [current["data_parallel_size"]]
    ep_values = constraints.allowed_expert_parallel_sizes or [current["effective_expert_parallel_size"]]
    candidates: list[dict[str, int | bool]] = []
    for tp in sorted(set(tp_values)):
        for dp in sorted(set(dp_values)):
            tp_dp_product = tp * dp
            if (
                constraints.allowed_tp_dp_products
                and tp_dp_product not in constraints.allowed_tp_dp_products
            ):
                continue
            if (
                not constraints.allowed_tp_dp_products
                and tp_dp_product > study.hardware.gpu_count
            ):
                continue
            if (
                constraints.require_tp_dp_product_equals_gpu_count
                and tp_dp_product != study.hardware.gpu_count
            ):
                continue
            for ep in sorted(set(ep_values)):
                if constraints.require_ep_size_leq_tp_dp_product and ep > tp_dp_product:
                    continue
                if (
                    constraints.require_ep_size_divides_tp_dp_product
                    and tp_dp_product % ep != 0
                ):
                    continue
                enable_ep = base_enable_ep or ep > 1
                if (
                    constraints.require_enable_expert_parallel_when_ep_gt_one
                    and ep > 1
                    and not enable_ep
                ):
                    continue
                candidate = {
                    "tensor_parallel_size": tp,
                    "data_parallel_size": dp,
                    "expert_parallel_size": ep,
                    "enable_expert_parallel": enable_ep,
                }
                if (
                    candidate["tensor_parallel_size"] == current["tensor_parallel_size"]
                    and candidate["data_parallel_size"] == current["data_parallel_size"]
                    and candidate["expert_parallel_size"] == current["effective_expert_parallel_size"]
                ):
                    continue
                candidates.append(candidate)
    return candidates


def _descriptor_payload(study: StudySpec) -> list[dict[str, Any]]:
    if study.engine.engine_name.lower() != "vllm":
        return []
    return [
        to_jsonable(descriptor)
        for descriptor in default_vllm_descriptors(tunable_flags=study.engine.tunable_flags)
    ]


def build_initial_config_review_prompt(
    *,
    study: StudySpec,
    window_summary: dict[str, Any],
    capability_profile: dict[str, Any] | None,
    workload_profile: "WorkloadProfile | None" = None,
) -> str:
    """Build the static pre-flight audit prompt for the study's base config."""

    payload: dict[str, Any] = {
        "study_id": study.study_id,
        "objective": "static initial-config audit before any measurement or repair",
        "review_mode": study.llm.initial_config_review.mode,
        "output_contract": {
            "verdict": "one of ok, risky, invalid, unknown",
            "issues": [
                {
                    "knob": "flag/env knob name or mechanism name",
                    "mechanism": "affected mechanism if known",
                    "reason": "short reason grounded in the supplied descriptors/context",
                    "severity": "low|medium|high",
                }
            ],
            "minimal_repair_patch": {
                "env_patch": {},
                "flag_patch": {},
            },
            "do_not_change": ["knob names that should remain fixed during repair"],
            "confidence": "number in [0, 1]",
            "requires_harness_validation": True,
        },
        "hardware": {
            "gpu_count": study.hardware.gpu_count,
            "gpu_model": study.hardware.gpu_model,
        },
        "model": {
            "model_id": study.model.model_id,
            "served_model_name": study.model.served_model_name,
        },
        "engine": {
            "engine_name": study.engine.engine_name,
            "engine_version": study.engine.engine_version,
            "base_envs": study.engine.base_envs,
            "base_flags": study.engine.base_flags,
            "allowed_env_keys": study.engine.tunable_envs,
            "allowed_flag_keys": study.engine.tunable_flags,
            "topology_constraints": (
                study.engine.topology_constraints.__dict__
                if study.engine.topology_constraints is not None
                else None
            ),
            "effective_topology": _effective_topology(study),
        },
        "trace": {
            "window_id": study.trace.window_id,
            "request_mode": study.trace.request_mode,
            "completion_tokens_override": study.trace.completion_tokens_override,
            "input_length_filter": (
                {
                    "min_input_tokens": study.trace.input_length_filter.min_input_tokens,
                    "max_input_tokens": study.trace.input_length_filter.max_input_tokens,
                }
                if study.trace.input_length_filter is not None
                else None
            ),
            "window_summary": window_summary,
            "workload_lca_profile": (
                workload_profile.to_dict() if workload_profile is not None else None
            ),
        },
        "slo": {
            "target_pass_rate": study.slo.target_pass_rate,
            "ttft_rule": study.slo.ttft_rule,
            "tpot_rule": study.slo.tpot_rule,
        },
        "capability_profile": capability_profile or {},
        "knob_descriptors": _descriptor_payload(study),
    }
    sections = [
        "You are doing a pre-flight review of an LLM serving initial configuration.",
        "This is not a tuning proposal. Do not claim measured performance.",
        "Identify obviously risky or inconsistent initial settings before the first baseline trial.",
        "Use the supplied knob descriptors and constraints; do not use case-specific memorized values.",
        "Return exactly one JSON object matching output_contract. Do not wrap it in markdown.",
        "If uncertain, use verdict=unknown and an empty minimal_repair_patch.",
        "The current warn mode records your audit only; the patch will not be applied automatically.",
        "",
        "Audit context:",
        json.dumps(payload, default=lambda value: value.__dict__, ensure_ascii=False, indent=2),
    ]
    return "\n".join(sections)


def build_prompt(
    *,
    study: StudySpec,
    window_summary: dict[str, Any],
    state: StudyState,
    capability_profile: dict[str, Any] | None,
    workload_profile: "WorkloadProfile | None" = None,
) -> str:
    objective_notes: list[str] = []
    if study.trace.request_mode == "decode_only":
        objective_notes.append(
            "This study is decode-only. The engine uses a KV decode benchmark connector, so TTFT is informational only unless an explicit TTFT rule is configured."
        )
    if study.trace.completion_tokens_override is not None:
        objective_notes.append(
            "The trace replay overrides completion length for every request. "
            f"min_tokens=max_tokens={study.trace.completion_tokens_override}."
        )
    if study.slo.ttft_rule is None:
        objective_notes.append("There is no TTFT SLO for this study.")
    if study.slo.tpot_rule is None:
        objective_notes.append("There is no TPOT SLO for this study.")
    history = []
    for trial in state.trials[-study.llm.max_history_trials :]:
        history.append(
            {
                "trial_id": trial.trial_id,
                "status": trial.status,
                "parallel_size": trial.parallel_size,
                "best_sampling_u": trial.best_sampling_u,
                "best_request_rate": trial.best_request_rate,
                "best_request_rate_per_gpu": trial.best_request_rate_per_gpu,
                "best_pass_rate": trial.best_pass_rate,
                "diagnosis": trial.diagnosis,
                "config_patch": trial.config_patch,
                "failure_reason": trial.failure_reason,
                "failure_stage": trial.failure_stage,
            }
        )
    launch_failures = _launch_failure_history(state)
    parallel_candidates = _enumerate_parallel_candidates(study)
    common_preamble = [
        "You are tuning an OpenAI-compatible serving engine.",
        "Return exactly one JSON object with keys: observation, diagnosis, config_patch, expected_effects, why_not_previous_failures, should_stop.",
        "config_patch must contain env_patch and flag_patch.",
        "expected_effects must be a JSON array of short strings, not an object.",
        (
            "should_stop must be a boolean. Use true only when the harness convergence guard says another adjacent probe is not justified."
            if study.llm.use_harness
            else "should_stop must be a boolean. Use false unless no valid config can be proposed."
        ),
        "Only use allowed tunable env keys and allowed tunable flag keys.",
        "Do not wrap the JSON in markdown fences or any extra text.",
        "Do not repeat a config that previously failed to launch unless the new patch explicitly removes the failing knob.",
        "Treat previous engine launch failures as hard negative evidence. If you touch TP/DP/EP, keep the proposal inside the topology constraints exactly.",
    ]
    if not study.llm.use_harness:
        sections = [
            *common_preamble,
            "",
            "Study context:",
            json.dumps(
                {
                    "study_id": study.study_id,
                    "objective": "maximize feasible request_rate_per_gpu at the SLO target",
                    "current_best": {
                        "trial_id": state.best_trial_id,
                        "best_parallel_size": state.best_parallel_size,
                        "best_sampling_u": state.best_sampling_u,
                        "best_request_rate": state.best_request_rate,
                        "best_request_rate_per_gpu": state.best_request_rate_per_gpu,
                    },
                    "hardware": {
                        "gpu_count": study.hardware.gpu_count,
                        "gpu_model": study.hardware.gpu_model,
                    },
                    "model": {
                        "model_id": study.model.model_id,
                        "served_model_name": study.model.served_model_name,
                    },
                    "trace": {
                        "window_id": study.trace.window_id,
                        "request_mode": study.trace.request_mode,
                        "completion_tokens_override": study.trace.completion_tokens_override,
                        "input_length_filter": (
                            {
                                "min_input_tokens": study.trace.input_length_filter.min_input_tokens,
                                "max_input_tokens": study.trace.input_length_filter.max_input_tokens,
                            }
                            if study.trace.input_length_filter is not None
                            else None
                        ),
                    },
                    "engine": {
                        "engine_name": study.engine.engine_name,
                        "engine_version": study.engine.engine_version,
                        "base_flags": study.engine.base_flags,
                        "base_envs": study.engine.base_envs,
                        "allowed_flag_keys": study.engine.tunable_flags,
                        "allowed_env_keys": study.engine.tunable_envs,
                        "topology_constraints": (
                            study.engine.topology_constraints.__dict__
                            if study.engine.topology_constraints is not None
                            else None
                        ),
                    },
                },
                ensure_ascii=False,
                indent=2,
            ),
            "",
            "SLO:",
            json.dumps(
                {
                    "target_pass_rate": study.slo.target_pass_rate,
                    "ttft_rule": study.slo.ttft_rule,
                    "tpot_rule": study.slo.tpot_rule,
                    "objective_notes": objective_notes,
                },
                default=lambda value: value.__dict__,
                ensure_ascii=False,
                indent=2,
            ),
            "",
            "Trial history:",
            json.dumps(history, ensure_ascii=False, indent=2),
            "",
            "Known launch failures:",
            json.dumps(launch_failures, ensure_ascii=False, indent=2),
            "",
            "Tested config signatures:",
            json.dumps(_tested_config_signatures(study, state), ensure_ascii=False, indent=2),
        ]
        return "\n".join(sections)

    sections = [
        *common_preamble,
        (
            "TP/DP/EP are part of the tunable space for this study. Prioritize exploring legal topology changes in parallel space before runtime-only knobs unless recent history already proves a topology variant is worse or fails to launch."
            if parallel_candidates
            else "If TP/DP/EP are not tunable, focus on the remaining launch-safe runtime knobs."
        ),
        (
            "Harness candidate policy is advisory: prefer a high-scoring harness candidate when it matches your diagnosis, but you may propose an out-of-set config when the harness candidate family appears to miss the right step; such proposals are audited as candidate-family gaps."
            if study.llm.harness_candidate_policy == "advisory"
            else "Harness candidate policy is strict: your config_patch must match one of the harness eligible candidates after effective-config materialization."
        ),
        "",
        "Study stack:",
        json.dumps(
            {
                "study_id": study.study_id,
                "current_best": {
                    "trial_id": state.best_trial_id,
                    "best_parallel_size": state.best_parallel_size,
                    "best_sampling_u": state.best_sampling_u,
                    "best_request_rate": state.best_request_rate,
                    "best_request_rate_per_gpu": state.best_request_rate_per_gpu,
                },
                "best_by_parallel_size": state.best_by_parallel_size,
                "hardware": {
                    "gpu_count": study.hardware.gpu_count,
                    "gpu_model": study.hardware.gpu_model,
                },
                "model": {
                    "model_id": study.model.model_id,
                    "served_model_name": study.model.served_model_name,
                },
                "trace": {
                    "window_id": study.trace.window_id,
                    "request_mode": study.trace.request_mode,
                    "completion_tokens_override": study.trace.completion_tokens_override,
                    "early_stop_max_lag_s": study.trace.early_stop_max_lag_s,
                    "early_stop_max_elapsed_s": study.trace.early_stop_max_elapsed_s,
                    "restart_engine_after_early_stop": study.trace.restart_engine_after_early_stop,
                    "input_length_filter": (
                        {
                            "min_input_tokens": study.trace.input_length_filter.min_input_tokens,
                            "max_input_tokens": study.trace.input_length_filter.max_input_tokens,
                        }
                        if study.trace.input_length_filter is not None
                        else None
                    ),
                },
                "engine": {
                    "engine_name": study.engine.engine_name,
                    "engine_version": study.engine.engine_version,
                    "base_flags": study.engine.base_flags,
                    "base_envs": study.engine.base_envs,
                    "allowed_flag_keys": study.engine.tunable_flags,
                    "allowed_env_keys": study.engine.tunable_envs,
                    "topology_constraints": (
                        study.engine.topology_constraints.__dict__
                        if study.engine.topology_constraints is not None
                        else None
                    ),
                    "effective_topology": _effective_topology(study),
                },
            },
            ensure_ascii=False,
            indent=2,
        ),
        "",
        "Window summary:",
        json.dumps(window_summary, ensure_ascii=False, indent=2),
        "",
        "SLO:",
        json.dumps(
            {
                "target_pass_rate": study.slo.target_pass_rate,
                "ttft_rule": study.slo.ttft_rule,
                "tpot_rule": study.slo.tpot_rule,
                "objective_notes": objective_notes,
            },
            default=lambda value: value.__dict__,
            ensure_ascii=False,
            indent=2,
        ),
        "",
        "Capability profile:",
        json.dumps(capability_profile or {}, ensure_ascii=False, indent=2),
        "",
        "Trial history:",
        json.dumps(history, ensure_ascii=False, indent=2),
        "",
        "Known launch failures:",
        json.dumps(launch_failures, ensure_ascii=False, indent=2),
        "",
        "Parallel space candidates:",
        json.dumps(parallel_candidates, ensure_ascii=False, indent=2),
        "",
        "Tested config signatures:",
        json.dumps(_tested_config_signatures(study, state), ensure_ascii=False, indent=2),
    ]
    sections.extend(
        [
            "",
            "Harnesses:",
            render_harness_context(
                build_harness_context(
                    study=study,
                    window_summary=window_summary,
                    state=state,
                    workload_profile=workload_profile,
                )
            ),
            "",
        ]
    )
    sections.extend(
        [
        "The primary cross-topology comparison metric is request_rate_per_gpu, not raw request_rate.",
        "The proposal should beat the incumbent on request_rate_per_gpu under the 95%+ SLO target.",
        (
            "The evaluator may use the same tp_dp_product incumbent as the search floor when search.inherit_incumbent_floor=true."
            if study.search.inherit_incumbent_floor
            else "The evaluator runs each proposal over the full configured search range so raw per-iteration performance is measured directly."
        ),
        "Do not assume a configuration with fewer GPUs should inherit the global incumbent sampling_u.",
        "Follow the active harness. Prefer stop over a weak exploratory proposal once a good incumbent has converged.",
        ]
    )
    return "\n".join(sections)


def _tested_config_signatures(study: StudySpec, state: StudyState) -> list[dict[str, Any]]:
    signatures: list[dict[str, Any]] = []
    seen: set[str] = set()
    for trial in state.trials:
        config_patch = trial.config_patch or {}
        signature = _effective_config_signature(study, config_patch)
        if signature in seen:
            continue
        seen.add(signature)
        signatures.append(
            {
                "trial_id": trial.trial_id,
                "status": trial.status,
                "best_request_rate_per_gpu": trial.best_request_rate_per_gpu,
                "effective_config_signature": signature,
                "config_patch": config_patch,
            }
        )
    return signatures


def load_capability_profile(study: StudySpec, *, study_spec_path: Path) -> dict[str, Any] | None:
    if not study.capability_profile_path:
        return None
    path = Path(study.capability_profile_path)
    if not path.is_absolute():
        path = (study_spec_path.parent / path).resolve()
    return json.loads(path.read_text(encoding="utf-8"))


def validate_proposal(proposal: Proposal, study: StudySpec) -> Proposal:
    unknown_envs = sorted(set(proposal.config_patch.env_patch) - set(study.engine.tunable_envs))
    unknown_flags = sorted(
        set(proposal.config_patch.flag_patch) - set(study.engine.tunable_flags)
    )
    if unknown_envs:
        raise SpecError(f"Proposal uses unsupported env keys: {', '.join(unknown_envs)}")
    if unknown_flags:
        raise SpecError(f"Proposal uses unsupported flag keys: {', '.join(unknown_flags)}")
    topology = _effective_topology(study, proposal)
    tp = topology["tensor_parallel_size"]
    dp = topology["data_parallel_size"]
    effective_ep = topology["effective_expert_parallel_size"]
    if tp < 1 or dp < 1:
        raise SpecError("Proposal violates topology constraints: TP/DP must be >= 1.")
    constraints = study.engine.topology_constraints
    if constraints is None:
        return proposal
    if (
        constraints.allowed_tensor_parallel_sizes
        and tp not in constraints.allowed_tensor_parallel_sizes
    ):
        raise SpecError(
            "Proposal violates topology constraints: "
            f"tensor-parallel-size={tp} not in {constraints.allowed_tensor_parallel_sizes}."
        )
    if constraints.allowed_data_parallel_sizes and dp not in constraints.allowed_data_parallel_sizes:
        raise SpecError(
            "Proposal violates topology constraints: "
            f"data-parallel-size={dp} not in {constraints.allowed_data_parallel_sizes}."
        )
    if (
        constraints.allowed_expert_parallel_sizes
        and effective_ep not in constraints.allowed_expert_parallel_sizes
    ):
        raise SpecError(
            "Proposal violates topology constraints: "
            f"expert-parallel-size={effective_ep} not in {constraints.allowed_expert_parallel_sizes}."
        )
    tp_dp_product = topology["tp_dp_product"]
    if (
        constraints.allowed_tp_dp_products
        and tp_dp_product not in constraints.allowed_tp_dp_products
    ):
        raise SpecError(
            "Proposal violates topology constraints: "
            f"tensor-parallel-size * data-parallel-size = {tp_dp_product} not in "
            f"{constraints.allowed_tp_dp_products}."
        )
    if (
        not constraints.allowed_tp_dp_products
        and tp_dp_product > study.hardware.gpu_count
    ):
        raise SpecError(
            "Proposal violates topology constraints: "
            f"tensor-parallel-size * data-parallel-size = {tp_dp_product} exceeds "
            f"hardware.gpu_count={study.hardware.gpu_count}."
        )
    if (
        constraints.require_tp_dp_product_equals_gpu_count
        and tp_dp_product != study.hardware.gpu_count
    ):
        raise SpecError(
            "Proposal violates topology constraints: "
            f"tensor-parallel-size * data-parallel-size must equal hardware.gpu_count="
            f"{study.hardware.gpu_count}, got {tp_dp_product}."
        )
    if (
        constraints.require_enable_expert_parallel_when_ep_gt_one
        and effective_ep > 1
        and not topology["enable_expert_parallel"]
    ):
        raise SpecError(
            "Proposal violates topology constraints: expert-parallel-size > 1 requires "
            "enable-expert-parallel=true."
        )
    if (
        constraints.require_ep_size_leq_tp_dp_product
        and effective_ep > tp_dp_product
    ):
        raise SpecError(
            "Proposal violates topology constraints: "
            f"expert-parallel-size={effective_ep} must be <= tensor-parallel-size * "
            f"data-parallel-size={tp_dp_product}."
        )
    if (
        constraints.require_ep_size_divides_tp_dp_product
        and tp_dp_product % effective_ep != 0
    ):
        raise SpecError(
            "Proposal violates topology constraints: "
            f"expert-parallel-size={effective_ep} must divide tensor-parallel-size * "
            f"data-parallel-size={tp_dp_product}."
        )
    if (
        constraints.validate_cuda_graph_sizes_divisible_by_tp_when_tp_ep_reduce_scatter
        and topology["vllm_tpep_reduce_scatter"]
        and topology["enable_expert_parallel"]
        and effective_ep > 1
        and tp > 1
    ):
        flags = _effective_engine_flags(study, proposal)
        graph_sizes = flags.get("cuda-graph-sizes")
        if isinstance(graph_sizes, list):
            invalid_sizes = [
                _parse_int_like(item, context="cuda-graph-sizes")
                for item in graph_sizes
                if _parse_int_like(item, context="cuda-graph-sizes") % tp != 0
            ]
            if invalid_sizes:
                raise SpecError(
                    "Proposal violates topology constraints: "
                    f"cuda-graph-sizes must be divisible by tensor-parallel-size={tp} "
                    f"when VLLM_TPEP_REDUCE_SCATTER=1 and expert parallel is enabled. "
                    f"First invalid sizes: {invalid_sizes[:8]}"
                )
    return proposal


def _repair_json_text(text: str) -> str:
    candidate = text.strip()
    if not candidate:
        return candidate
    stack: list[str] = []
    in_string = False
    escape = False
    for char in candidate:
        if in_string:
            if escape:
                escape = False
            elif char == "\\":
                escape = True
            elif char == '"':
                in_string = False
            continue
        if char == '"':
            in_string = True
        elif char == "{":
            stack.append("}")
        elif char == "[":
            stack.append("]")
        elif char in {"}", "]"} and stack and stack[-1] == char:
            stack.pop()
    if in_string:
        candidate += '"'
    if stack:
        candidate += "".join(reversed(stack))
    return candidate


def _parse_json_object_text(text: str) -> dict[str, Any]:
    try:
        payload = json.loads(text)
    except json.JSONDecodeError:
        start = text.find("{")
        if start < 0:
            raise
        end = text.rfind("}")
        candidate = text[start:] if end < start else text[start : end + 1]
        try:
            payload = json.loads(candidate)
        except json.JSONDecodeError:
            payload = json.loads(_repair_json_text(candidate))
    if not isinstance(payload, dict):
        raise SpecError("proposal payload must be a JSON object")
    return payload


def parse_proposal_text(text: str, study: StudySpec) -> Proposal:
    payload = _parse_json_object_text(text)
    proposal = Proposal.from_dict(payload)
    return validate_proposal(proposal, study)


def parse_initial_config_review_text(text: str, study: StudySpec) -> dict[str, Any]:
    payload = _parse_json_object_text(text)
    verdict = str(payload.get("verdict") or "unknown").strip().lower()
    if verdict not in {"ok", "risky", "invalid", "unknown"}:
        raise SpecError(
            "initial-config review verdict must be one of: ok, risky, invalid, unknown."
        )

    issues_payload = payload.get("issues", [])
    if not isinstance(issues_payload, list):
        raise SpecError("initial-config review issues must be a list.")
    issues: list[dict[str, str]] = []
    for idx, item in enumerate(issues_payload):
        if isinstance(item, str):
            issues.append(
                {
                    "knob": "",
                    "mechanism": "",
                    "reason": item.strip(),
                    "severity": "medium",
                }
            )
            continue
        if not isinstance(item, dict):
            raise SpecError(f"initial-config review issues[{idx}] must be an object.")
        severity = str(item.get("severity") or "medium").strip().lower()
        if severity not in {"low", "medium", "high"}:
            severity = "medium"
        issues.append(
            {
                "knob": str(item.get("knob") or "").strip(),
                "mechanism": str(item.get("mechanism") or "").strip(),
                "reason": str(item.get("reason") or "").strip(),
                "severity": severity,
            }
        )

    repair_payload = payload.get("minimal_repair_patch") or {}
    if not isinstance(repair_payload, dict):
        raise SpecError("initial-config review minimal_repair_patch must be an object.")
    if "env_patch" not in repair_payload and "flag_patch" not in repair_payload:
        repair_payload = {"env_patch": {}, "flag_patch": repair_payload}
    repair_proposal = Proposal.from_dict(
        {
            "observation": "Initial-config pre-flight repair candidate.",
            "diagnosis": "Validate the LLM audit's minimal repair patch against study constraints.",
            "config_patch": repair_payload,
            "expected_effects": ["pre-flight audit only"],
            "should_stop": False,
        }
    )
    validate_proposal(repair_proposal, study)

    do_not_change_payload = payload.get("do_not_change", [])
    if isinstance(do_not_change_payload, str):
        do_not_change = [do_not_change_payload.strip()] if do_not_change_payload.strip() else []
    elif isinstance(do_not_change_payload, list):
        do_not_change = [str(item).strip() for item in do_not_change_payload if str(item).strip()]
    else:
        raise SpecError("initial-config review do_not_change must be a list.")

    raw_confidence = payload.get("confidence", 0.0)
    confidence = float(raw_confidence) if isinstance(raw_confidence, (int, float)) else 0.0
    confidence = max(0.0, min(1.0, confidence))
    requires_validation = payload.get("requires_harness_validation")
    if requires_validation is None:
        requires_validation = True
    if not isinstance(requires_validation, bool):
        raise SpecError("initial-config review requires_harness_validation must be boolean.")

    return {
        "schema_version": 1,
        "verdict": verdict,
        "issues": issues,
        "minimal_repair_patch": to_jsonable(repair_proposal.config_patch),
        "do_not_change": do_not_change,
        "confidence": confidence,
        "requires_harness_validation": requires_validation,
    }


def _extract_response_text(response: dict[str, Any]) -> str:
    output_text = response.get("output_text")
    if isinstance(output_text, str) and output_text:
        return output_text
    choices = response.get("choices")
    if isinstance(choices, list) and choices:
        message = choices[0].get("message", {})
        if isinstance(message, dict):
            content = message.get("content")
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                text = "".join(
                    item.get("text", "")
                    for item in content
                    if isinstance(item, dict) and isinstance(item.get("text"), str)
                )
                if text:
                    return text
    output = response.get("output")
    if isinstance(output, list):
        parts: list[str] = []
        for item in output:
            if not isinstance(item, dict):
                continue
            content = item.get("content")
            if not isinstance(content, list):
                continue
            for block in content:
                if not isinstance(block, dict):
                    continue
                text = block.get("text")
                if isinstance(text, str) and text:
                    parts.append(text)
        if parts:
            return "".join(parts)
    raise RuntimeError("LLM response content is empty")


def _call_llm_text(
    *,
    policy: LLMPolicySpec,
    prompt: str,
    system_prompt: str = "",
) -> str:
    if policy.endpoint is None:
        raise RuntimeError("study.llm.endpoint is not configured")
    last_error: Exception | None = None
    max_attempts = 4
    for attempt in range(max_attempts):
        try:
            if policy.endpoint.stream:
                text = stream_text_completion(
                    base_url=policy.endpoint.base_url,
                    api_key_env=policy.endpoint.api_key_env,
                    provider=policy.endpoint.provider,
                    wire_api=policy.endpoint.wire_api,
                    model=policy.endpoint.model,
                    messages=[{"role": "user", "content": prompt}],
                    timeout_s=policy.endpoint.timeout_s,
                    system_prompt=system_prompt,
                    reasoning_effort=policy.endpoint.reasoning_effort,
                )
            else:
                response = chat_completion(
                    base_url=policy.endpoint.base_url,
                    api_key_env=policy.endpoint.api_key_env,
                    provider=policy.endpoint.provider,
                    wire_api=policy.endpoint.wire_api,
                    model=policy.endpoint.model,
                    messages=[{"role": "user", "content": prompt}],
                    timeout_s=policy.endpoint.timeout_s,
                    system_prompt=system_prompt,
                    reasoning_effort=policy.endpoint.reasoning_effort,
                )
                text = _extract_response_text(response)
            if text.strip():
                return text
            last_error = RuntimeError("LLM response content is empty")
        except Exception as exc:  # noqa: BLE001
            last_error = exc
        if attempt < max_attempts - 1:
            time.sleep(min(30.0, 2.0 * (2**attempt)))
            continue
    raise RuntimeError(f"LLM proposal failed after retry: {last_error}") from last_error


def call_llm_for_proposal(
    *,
    policy: LLMPolicySpec,
    prompt: str,
    use_harness: bool = True,
) -> str:
    system_prompt = policy.system_prompt if use_harness else ""
    return _call_llm_text(policy=policy, prompt=prompt, system_prompt=system_prompt)


def call_llm_for_initial_config_review(
    *,
    policy: LLMPolicySpec,
    prompt: str,
) -> str:
    review_system = "\n".join(
        item
        for item in (
            policy.system_prompt,
            "You are auditing an initial serving config. Return only the requested JSON audit.",
        )
        if item.strip()
    )
    return _call_llm_text(policy=policy, prompt=prompt, system_prompt=review_system)
