from __future__ import annotations

import tomllib
from dataclasses import dataclass
from pathlib import Path
from typing import Any


@dataclass(frozen=True)
class RunConfig:
    name: str
    seed: int
    artifact_dir: Path


@dataclass(frozen=True)
class DataConfig:
    root: Path
    train_cases: int
    val_cases: int
    test_cases: int
    points_per_case: int
    batch_size: int


@dataclass(frozen=True)
class ModelConfig:
    type: str
    hidden_width: int
    depth: int
    activation: str
    coordinate_features: tuple[str, ...]
    fourier_scales: tuple[float, ...]
    condition_width: int
    condition_depth: int
    condition_dim: int


@dataclass(frozen=True)
class OptimConfig:
    lr: float
    weight_decay: float
    steps: int
    log_interval: int | None = None


@dataclass(frozen=True)
class CheckpointConfig:
    interval_seconds: int


@dataclass(frozen=True)
class StabilityConfig:
    max_grad_norm: float | None


@dataclass(frozen=True)
class PrecisionConfig:
    dtype: str


@dataclass(frozen=True)
class DeviceConfig:
    type: str
    allow_cpu_fallback: bool
    benchmark_kernels: bool


@dataclass(frozen=True)
class LossConfig:
    type: str


@dataclass(frozen=True)
class ObservabilityConfig:
    backend: str
    project: str
    entity: str | None
    mode: str
    tags: tuple[str, ...]


@dataclass(frozen=True)
class TrainingConfig:
    path: Path
    config_text: str
    run: RunConfig
    data: DataConfig
    model: ModelConfig
    optim: OptimConfig
    device: DeviceConfig
    loss: LossConfig
    checkpoint: CheckpointConfig
    stability: StabilityConfig
    precision: PrecisionConfig
    observability: ObservabilityConfig

_REQUIRED_SECTIONS = ("run", "data", "model", "optim", "device", "loss")


def load_training_config(path: str | Path) -> TrainingConfig:
    config_path = Path(path).expanduser()
    if not config_path.exists():
        raise FileNotFoundError(f"Training config not found: {config_path}")
    if not config_path.is_file():
        raise ValueError(f"Training config is not a file: {config_path}")

    text = config_path.read_text()
    try:
        raw = tomllib.loads(text)
    except tomllib.TOMLDecodeError as exc:
        raise ValueError(f"Invalid TOML config {config_path}: {exc}") from exc

    for section_name in _REQUIRED_SECTIONS:
        if section_name not in raw or not isinstance(raw[section_name], dict):
            raise ValueError(f"Training config missing [{section_name}] section")

    run_raw = raw["run"]
    data_raw = raw["data"]
    model_raw = raw["model"]
    optim_raw = raw["optim"]
    device_raw = raw["device"]
    loss_raw = raw["loss"]
    checkpoint_raw = raw.get("checkpoint", {})
    if checkpoint_raw is None:
        checkpoint_raw = {}
    if not isinstance(checkpoint_raw, dict):
        raise ValueError("Training config [checkpoint] section must be a table")
    stability_raw = raw.get("stability", {})
    if stability_raw is None:
        stability_raw = {}
    if not isinstance(stability_raw, dict):
        raise ValueError("Training config [stability] section must be a table")
    precision_raw = raw.get("precision", {})
    if precision_raw is None:
        precision_raw = {}
    if not isinstance(precision_raw, dict):
        raise ValueError("Training config [precision] section must be a table")
    observability_raw = raw.get("observability", {})
    if observability_raw is None:
        observability_raw = {}
    if not isinstance(observability_raw, dict):
        raise ValueError("Training config [observability] section must be a table")


    run = RunConfig(
        name=_string(run_raw, "name"),
        seed=_integer(run_raw, "seed", minimum=0),
        artifact_dir=_path(run_raw, "artifact_dir"),
    )
    data = DataConfig(
        root=_path(data_raw, "root"),
        train_cases=_integer(data_raw, "train_cases", minimum=1),
        val_cases=_integer(data_raw, "val_cases", minimum=0),
        test_cases=_integer(data_raw, "test_cases", minimum=0),
        points_per_case=_integer(data_raw, "points_per_case", minimum=1),
        batch_size=_integer(data_raw, "batch_size", minimum=1),
    )
    model = ModelConfig(
        type=_choice(_string(model_raw, "type"), {"mlp", "film_fourier_mlp"}, "model.type"),
        hidden_width=_integer(model_raw, "hidden_width", minimum=1),
        depth=_integer(model_raw, "depth", minimum=1),
        activation=_choice(_string(model_raw, "activation").lower(), {"gelu", "relu", "silu", "tanh"}, "model.activation"),
        coordinate_features=_string_tuple(model_raw, "coordinate_features", default=("x", "y", "sdf")),
        fourier_scales=_number_tuple(model_raw, "fourier_scales", default=(1.0, 2.0, 4.0, 8.0, 16.0)),
        condition_width=_integer(model_raw, "condition_width", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)),
        condition_depth=_integer(model_raw, "condition_depth", minimum=1, default=2),
        condition_dim=_integer(model_raw, "condition_dim", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)),
    )
    optim = OptimConfig(
        lr=_number(optim_raw, "lr", minimum=0.0, exclusive_minimum=True),
        weight_decay=_number(optim_raw, "weight_decay", minimum=0.0),
        steps=_integer(optim_raw, "steps", minimum=1),
        log_interval=_optional_integer(optim_raw, "log_interval", minimum=1),
    )
    device = DeviceConfig(
        type=_choice(_string(device_raw, "type").lower(), {"cuda", "cpu", "auto"}, "device.type"),
        allow_cpu_fallback=_boolean(device_raw, "allow_cpu_fallback"),
        benchmark_kernels=_boolean(device_raw, "benchmark_kernels"),
    )
    loss = LossConfig(type=_choice(_string(loss_raw, "type"), {"normalized_mse"}, "loss.type"))
    checkpoint = CheckpointConfig(
        interval_seconds=_integer(checkpoint_raw, "interval_seconds", minimum=0, default=1800),
    )
    stability = StabilityConfig(
        max_grad_norm=_optional_number(stability_raw, "max_grad_norm", minimum=0.0, exclusive_minimum=True),
    )
    precision = PrecisionConfig(
        dtype=_choice(_string(precision_raw, "dtype", default="float32").lower(), {"float32", "bf16"}, "precision.dtype"),
    )
    observability = ObservabilityConfig(
        backend=_choice(_string(observability_raw, "backend", default="none").lower(), {"none", "wandb"}, "observability.backend"),
        project=_string(observability_raw, "project", default="airfrans"),
        entity=_optional_string(observability_raw, "entity"),
        mode=_choice(_string(observability_raw, "mode", default="online").lower(), {"online", "offline", "disabled"}, "observability.mode"),
        tags=_string_tuple(observability_raw, "tags", default=()),
    )


    requested_cases = data.train_cases + data.val_cases + data.test_cases
    if requested_cases <= 0:
        raise ValueError("Training config must request at least one case")

    return TrainingConfig(
        path=config_path,
        config_text=text,
        run=run,
        data=data,
        model=model,
        optim=optim,
        device=device,
        loss=loss,
        checkpoint=checkpoint,
        stability=stability,
        precision=precision,
        observability=observability,
    )


