from __future__ import annotations

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

from .config_signature import materialize_proposal_for_execution
from .spec import (
    Proposal,
    SamplingSearchSpec,
    StudySpec,
    StudyState,
    TrialSpec,
    TrialSummary,
    to_jsonable,
)


class StudyStore:
    def __init__(self, root: Path | None = None):
        base = root or Path(".aituner") / "studies"
        self.root = base.resolve()

    def study_root(self, study_id: str) -> Path:
        return self.root / study_id

    def init_study(self, *, spec_path: Path, study: StudySpec) -> Path:
        root = self.study_root(study.study_id)
        for rel in (
            "prompts",
            "proposals",
            "proposal_attributions",
            "trials",
            "results",
            "harness",
            "candidate_family_gaps",
            "preflight_audits",
        ):
            (root / rel).mkdir(parents=True, exist_ok=True)
        (root / "study_spec.source").write_text(str(spec_path.resolve()) + "\n", encoding="utf-8")
        self.write_json(root / "study_spec.snapshot.json", to_jsonable(study))
        if not (root / "state.json").exists():
            self.write_json(root / "state.json", to_jsonable(StudyState(study_id=study.study_id)))
        return root

    def load_state(self, study_id: str) -> StudyState:
        payload = json.loads((self.study_root(study_id) / "state.json").read_text(encoding="utf-8"))
        trials = [TrialSummary(**item) for item in payload.get("trials", [])]
        return StudyState(
            study_id=str(payload["study_id"]),
            best_trial_id=payload.get("best_trial_id"),
            best_parallel_size=payload.get("best_parallel_size"),
            best_sampling_u=payload.get("best_sampling_u"),
            best_request_rate=payload.get("best_request_rate"),
            best_request_rate_per_gpu=payload.get("best_request_rate_per_gpu"),
            next_trial_index=int(payload.get("next_trial_index", 1)),
            tuning_stop_reason=str(payload.get("tuning_stop_reason") or ""),
            tuning_stop_diagnosis=str(payload.get("tuning_stop_diagnosis") or ""),
            tuning_stop_details=dict(payload.get("tuning_stop_details") or {}),
            best_by_parallel_size={
                str(key): value
                for key, value in (payload.get("best_by_parallel_size") or {}).items()
                if isinstance(value, dict)
            },
            trials=trials,
        )

    def save_state(self, state: StudyState) -> None:
        self.write_json(self.study_root(state.study_id) / "state.json", to_jsonable(state))

    def write_prompt(self, study_id: str, prompt_name: str, prompt_text: str) -> Path:
        path = self.study_root(study_id) / "prompts" / f"{prompt_name}.txt"
        path.write_text(prompt_text, encoding="utf-8")
        return path

    def write_proposal(self, study_id: str, proposal_name: str, proposal: Proposal) -> Path:
        path = self.study_root(study_id) / "proposals" / f"{proposal_name}.json"
        self.write_json(path, to_jsonable(proposal))
        return path

    def write_harness_snapshot(
        self,
        study_id: str,
        snapshot_name: str,
        payload: dict[str, Any],
    ) -> Path:
        path = self.study_root(study_id) / "harness" / f"{snapshot_name}.json"
        self.write_json(path, payload)
        return path

    def write_proposal_attribution(
        self,
        study_id: str,
        proposal_name: str,
        payload: dict[str, Any],
    ) -> Path:
        path = self.study_root(study_id) / "proposal_attributions" / f"{proposal_name}.json"
        self.write_json(path, payload)
        return path

    def write_candidate_family_gap(
        self,
        study_id: str,
        trial_id: str,
        payload: dict[str, Any],
    ) -> Path:
        path = self.study_root(study_id) / "candidate_family_gaps" / f"{trial_id}.json"
        self.write_json(path, payload)
        return path

    def write_preflight_audit(
        self,
        study_id: str,
        audit_name: str,
        payload: dict[str, Any],
    ) -> Path:
        path = self.study_root(study_id) / "preflight_audits" / f"{audit_name}.json"
        self.write_json(path, payload)
        return path

    def materialize_trial(
        self,
        *,
        study: StudySpec,
        state: StudyState,
        proposal: Proposal,
    ) -> tuple[TrialSpec, StudyState]:
        proposal = materialize_proposal_for_execution(
            study=study,
            state=state,
            proposal=proposal,
        )
        trial_id = f"trial-{state.next_trial_index:04d}"
        trial_root = self.study_root(study.study_id) / "trials" / trial_id
        trial_root.mkdir(parents=True, exist_ok=True)
        parallel_size = _parallel_size_for_proposal(study=study, proposal=proposal)
        search = study.search
        if study.search.inherit_incumbent_floor:
            search = replace(
                study.search,
                low=_derive_search_floor(
                    study=study,
                    state=state,
                    parallel_size=parallel_size,
                ),
            )
        search, search_evidence = resolve_auto_high_search(
            search=search,
            sampling_us=_sampling_us_for_study_source(
                study=study,
                study_spec_source_path=self.study_root(study.study_id) / "study_spec.source",
            ),
        )
        spec = TrialSpec(
            study_id=study.study_id,
            trial_id=trial_id,
            config_patch=proposal.config_patch,
            search=search,
            study_spec_path=str((self.study_root(study.study_id) / "study_spec.source").resolve()),
            artifact_dir=str(trial_root),
            probe_log_path=str(trial_root / "probe_history.json"),
            engine_log_path=str(trial_root / "engine.log"),
            result_path=str(trial_root / "result.json"),
            search_evidence=search_evidence,
        )
        self.write_json(trial_root / "trial_spec.json", to_jsonable(spec))
        next_trial = (
            TrialSummary(
                trial_id=trial_id,
                status="queued",
                parallel_size=parallel_size,
                diagnosis=proposal.diagnosis,
                config_patch=to_jsonable(proposal.config_patch),
            )
        )
        next_state = replace(
            state,
            next_trial_index=state.next_trial_index + 1,
            trials=[*state.trials, next_trial],
        )
        self.save_state(next_state)
        return spec, next_state

    def ingest_trial_results(self, study_id: str) -> StudyState:
        study = self.load_study_snapshot(study_id)
        state = self.load_state(study_id)
        by_id = {item.trial_id: item for item in state.trials}
        trials_dir = self.study_root(study_id) / "trials"
        best_trial_id = state.best_trial_id
        best_parallel_size = state.best_parallel_size
        best_sampling_u = state.best_sampling_u
        best_rate = state.best_request_rate
        best_rate_per_gpu = state.best_request_rate_per_gpu
        best_by_parallel_size = {
            str(key): dict(value)
            for key, value in (state.best_by_parallel_size or {}).items()
            if isinstance(value, dict)
        }
        for trial_dir in sorted(trials_dir.glob("trial-*")):
            result_path = trial_dir / "result.json"
            if not result_path.exists():
                continue
            payload = json.loads(result_path.read_text(encoding="utf-8"))
            trial_id = str(payload["trial_id"])
            summary = by_id.get(trial_id)
            if summary is None:
                summary = TrialSummary(trial_id=trial_id, status="unknown")
                state.trials.append(summary)
                by_id[trial_id] = summary
            summary.status = str(payload.get("status") or "completed")
            if summary.parallel_size is None:
                summary.parallel_size = _parallel_size_for_trial_id(
                    study=study,
                    study_root=self.study_root(study_id),
                    trial_id=trial_id,
                )
            summary.best_sampling_u = payload.get("best_sampling_u")
            summary.best_request_rate = payload.get("best_request_rate")
            summary.best_request_rate_per_gpu = _request_rate_per_gpu(
                summary.best_request_rate,
                summary.parallel_size,
            )
            summary.best_pass_rate = payload.get("best_pass_rate")
            summary.result_path = str(result_path)
            summary.failure_reason = str(payload.get("failure_reason") or "").strip()
            summary.failure_stage = str(payload.get("failure_stage") or "").strip()
            if (
                isinstance(summary.parallel_size, int)
                and summary.parallel_size > 0
                and isinstance(summary.best_request_rate_per_gpu, (int, float))
            ):
                group_key = str(summary.parallel_size)
                incumbent = best_by_parallel_size.get(group_key)
                if incumbent is None or float(summary.best_request_rate_per_gpu) > float(
                    incumbent.get("best_request_rate_per_gpu") or float("-inf")
                ):
                    best_by_parallel_size[group_key] = {
                        "trial_id": trial_id,
                        "parallel_size": summary.parallel_size,
                        "best_sampling_u": summary.best_sampling_u,
                        "best_request_rate": summary.best_request_rate,
                        "best_request_rate_per_gpu": summary.best_request_rate_per_gpu,
                        "best_pass_rate": summary.best_pass_rate,
                    }
            if (
                isinstance(summary.best_request_rate_per_gpu, (int, float))
                and (
                    best_rate_per_gpu is None
                    or float(summary.best_request_rate_per_gpu) > float(best_rate_per_gpu)
                )
            ):
                best_rate = (
                    float(summary.best_request_rate)
                    if isinstance(summary.best_request_rate, (int, float))
                    else None
                )
                best_rate_per_gpu = float(summary.best_request_rate_per_gpu)
                best_sampling_u = (
                    float(summary.best_sampling_u)
                    if isinstance(summary.best_sampling_u, (int, float))
                    else None
                )
                best_trial_id = trial_id
                best_parallel_size = summary.parallel_size
        state.best_sampling_u = best_sampling_u
        state.best_request_rate = best_rate
        state.best_request_rate_per_gpu = best_rate_per_gpu
        state.best_trial_id = best_trial_id
        state.best_parallel_size = best_parallel_size
        state.best_by_parallel_size = best_by_parallel_size
        self.save_state(state)
        return state

    def load_study_snapshot(self, study_id: str) -> StudySpec:
        path = self.study_root(study_id) / "study_spec.snapshot.json"
        payload = json.loads(path.read_text(encoding="utf-8"))
        return StudySpec.from_dict(payload)

    @staticmethod
    def write_json(path: Path, payload: Any) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def _parse_int_like(value: Any, *, default: int = 1) -> int:
    if value is None:
        return default
    if isinstance(value, bool):
        raise ValueError("Boolean values are not valid integer 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 ValueError(f"Unable to parse integer value from {value!r}.")


def _parallel_size_for_config(*, study: StudySpec, flag_patch: dict[str, Any]) -> int:
    flags = dict(study.engine.base_flags)
    flags.update(flag_patch)
    tp = _parse_int_like(flags.get("tensor-parallel-size"), default=1)
    dp = _parse_int_like(flags.get("data-parallel-size"), default=1)
    return tp * dp


def _parallel_size_for_proposal(*, study: StudySpec, proposal: Proposal) -> int:
    return _parallel_size_for_config(study=study, flag_patch=proposal.config_patch.flag_patch)


def _parallel_size_for_trial_id(*, study: StudySpec, study_root: Path, trial_id: str) -> int | None:
    trial_spec_path = study_root / "trials" / trial_id / "trial_spec.json"
    if not trial_spec_path.exists():
        return None
    payload = json.loads(trial_spec_path.read_text(encoding="utf-8"))
    config_patch = payload.get("config_patch") or {}
    flag_patch = config_patch.get("flag_patch") if isinstance(config_patch, dict) else {}
    if not isinstance(flag_patch, dict):
        return None
    return _parallel_size_for_config(study=study, flag_patch=flag_patch)


def _request_rate_per_gpu(best_request_rate: Any, parallel_size: int | None) -> float | None:
    if not isinstance(best_request_rate, (int, float)) or not isinstance(parallel_size, int):
        return None
    if parallel_size <= 0:
        return None
    return float(best_request_rate) / float(parallel_size)


def _derive_search_floor(*, study: StudySpec, state: StudyState, parallel_size: int) -> float:
    low = study.search.low
    high = study.search.high
    group_incumbent = (state.best_by_parallel_size or {}).get(str(parallel_size))
    if isinstance(group_incumbent, dict) and isinstance(
        group_incumbent.get("best_sampling_u"), (int, float)
    ):
        candidate = float(group_incumbent["best_sampling_u"])
    else:
        candidate = low
    return min(high, max(low, candidate))


def _sampling_us_for_study_source(
    *,
    study: StudySpec,
    study_spec_source_path: Path,
) -> list[float]:
    if not study.search.auto_high.enabled:
        return []
    from .trace import load_trace_requests

    study_spec_path = Path(study_spec_source_path.read_text(encoding="utf-8").strip())
    _, requests = load_trace_requests(study, study_spec_path=study_spec_path)
    return [float(request.sampling_u) for request in requests]


def resolve_auto_high_search(
    *,
    search: SamplingSearchSpec,
    sampling_us: list[float],
) -> tuple[SamplingSearchSpec, dict[str, Any]]:
    policy = search.auto_high
    trace_max_sampling_u = max(sampling_us) if sampling_us else None
    evidence = {
        "enabled": policy.enabled,
        "original_high": search.high,
        "effective_high": search.high,
        "trace_max_sampling_u": trace_max_sampling_u,
        "max_sampling_u": policy.max_sampling_u,
        "require_human_confirmation_beyond_trace": (
            policy.require_human_confirmation_beyond_trace
        ),
        "reason": "auto_high_disabled",
    }
    if not policy.enabled:
        return search, evidence
    if trace_max_sampling_u is None:
        evidence["reason"] = "trace_has_no_sampling_u"
        return search, evidence
    ceiling = min(float(policy.max_sampling_u), 1.0, float(trace_max_sampling_u))
    evidence["effective_ceiling"] = ceiling
    if ceiling < float(search.low):
        evidence["reason"] = "auto_high_ceiling_below_search_low"
        return search, evidence
    if abs(float(search.high) - ceiling) <= 1e-12:
        evidence["reason"] = "search_high_already_at_auto_high_ceiling"
        return search, evidence
    updated = replace(search, high=ceiling)
    evidence["effective_high"] = updated.high
    evidence["reason"] = (
        "search_high_raised_to_trace_ceiling"
        if float(search.high) < ceiling
        else "search_high_lowered_to_trace_ceiling"
    )
    return updated, evidence
