"""Deterministic paired-trial plan expansion."""

from __future__ import annotations

import itertools
import json
from pathlib import Path
from typing import Any, Mapping

from .hashing import stable_hash, to_jsonable
from .rng import named_seed_record
from .schema import (
    SCHEMA_VERSION,
    STORAGE_FORMATS,
    ContractError,
    require_nonempty_string,
    require_schema,
    require_sequence,
)


def load_document(path: Path) -> dict[str, Any]:
    """Load JSON, with optional YAML support when PyYAML is available."""
    path = Path(path)
    suffix = path.suffix.lower()
    with path.open("r", encoding="utf-8") as stream:
        if suffix == ".json":
            document = json.load(stream)
        elif suffix in {".yaml", ".yml"}:
            try:
                import yaml
            except ImportError as exc:
                raise RuntimeError(
                    "YAML input requires PyYAML; use JSON in minimal environments"
                ) from exc
            document = yaml.safe_load(stream)
        else:
            raise ValueError("experiment documents must use .json, .yaml, or .yml")
    if not isinstance(document, dict):
        raise ContractError(f"{path} must contain a mapping")
    return document


def _normalize_methods(methods: Any) -> list[dict[str, Any]]:
    normalized: list[dict[str, Any]] = []
    for entry in require_sequence(methods, "methods"):
        if isinstance(entry, str):
            method = {"method_id": require_nonempty_string(entry, "method")}
        elif isinstance(entry, Mapping):
            method = dict(to_jsonable(entry))
            source_id = method.pop("id", method.get("method_id"))
            method["method_id"] = require_nonempty_string(
                source_id, "method.method_id"
            )
        else:
            raise ContractError("each method must be a string or mapping")
        normalized.append(method)
    identifiers = [method["method_id"] for method in normalized]
    if len(set(identifiers)) != len(identifiers):
        raise ContractError("method identifiers must be unique")
    return normalized


def _normalize_trajectories(trajectories: Any) -> list[dict[str, Any]]:
    normalized: list[dict[str, Any]] = []
    for entry in require_sequence(trajectories, "trajectories"):
        if isinstance(entry, str):
            trajectory = {
                "trajectory_id": require_nonempty_string(entry, "trajectory")
            }
        elif isinstance(entry, Mapping):
            trajectory = dict(to_jsonable(entry))
            source_id = trajectory.pop("id", trajectory.get("trajectory_id"))
            trajectory["trajectory_id"] = require_nonempty_string(
                source_id, "trajectory.trajectory_id"
            )
        else:
            raise ContractError("each trajectory must be a string or mapping")
        normalized.append(trajectory)
    identifiers = [trajectory["trajectory_id"] for trajectory in normalized]
    if len(set(identifiers)) != len(identifiers):
        raise ContractError("trajectory identifiers must be unique")
    return normalized


def _factor_cells(factors: Any) -> list[dict[str, Any]]:
    if factors is None:
        return [{}]
    if not isinstance(factors, Mapping):
        raise ContractError("factors must be a mapping from name to levels")
    names = sorted(factors)
    levels = [list(require_sequence(factors[name], f"factors.{name}")) for name in names]
    return [
        dict(zip(names, to_jsonable(combination)))
        for combination in itertools.product(*levels)
    ]


