from pathlib import Path
from typing import Any
from unittest.mock import patch

import numpy as np

from llg3d import parameters
from llg3d.io import (
    Metrics,
    Observables,
    RecordsBuffer,
    RunResults,
    load_results,
    save_results,
    get_tqdm_file,
)


def test_get_tqdm_file():
    """Test that the correct file object is returned for tqdm output."""
    with patch("llg3d.io.size", 4), patch("llg3d.io.rank", 1):
        file_obj = get_tqdm_file()
        assert file_obj is None

    with patch("llg3d.io.size", 4), patch("llg3d.io.rank", 0):
        file_obj = get_tqdm_file()
        assert file_obj is not None

    with patch("llg3d.io.size", 1), patch("llg3d.io.rank", 0):
        file_obj = get_tqdm_file()
        assert file_obj is not None


def test_load_results(
    TEST_DIR: Path, run_ref: RunResults, element_name: str, precision: str
):
    """Test loading results from a .npz file."""

    ref_filename = (
        TEST_DIR / "ref" / f"run_T1100_300x11x11_{element_name}_{precision}.npz"
    )
    data = load_results(ref_filename)
    assert data.params == run_ref.params
    # Handle both legacy and new structure
    if hasattr(data, "results"):
        # New structure
        assert "metrics" in data.results
    else:
        # Legacy structure (flat keys)
        pass  # No assertions for legacy format


def test_save_results(tmp_path: Path, run_params: parameters.RunParameters):
    """Test saving and loading results with new hierarchical structure."""

    metrics: Metrics = {
        "total_time": 123.456,
        "time_per_ite": 0.1234,
        "efficiency": 1e-9,
        "CFL": 0.02,
    }

    records: RecordsBuffer = {
        "xyz_average": list(np.random.rand(10, 2)),
        "x_profiles": {
            "t": list(np.random.rand(10)),
            "m1": list(np.random.rand(10)),
            "m2": list(np.random.rand(10)),
            "m3": list(np.random.rand(10)),
        },
    }

    results_file = tmp_path / "results.npz"
    save_results(results_file, run_params, metrics, records_buffer=records)

    loaded_data = load_results(results_file)
    assert loaded_data.params == run_params
    assert hasattr(loaded_data, "results")
    assert loaded_data.results["metrics"] == metrics

    assert "records" in loaded_data.results
    assert "xyz_average" in loaded_data.results["records"]
    assert "x_profiles" in loaded_data.results["records"]
    assert np.allclose(
        loaded_data.results["records"]["xyz_average"],
        records["xyz_average"],
    )
    assert np.allclose(
        loaded_data.results["records"]["x_profiles"]["m1"],
        records["x_profiles"]["m1"],
    )


def test_extensibility(tmp_path: Path, run_params: parameters.RunParameters):
    """Test that custom fields can be added without modifying io.py."""

    # Custom metrics (any fields allowed) - no type annotation for extensibility
    metrics: Metrics = {
        "total_time": 100.0,
        "time_per_ite": 1.0,
        "efficiency": 0.01,
        "CFL": 0.01,
        "m1_mean": 0.5,
        "custom_metric": 999.0,  # Custom field
        "convergence_rate": 0.95,  # Another custom field
    }  # type: ignore

    # Custom records with various structures
    records: RecordsBuffer = {
        "xyz_average": list(np.random.rand(5, 2)),
        "x_profiles": {
            "t": list(np.random.rand(3)),
            "m1": list(np.random.rand(3)),
            "m2": list(np.random.rand(3)),
            "m3": list(np.random.rand(3)),
        },
        # New top-level custom field
        "energy_evolution": list(np.linspace(0, 100, 20)),
        # Nested custom structure
        "field_snapshots": {
            "times": np.array([0.0, 1.0, 2.0]),
            "data": np.random.rand(3, 10, 5, 5),
        },
        # Deeply nested structure
        "analysis": {
            "spectra": {
                "fourier": np.random.rand(50),
                "wavelets": np.random.rand(50),
            },
            "statistics": {
                "mean": np.array([1.0, 2.0, 3.0]),
                "std": np.array([0.1, 0.2, 0.3]),
            },
        },
    }  # type: ignore

    # Save
    results_file = tmp_path / "extensibility_test.npz"
    save_results(results_file, run_params, metrics, records_buffer=records)

    # Load and verify
    loaded: RunResults = load_results(results_file)

    # Verify custom metrics
    assert loaded.results["metrics"]["custom_metric"] == 999.0
    assert loaded.results["metrics"]["convergence_rate"] == 0.95

    # Verify standard records still work
    assert "xyz_average" in loaded.results["records"]
    assert "x_profiles" in loaded.results["records"]

    # Verify custom top-level record
    assert "energy_evolution" in loaded.results["records"]
    assert loaded.results["records"]["energy_evolution"].shape == (20,)
    np.testing.assert_array_equal(
        loaded.results["records"]["energy_evolution"],
        records["energy_evolution"],  # type: ignore
    )

    # Verify nested custom structure
    assert "field_snapshots" in loaded.results["records"]
    assert "times" in loaded.results["records"]["field_snapshots"]
    assert "data" in loaded.results["records"]["field_snapshots"]
    np.testing.assert_array_equal(
        loaded.results["records"]["field_snapshots"]["times"],
        records["field_snapshots"]["times"],  # type: ignore
    )

    # Verify deeply nested structure
    assert "analysis" in loaded.results["records"]
    assert "spectra" in loaded.results["records"]["analysis"]
    assert "fourier" in loaded.results["records"]["analysis"]["spectra"]
    assert "wavelets" in loaded.results["records"]["analysis"]["spectra"]
    assert "statistics" in loaded.results["records"]["analysis"]
    np.testing.assert_array_equal(
        loaded.results["records"]["analysis"]["spectra"]["fourier"],
        records["analysis"]["spectra"]["fourier"],  # type: ignore
    )


