from __future__ import annotations

import json
import random
from dataclasses import asdict, replace
from pathlib import Path
from typing import Any

from .events import sha256_json
from .scenarios import Scenario
from .study import StudyTrial, build_schedule


def build_extension_schedule(
    scenarios: list[Scenario],
    repetitions: int,
    seed: int,
    gate_scenario_ids: list[str],
) -> list[StudyTrial]:
    base = build_schedule(scenarios, repetitions, seed)
    known = {item.scenario_id for item in scenarios}
    if len(gate_scenario_ids) != len(set(gate_scenario_ids)):
        raise ValueError("Gate scenario identifiers must be unique")
    if not set(gate_scenario_ids).issubset(known):
        raise ValueError("Gate references an unknown external scenario")

    gate_trial_ids = [f"R01-{item}" for item in gate_scenario_ids]
    by_trial = {item.trial_id: item for item in base}
    gate = [by_trial[item] for item in gate_trial_ids]
    gate_set = set(gate_trial_ids)
    remaining = [item for item in base if item.trial_id not in gate_set]
    ordered = gate + remaining
    return [
        replace(item, schedule_position=index)
        for index, item in enumerate(ordered, start=1)
    ]


def extension_schedule_document(
    scenarios: list[Scenario],
    repetitions: int,
    seed: int,
    gate_scenario_ids: list[str],
) -> dict[str, Any]:
    trials = build_extension_schedule(
        scenarios, repetitions, seed, gate_scenario_ids
    )
    rows = [asdict(item) for item in trials]
    gate_trial_ids = [f"R01-{item}" for item in gate_scenario_ids]
    return {
        "schema_version": "1.0",
        "schedule_type": "external_replication_with_embedded_gate",
        "seed": seed,
        "repetitions": repetitions,
        "scenario_count": len(scenarios),
        "trial_count": len(trials),
        "gate_trial_ids": gate_trial_ids,
        "gate_trial_count": len(gate_trial_ids),
        "schedule_hash": sha256_json(rows),
        "trials": rows,
    }


def freeze_extension_schedule(
    path: Path,
    scenarios: list[Scenario],
    repetitions: int,
    seed: int,
    gate_scenario_ids: list[str],
) -> dict[str, Any]:
    expected = extension_schedule_document(
        scenarios, repetitions, seed, gate_scenario_ids
    )
    if path.is_file():
        existing = json.loads(path.read_text(encoding="utf-8"))
        if existing.get("schedule_hash") != expected["schedule_hash"]:
            raise ValueError(
                "Existing extension schedule does not match the frozen design"
            )
        return existing
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(expected, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return expected


def read_finished_trials(path: Path) -> dict[str, dict[str, Any]]:
    finished: dict[str, dict[str, Any]] = {}
    if not path.is_file():
        return finished
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        event = json.loads(line)
        if event.get("event_type") != "run_finished":
            continue
        payload = dict(event.get("payload", {}))
        trial_id = str(payload.get("trial_id", ""))
        if trial_id:
            finished[trial_id] = payload
    return finished
