from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

from airfrans_frontier.training.config import load_training_config


class TrainingConfigTests(unittest.TestCase):
    def test_config_loader_accepts_mlp_tiny(self) -> None:
        config = load_training_config("configs/mlp_tiny.toml")

        self.assertEqual(config.run.name, "mlp_tiny")
        self.assertEqual(config.model.type, "mlp")
        self.assertEqual(config.loss.type, "normalized_mse")
        self.assertEqual(config.device.type, "cuda")
        self.assertTrue(config.data.root.is_absolute())
        self.assertEqual(config.data.source, "local")
        self.assertIsNone(config.data.hf_repo_id)
        self.assertIsNone(config.data.cache_dir)

    def test_config_loader_accepts_huggingface_data_source(self) -> None:
        config = load_training_config("configs/aggressive_smoke.toml")

        self.assertEqual(config.data.source, "huggingface")
        self.assertEqual(config.data.hf_repo_id, "zacheryasc/airfrans-processed")
        self.assertEqual(config.data.hf_path_prefix, "processed/full")
        self.assertTrue(config.data.cache_dir is not None)

    def test_config_loader_parses_model_metadata_and_defaults_old_configs(self) -> None:
        config = load_training_config("configs/mlp_tiny.toml")
        self.assertEqual(config.model_metadata.reported_family, "mlp")
        self.assertFalse(config.model_metadata.is_proxy)

        with tempfile.TemporaryDirectory() as tmp:
            config_path = Path(tmp) / "proxy.toml"
            config_path.write_text(
                f"""
[run]
name = "proxy"
seed = 0
artifact_dir = "{Path(tmp) / "runs"}"

[data]
root = "{Path(tmp) / "data"}"
train_cases = 1
val_cases = 0
test_cases = 0
points_per_case = 1
batch_size = 1

[model]
type = "raster_fno_unet"
hidden_width = 8
depth = 1
activation = "gelu"

[model_metadata]
requested_family = "raster_fno_unet"
implementation_family = "raster_fno_unet"
reported_family = "raster_fno_unet_proxy"
is_proxy = true
proxy_for = "fno_or_raster_field_model"
proxy_notes = "proxy"
coordinate_encoding_compatibility = "raw_only"

[optim]
lr = 0.001
weight_decay = 0.0
steps = 1

[device]
type = "cpu"
allow_cpu_fallback = false
benchmark_kernels = false

[loss]
type = "normalized_mse"
""".strip()
                + "\n"
            )

            parsed = load_training_config(config_path)

            self.assertEqual(parsed.model_metadata.reported_family, "raster_fno_unet_proxy")
            self.assertTrue(parsed.model_metadata.is_proxy)
            self.assertEqual(parsed.model_metadata.coordinate_encoding_compatibility, "raw_only")

    def test_config_loader_accepts_all_points_dead_curve_and_lightweight_checkpoint(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            config_path = Path(tmp) / "all_points.toml"
            config_path.write_text(
                f"""
[run]
name = "all_points"
seed = 0
artifact_dir = "{Path(tmp) / "runs"}"

[data]
root = "{Path(tmp) / "data"}"
train_cases = 1
val_cases = 0
test_cases = 0
all_points_per_case = true
batch_size = 1

[coordinate_encoding]
type = "nerf_multires"
features = ["x", "y", "sdf"]
levels = 4

[model]
type = "mlp_encoded_baseline"
hidden_width = 8
depth = 1
activation = "gelu"
coordinate_features = ["x", "y", "sdf"]

[optim]
lr = 0.001
weight_decay = 0.0
steps = 1

[device]
type = "cpu"
allow_cpu_fallback = false
benchmark_kernels = false

[loss]
type = "normalized_mse"

[checkpoint]
policy = "lightweight_scaling_probe"
include_optimizer_state = false
include_rng_state = false
interval_seconds = 0

[stability]
dead_curve_patience_evals = 3
dead_curve_min_relative_improvement = 0.01
dead_curve_warmup_steps = 5
""".strip()
                + "\n"
            )

            parsed = load_training_config(config_path)

        self.assertTrue(parsed.data.all_points_per_case)
        self.assertIsNone(parsed.data.points_per_case)
        self.assertEqual(parsed.model.type, "mlp_encoded_baseline")
        self.assertEqual(parsed.checkpoint.policy, "lightweight_scaling_probe")
        self.assertFalse(parsed.checkpoint.include_optimizer_state)
        self.assertEqual(parsed.stability.dead_curve_patience_evals, 3)


    def test_config_loader_accepts_huggingface_streaming_data_source(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            config_path = Path(tmp) / "hf_streaming.toml"
            config_path.write_text(
                f"""
[run]
name = "hf_streaming"
seed = 0
artifact_dir = "{Path(tmp) / "runs"}"

[data]
root = "{Path(tmp) / "cache" / "processed" / "full"}"
source = "huggingface_streaming"
hf_repo_id = "zacheryasc/airfrans-processed"
hf_repo_type = "dataset"
hf_path_prefix = "processed/full"
cache_dir = "{Path(tmp) / "cache"}"
train_cases = 1
val_cases = 0
test_cases = 0
all_points_per_case = true
batch_size = 1
streaming_queue_max_cases = 8
streaming_normalization_cases = 1

[model]
type = "mlp"
hidden_width = 8
depth = 1
activation = "gelu"

[optim]
lr = 0.001
weight_decay = 0.0
steps = 1

[device]
type = "cpu"
allow_cpu_fallback = false
benchmark_kernels = false

[loss]
type = "normalized_mse"

[checkpoint]
interval_seconds = 0
""".strip()
                + "\n"
            )

            parsed = load_training_config(config_path)

        self.assertEqual(parsed.data.source, "huggingface_streaming")
        self.assertEqual(parsed.data.hf_repo_id, "zacheryasc/airfrans-processed")
        self.assertTrue(parsed.data.all_points_per_case)

    def test_config_loader_rejects_missing_section(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            config_path = Path(tmp) / "bad.toml"
            config_path.write_text("[run]\nname = 'bad'\n")

            with self.assertRaisesRegex(ValueError, r"missing \[data\] section"):
                load_training_config(config_path)


if __name__ == "__main__":
    unittest.main()