from __future__ import annotations

import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from airfrans_frontier.runtime import remove_pythonpath_entries

remove_pythonpath_entries()

import numpy as np
import torch

from airfrans_frontier.training.config import load_training_config
from airfrans_frontier.training.loop import select_device


FEATURE_NAMES = np.array(["re_norm", "aoa_norm", "x", "y", "sdf"])
TARGET_NAMES = np.array(["velocity_x", "velocity_y", "pressure", "turbulent_viscosity"])


def write_toy_simulator_dataset(root: Path, *, cases: int = 4, points: int = 64) -> None:
    root.mkdir(parents=True)
    rng = np.random.default_rng(123)
    for case_index in range(cases):
        re_norm = np.full(points, -0.5 + 0.25 * case_index, dtype=np.float32)
        aoa_norm = np.full(points, -0.2 + 0.15 * case_index, dtype=np.float32)
        x = rng.uniform(-1.0, 1.0, size=points).astype(np.float32)
        y = rng.uniform(-1.0, 1.0, size=points).astype(np.float32)
        sdf = (np.sqrt(x * x + y * y) - 0.5).astype(np.float32)
        features = np.stack([re_norm, aoa_norm, x, y, sdf], axis=1).astype(np.float32)
        targets = np.stack(
            [
                0.5 * x + 0.2 * y + 0.1 * aoa_norm,
                -0.3 * x + 0.7 * sdf,
                x * y + 0.05 * re_norm,
                sdf**2 + 0.1 * y,
            ],
            axis=1,
        ).astype(np.float32)
        np.savez(
            root / f"case_{case_index:02d}.npz",
            features=features,
            targets=targets,
            feature_names=FEATURE_NAMES,
            target_names=TARGET_NAMES,
        )


def write_training_config(
    path: Path,
    *,
    data_root: Path,
    artifact_dir: Path,
    device_type: str,
    allow_cpu_fallback: bool = False,
) -> None:
    path.write_text(
        f"""
[run]
name = "test_mlp"
seed = 0
artifact_dir = "{artifact_dir}"

[data]
root = "{data_root}"
train_cases = 2
val_cases = 1
test_cases = 1
points_per_case = 64
batch_size = 64

[model]
type = "mlp"
hidden_width = 128
depth = 4
activation = "gelu"

[optim]
lr = 0.01
weight_decay = 0.0
steps = 1000
log_interval = 250

[device]
type = "{device_type}"
allow_cpu_fallback = {str(allow_cpu_fallback).lower()}
benchmark_kernels = true

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


class TrainingLoopTests(unittest.TestCase):
    def test_cuda_config_fails_clearly_when_cuda_unavailable(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            tmp_path = Path(tmp)
            config_path = tmp_path / "config.toml"
            write_training_config(
                config_path,
                data_root=tmp_path / "data",
                artifact_dir=tmp_path / "artifacts",
                device_type="cuda",
            )
            config = load_training_config(config_path)

            with patch("torch.cuda.is_available", return_value=False):
                with self.assertRaisesRegex(RuntimeError, "CUDA requested"):
                    select_device(config)

    def test_mlp_training_smoke_learns_tiny_deterministic_simulator(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            tmp_path = Path(tmp)
            data_root = tmp_path / "data"
            artifact_dir = tmp_path / "artifacts"
            config_path = tmp_path / "config.toml"
            write_toy_simulator_dataset(data_root, cases=4, points=64)
            device_type = "cuda" if torch.cuda.is_available() else "cpu"
            write_training_config(
                config_path,
                data_root=data_root,
                artifact_dir=artifact_dir,
                device_type=device_type,
            )

            result = subprocess.run(
                [sys.executable, "-m", "airfrans_frontier.cli", "train", str(config_path)],
                text=True,
                capture_output=True,
                check=False,
            )
            self.assertEqual(result.returncode, 0, msg=f"stdout={result.stdout}\nstderr={result.stderr}")
            run_dir_line = next(line for line in result.stdout.splitlines() if line.startswith("run_dir: "))
            run_dir = Path(run_dir_line.removeprefix("run_dir: "))
            final_metrics = json.loads((run_dir / "final_metrics.json").read_text())

            self.assertTrue(np.isfinite(final_metrics["train_loss"]))
            self.assertLess(final_metrics["train_loss"], final_metrics["initial_train_loss"] * 0.1)
            self.assertLess(final_metrics["train_loss"], 1e-2)
            self.assertTrue((run_dir / "checkpoint.pt").exists())
            self.assertTrue((run_dir / "metrics.jsonl").exists())
            self.assertEqual(final_metrics["device"].startswith("cuda"), torch.cuda.is_available())
            if torch.cuda.is_available():
                self.assertIn("T550", final_metrics["gpu_name"])


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