def _path(section: dict[str, Any], key: str) -> Path:
    value = _string(section, key)
    path = Path(value).expanduser()
    if path.is_absolute():
        return path
    return Path.cwd() / path


def _string(section: dict[str, Any], key: str, *, default: str | None = None) -> str:
    if key not in section:
        if default is not None:
            return default
        raise ValueError(f"Training config missing key: {key}")
    value = section[key]
    if not isinstance(value, str) or not value:
        raise ValueError(f"Expected non-empty string for {key}")
    return value


def _optional_string(section: dict[str, Any], key: str) -> str | None:
    if key not in section:
        return None
    return _string(section, key)


def _integer(section: dict[str, Any], key: str, *, minimum: int | None = None, default: int | None = None) -> int:
    if key not in section:
        if default is not None:
            return default
        raise ValueError(f"Training config missing key: {key}")
    value = section[key]
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError(f"Expected integer for {key}")
    if minimum is not None and value < minimum:
        raise ValueError(f"Expected {key} >= {minimum}")
    return value


def _optional_integer(section: dict[str, Any], key: str, *, minimum: int | None = None) -> int | None:
    if key not in section:
        return None
    return _integer(section, key, minimum=minimum)


def _number(
    section: dict[str, Any],
    key: str,
    *,
    minimum: float | None = None,
    exclusive_minimum: bool = False,
) -> float:
    value = _required(section, key)
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"Expected number for {key}")
    result = float(value)
    if minimum is not None:
        if exclusive_minimum and result <= minimum:
            raise ValueError(f"Expected {key} > {minimum}")
        if not exclusive_minimum and result < minimum:
            raise ValueError(f"Expected {key} >= {minimum}")
    return result


def _optional_number(
    section: dict[str, Any],
    key: str,
    *,
    minimum: float | None = None,
    exclusive_minimum: bool = False,
) -> float | None:
    if key not in section:
        return None
    return _number(section, key, minimum=minimum, exclusive_minimum=exclusive_minimum)


def _number_tuple(section: dict[str, Any], key: str, *, default: tuple[float, ...]) -> tuple[float, ...]:
    if key not in section:
        return default
    value = section[key]
    if not isinstance(value, list) or not value:
        raise ValueError(f"Expected non-empty number array for {key}")
    result: list[float] = []
    for item in value:
        if isinstance(item, bool) or not isinstance(item, (int, float)):
            raise ValueError(f"Expected number array for {key}")
        result.append(float(item))
    return tuple(result)


def _string_tuple(section: dict[str, Any], key: str, *, default: tuple[str, ...]) -> tuple[str, ...]:
    if key not in section:
        return default
    value = section[key]
    if not isinstance(value, list) or not value:
        raise ValueError(f"Expected non-empty string array for {key}")
    result: list[str] = []
    for item in value:
        if not isinstance(item, str) or not item:
            raise ValueError(f"Expected non-empty string array for {key}")
        result.append(item)
    if len(set(result)) != len(result):
        raise ValueError(f"Expected unique strings for {key}")
    return tuple(result)


def _boolean(section: dict[str, Any], key: str) -> bool:
    value = _required(section, key)
    if not isinstance(value, bool):
        raise ValueError(f"Expected boolean for {key}")
    return value


def _choice(value: str, allowed: set[str], key: str) -> str:
    if value not in allowed:
        allowed_text = ", ".join(sorted(allowed))
        raise ValueError(f"Expected {key} to be one of: {allowed_text}")
    return value


def _required(section: dict[str, Any], key: str) -> Any:
    if key not in section:
        raise ValueError(f"Training config missing key: {key}")
    return section[key]