from __future__ import annotations

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

from .events import sha256_json
from .scenarios import Scenario


@dataclass(frozen=True)
class StudyTrial:
    trial_id: str
    repetition: int
    within_repetition_position: int
    schedule_position: int
    scenario_id: str
    family: str
    variant: str
    expected_decision: str

    def as_dict(self) -> dict[str, Any]:
        return asdict(self)


def build_schedule(
    scenarios: list[Scenario], repetitions: int, seed: int
) -> list[StudyTrial]:
    if repetitions <= 0:
        raise ValueError("repetitions must be positive")
    if not scenarios:
        raise ValueError("at least one scenario is required")
    schedule: list[StudyTrial] = []
    global_position = 0
    for repetition in range(1, repetitions + 1):
        ordered = list(scenarios)
        random.Random(seed + repetition * 100_003).shuffle(ordered)
        for within_position, scenario in enumerate(ordered, start=1):
            global_position += 1
            schedule.append(
                StudyTrial(
                    trial_id=f"R{repetition:02d}-{scenario.scenario_id}",
                    repetition=repetition,
                    within_repetition_position=within_position,
                    schedule_position=global_position,
                    scenario_id=scenario.scenario_id,
                    family=scenario.family,
                    variant=scenario.variant,
                    expected_decision=scenario.expected_decision,
                )
            )
    return schedule


def schedule_document(
    scenarios: list[Scenario], repetitions: int, seed: int
) -> dict[str, Any]:
    trials = build_schedule(scenarios, repetitions, seed)
    trial_rows = [trial.as_dict() for trial in trials]
    return {
        "schema_version": "1.0",
        "seed": seed,
        "repetitions": repetitions,
        "scenario_count": len(scenarios),
        "trial_count": len(trials),
        "schedule_hash": sha256_json(trial_rows),
        "trials": trial_rows,
    }


def freeze_schedule(
    path: Path, scenarios: list[Scenario], repetitions: int, seed: int
) -> dict[str, Any]:
    expected = schedule_document(scenarios, repetitions, seed)
    if path.is_file():
        existing = json.loads(path.read_text(encoding="utf-8"))
        if existing.get("schedule_hash") != expected["schedule_hash"]:
            raise ValueError(
                "Existing schedule does not match the requested seed/repetitions"
            )
        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


@dataclass(frozen=True)
class LedgerState:
    completed_trial_ids: frozenset[str]
    terminal_trial_ids: frozenset[str]
    abort_attempts: dict[str, int]
    completed_cost_usd: float
    transient_abort_count: int
    harness_abort_count: int
    model_conformance_failure_count: int


def read_ledger(path: Path) -> LedgerState:
    if not path.is_file():
        return LedgerState(frozenset(), frozenset(), {}, 0.0, 0, 0, 0)
    completed: set[str] = set()
    terminal: set[str] = set()
    abort_attempts: dict[str, int] = {}
    completed_cost = 0.0
    transient = 0
    harness = 0
    conformance = 0
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        event = json.loads(line)
        payload = event.get("payload", {})
        event_type = event.get("event_type")
        if event_type == "run_finished":
            trial_id = str(payload.get("trial_id", ""))
            if trial_id and trial_id not in completed:
                completed.add(trial_id)
                completed_cost += float(payload.get("cost_usd", 0.0))
        elif event_type == "run_aborted":
            trial = payload.get("trial", {})
            trial_id = str(trial.get("trial_id", ""))
            if not trial_id:
                continue
            completed_cost += float(payload.get("cost_usd", 0.0))
            abort_attempts[trial_id] = abort_attempts.get(trial_id, 0) + 1
            abort_class = str(payload.get("abort_class", "harness_error"))
            if abort_class == "provider_transient":
                transient += 1
            elif abort_class == "model_conformance":
                conformance += 1
            else:
                harness += 1
            if bool(payload.get("terminal", False)):
                terminal.add(trial_id)
    return LedgerState(
        completed_trial_ids=frozenset(completed),
        terminal_trial_ids=frozenset(terminal),
        abort_attempts=abort_attempts,
        completed_cost_usd=completed_cost,
        transient_abort_count=transient,
        harness_abort_count=harness,
        model_conformance_failure_count=conformance,
    )


def study_status(
    provider: str,
    model: str,
    schedule: dict[str, Any],
    ledger: LedgerState,
) -> dict[str, Any]:
    total = int(schedule["trial_count"])
    primary_complete = len(
        ledger.completed_trial_ids | ledger.terminal_trial_ids
    )
    return {
        "schema_version": "1.0",
        "provider": provider,
        "model": model,
        "schedule_hash": schedule["schedule_hash"],
        "planned_trials": total,
        "primary_complete": primary_complete,
        "successful_run_records": len(ledger.completed_trial_ids),
        "terminal_conformance_records": len(ledger.terminal_trial_ids),
        "remaining": max(0, total - primary_complete),
        "completion_fraction": primary_complete / total if total else 0.0,
        "completed_cost_usd": round(ledger.completed_cost_usd, 8),
        "transient_abort_count": ledger.transient_abort_count,
        "harness_abort_count": ledger.harness_abort_count,
        "model_conformance_failure_count": (
            ledger.model_conformance_failure_count
        ),
    }