def test_observables(tmp_path: Path, run_params: parameters.RunParameters):
    """Test that observables are saved and loaded correctly."""

    metrics: Metrics = {
        "total_time": 100.0,
        "time_per_ite": 0.01,
        "efficiency": 1e-9,
        "CFL": 0.02,
    }

    observables: Observables = {
        "m1_mean": 0.87,
    }

    records: RecordsBuffer = {"xyz_average": list(np.random.rand(5, 2))}

    results_file = tmp_path / "observables_test.npz"
    save_results(
        results_file,
        run_params,
        metrics,
        observables=observables,
        records_buffer=records,
    )

    loaded: RunResults = load_results(results_file)

    # Verify observables
    assert "observables" in loaded.results
    assert "m1_mean" in loaded.results["observables"]
    assert loaded.results["observables"]["m1_mean"] == 0.87


def test_format_profiling_table_empty():
    """Test format_profiling_table with empty dict."""
    from llg3d.io import format_profiling_table

    result = format_profiling_table({})
    assert result == "(empty)"


def test_format_profiling_table_with_total_time():
    """Test format_profiling_table with total_time (includes % column)."""
    from llg3d.io import format_profiling_table

    profiling_dict = {
        "function_a": {"time": 5.0, "calls": 100},
        "function_b": {"time": 3.0, "calls": 50},
        "function_c": {"time": 2.0, "calls": 25},
    }
    total_time = 10.0

    result = format_profiling_table(profiling_dict, total_time)

    # Check that it contains expected elements
    assert "function_a" in result
    assert "function_b" in result
    assert "function_c" in result
    assert "%" in result  # Percentage column should be present
    assert "50.0%" in result  # function_a: 5.0/10.0 = 50%
    assert "30.0%" in result  # function_b: 3.0/10.0 = 30%
    assert "20.0%" in result  # function_c: 2.0/10.0 = 20%

    # Check sorting by time (descending)
    lines = result.split("\n")
    data_lines = [
        line
        for line in lines
        if line and not line.startswith("-") and "Function" not in line
    ]
    assert "function_a" in data_lines[0]  # Highest time first
    assert "function_b" in data_lines[1]
    assert "function_c" in data_lines[2]

    # Check average time calculation
    assert "0.050000" in result  # function_a avg: 5.0/100 = 0.05


def test_format_profiling_table_without_total_time():
    """Test format_profiling_table without total_time (no % column)."""
    from llg3d.io import format_profiling_table

    profiling_dict = {
        "kernel_x": {"time": 1.5, "calls": 1000},
        "kernel_y": {"time": 0.5, "calls": 500},
    }

    result = format_profiling_table(profiling_dict)

    # Check that functions are present
    assert "kernel_x" in result
    assert "kernel_y" in result

    # Check that percentage column is NOT present
    assert "%" not in result

    # Check that other columns are present
    assert "Calls" in result
    assert "total_time" in result
    assert "Avg Time" in result

    # Verify sorting (by time descending)
    lines = result.split("\n")
    data_lines = [
        line
        for line in lines
        if line and not line.startswith("-") and "Function" not in line
    ]
    assert "kernel_x" in data_lines[0]
    assert "kernel_y" in data_lines[1]


def test_format_profiling_table_zero_calls():
    """Test format_profiling_table handles zero calls gracefully."""
    from llg3d.io import format_profiling_table

    profiling_dict = {
        "never_called": {"time": 0.0, "calls": 0},
        "called_once": {"time": 1.0, "calls": 1},
    }

    result = format_profiling_table(profiling_dict, total_time=2.0)

    # Should handle division by zero for avg_time
    assert "never_called" in result
    assert "called_once" in result