def build_trial_plan(specification: Mapping[str, Any]) -> dict[str, Any]:
    """Expand a study specification into immutable paired trial records."""
    if not isinstance(specification, Mapping):
        raise ContractError("study specification must be a mapping")
    spec = dict(to_jsonable(specification))
    study_id = require_nonempty_string(spec.get("study_id"), "study_id")
    split = require_nonempty_string(spec.get("split"), "split")
    if split not in {"calibration", "pilot", "locked"}:
        raise ContractError("split must be calibration, pilot, or locked")
    root_seed = int(spec.get("root_seed", 0))
    if root_seed < 0:
        raise ContractError("root_seed must be non-negative")
    replicates = int(spec.get("replicates", 1))
    if replicates <= 0:
        raise ContractError("replicates must be positive")

    methods = _normalize_methods(spec.get("methods"))
    trajectories = _normalize_trajectories(spec.get("trajectories"))
    factor_cells = _factor_cells(spec.get("factors", {}))

    trials: list[dict[str, Any]] = []
    pair_ids: set[str] = set()
    for trajectory, factors, replicate in itertools.product(
        trajectories, factor_cells, range(replicates)
    ):
        pair_basis = {
            "study_id": study_id,
            "split": split,
            "trajectory": trajectory,
            "factors": factors,
            "replicate": replicate,
            "root_seed": root_seed,
        }
        pair_id = f"pair-{stable_hash(pair_basis, prefix='pair')[:16]}"
        seeds = named_seed_record(root_seed, pair_basis)
        pair_ids.add(pair_id)
        for method in methods:
            trial_basis = {
                "pair_id": pair_id,
                "method": method,
                "seeds": seeds,
            }
            trial_id = f"trial-{stable_hash(trial_basis, prefix='trial')[:16]}"
            trial = {
                "trial_index": len(trials),
                "trial_id": trial_id,
                "pair_id": pair_id,
                "study_id": study_id,
                "split": split,
                "method": method,
                "trajectory": trajectory,
                "factors": factors,
                "replicate": replicate,
                "seeds": seeds,
            }
            trial["trial_spec_hash"] = stable_hash(trial, prefix="trial-spec")
            trials.append(trial)

    spec_hash = stable_hash(spec, prefix="study-spec")
    plan = {
        "kind": "trial_plan",
        "schema_version": SCHEMA_VERSION,
        "study_id": study_id,
        "split": split,
        "specification": spec,
        "spec_hash": spec_hash,
        "storage_formats": dict(STORAGE_FORMATS),
        "pair_count": len(pair_ids),
        "trial_count": len(trials),
        "trials": trials,
    }
    plan["plan_hash"] = stable_hash(plan, prefix="trial-plan")
    validate_trial_plan(plan)
    return plan


def validate_trial_plan(plan: Mapping[str, Any]) -> None:
    require_schema(plan, "trial_plan")
    trials = require_sequence(plan.get("trials"), "trials")
    if int(plan.get("trial_count", -1)) != len(trials):
        raise ContractError("trial_count does not match trials")
    trial_ids: set[str] = set()
    pairs: dict[str, dict[str, Any]] = {}
    for trial in trials:
        if not isinstance(trial, Mapping):
            raise ContractError("each trial must be a mapping")
        trial_id = require_nonempty_string(trial.get("trial_id"), "trial_id")
        pair_id = require_nonempty_string(trial.get("pair_id"), "pair_id")
        if trial_id in trial_ids:
            raise ContractError(f"duplicate trial_id {trial_id}")
        trial_ids.add(trial_id)
        trial_without_hash = dict(trial)
        recorded_trial_hash = trial_without_hash.pop("trial_spec_hash", None)
        expected_trial_hash = stable_hash(trial_without_hash, prefix="trial-spec")
        if recorded_trial_hash != expected_trial_hash:
            raise ContractError(f"trial_spec_hash mismatch for {trial_id}")
        pair_signature = {
            "trajectory": trial.get("trajectory"),
            "factors": trial.get("factors"),
            "replicate": trial.get("replicate"),
            "seeds": trial.get("seeds"),
            "study_id": trial.get("study_id"),
            "split": trial.get("split"),
        }
        previous = pairs.setdefault(pair_id, pair_signature)
        if previous != pair_signature:
            raise ContractError(f"paired inputs differ inside {pair_id}")
    if int(plan.get("pair_count", -1)) != len(pairs):
        raise ContractError("pair_count does not match unique pair identifiers")
    plan_without_hash = dict(plan)
    recorded_hash = plan_without_hash.pop("plan_hash", None)
    expected_hash = stable_hash(plan_without_hash, prefix="trial-plan")
    if recorded_hash != expected_hash:
        raise ContractError("plan_hash mismatch")
