"""
Input/Output functions.

Extensibility Guide
===================

The I/O system is designed to be fully extensible. You can add arbitrary fields
to `metrics` and `records` without modifying this module.

Adding custom metrics
---------------------

    >>> solver.metrics['custom_value'] = 42.0
    >>> solver.metrics['convergence_rate'] = 0.95

Adding custom records
---------------------

    >>> # Simple array
    >>> solver.records['energy_evolution'] = np.array([...])

    >>> # Nested structure (automatically flattened in .npz)
    >>> solver.records['field_snapshots'] = {
    ...     'times': np.array([0, 1, 2]),
    ...     'data': np.array([...])
    ... }

    >>> # Deeply nested (unlimited depth)
    >>> solver.records['analysis'] = {
    ...     'spectra': {
    ...         'fourier': np.array([...]),
    ...         'wavelets': np.array([...])
    ...     }
    ... }

TypedDicts (`Metrics`, `Records`, etc.) are documentation only. Runtime accepts
any dict[str, Any] structure.
"""

from __future__ import annotations

import sys
from dataclasses import dataclass
from pathlib import Path
from typing import (
    Any,
    NotRequired,
    TextIO,
    TypedDict,
    cast,
)

import numpy as np

from .parameters import RunParameters
from .solvers import rank, size
from .solvers.profiling import ProfilingStats


def get_tqdm_file() -> TextIO | None:
    """Get a TQDM-compatible file for progress bar output in MPI."""
    if size == 1:
        return sys.stdout  # single process uses default
    if rank != 0:
        return None  # other ranks disable progress bar
    try:
        # For MPI, try to open /dev/tty for direct terminal output
        return open("/dev/tty", "w")
    except OSError:
        return sys.stdout  # fallback


class MetricsRequired(TypedDict):
    """Required fields for simulation metrics."""

    total_time: float  #: Total wall-clock time of the simulation
    time_per_ite: float  #: Average time per iteration
    efficiency: float  #: Time per iteration per cell
    CFL: float  #: CFL condition value


class Metrics(MetricsRequired, total=False):
    """
    Structure for simulation metrics (extensible).

    Required fields: total_time, time_per_ite, efficiency, CFL
    All other fields are optional and can be added dynamically.
    """

    #: Profiling statistics
    profiling_stats: NotRequired[ProfilingStats]


class Observables(TypedDict, total=False):
    """
    Physical observables from the simulation (extensible).

    All fields are optional and can be added dynamically.
    """

    m1_mean: float  #: Time-averaged magnetization in x direction


class XProfiles(TypedDict):
    """Structure for cross-sectional profiles (arrays in final records)."""

    t: np.ndarray  #: Time points for profiles
    m1: np.ndarray  #: m1 component profiles
    m2: np.ndarray  #: m2 component profiles
    m3: np.ndarray  #: m3 component profiles


class XProfilesBuffer(TypedDict):
    """Structure for cross-sectional profiles during accumulation (lists)."""

    t: list[float]  #: Time points for profiles
    m1: list[np.ndarray]  #: m1 component profiles
    m2: list[np.ndarray]  #: m2 component profiles
    m3: list[np.ndarray]  #: m3 component profiles


class RecordsBuffer(TypedDict, total=False):
    """
    Records during simulation (accumulation phase with lists).

    BaseSolver.records uses this structure during simulation, accumulating
    data as lists. Before saving, BaseSolver.save() converts lists to arrays.
    """

    xyz_average: list[tuple[float, float]]  #: Accumulated (t, value) pairs
    x_profiles: XProfilesBuffer  #: Profiles during accumulation


class Records(TypedDict, total=False):
    """
    Records after saving (finalized with numpy arrays).

    Returned by load_results() with all data as read-only numpy arrays.
    """

    xyz_average: np.ndarray  #: Space-averaged magnetization over time (shape: (2, N))
    x_profiles: XProfiles  #: Cross-sectional profiles in yz plane


class SimulationResults(TypedDict):
    """Structure for simulation results."""

    metrics: Metrics  #: Simulation metrics (performance, numerical quality)
    observables: NotRequired[Observables]  #: Physical observables (results)
    records: NotRequired[Records]  #: Optional time-series records


@dataclass
class RunResults:
    """Structure for loaded simulation results."""

    params: RunParameters
    results: SimulationResults
    file: str | Path

    def get_record(self, record_name: str) -> np.ndarray | dict[str, np.ndarray]:
        """Return a record by name from results or raise a descriptive error."""
        try:
            if "records" not in self.results:
                raise KeyError(f"RunResults from '{self.file}' has no records.")

            records = self.results["records"]
            # Cast to dict to allow dynamic access with a variable key for mypy
            record = cast(dict, records)[record_name]
            if isinstance(record, (np.ndarray, dict)):
                return record
            else:
                raise TypeError(f"'{record_name}' is not in the expected format.")
        except KeyError as e:
            msg = (
                f"RunResults from '{self.file}' does not contain the required record "
                f"'{record_name}'."
            )
            raise KeyError(msg) from e


def save_results(
    output_file: str | Path,
    params: RunParameters,
    metrics: Metrics,
    observables: Observables | None = None,
    records_buffer: RecordsBuffer | None = None,
) -> None:
    """
    Saves simulation results to a .npz file with hierarchical structure.

    Args:
        output_file: Path to the output .npz file.
        params: Dataclass of simulation parameters.
        metrics: Dictionary of simulation metrics. Must contain required fields:
            total_time, time_per_ite, CFL. Additional fields are allowed.
        observables: Dictionary of physical observables (results).
        records_buffer: RecordsBuffer or custom dict of records
            (accumulation phase with lists/dicts).

    Note:
        Both metrics and records accept any key-value pairs, allowing easy
        extension without modifying this function. See Metrics and Records
        TypedDict for recommended structure.

    Example:
        >>> records_buffer = {
        ...     'xyz_average': arr1,
        ...     'x_profiles': {'t': t, 'm1': m1, ...},
        ...     'custom_data': arr2  # Any new field works automatically
        ... }
        >>> save_results('run.npz', params, metrics, records_buffer=records_buffer)
    """
    data_to_save: dict[str, Any] = {}

    if records_buffer is None:
        records_buffer = {}
    if observables is None:
        observables = {}

    # Helper to convert lists to numpy arrays recursively
    def _convert_lists_to_arrays(d: Any) -> Any:
        """Recursively convert lists to numpy arrays."""
        if isinstance(d, list):
            return np.array(d)
        elif isinstance(d, dict):
            return {k: _convert_lists_to_arrays(v) for k, v in d.items()}
        else:
            return d

    # Convert lists in records_buffer to arrays
    records_buffer = _convert_lists_to_arrays(records_buffer)

    # Helper to flatten nested dicts into NPZ keys
    def _flatten_dict(d: dict | Any, prefix: str):
        """Recursively flatten nested dicts into NPZ keys."""
        if isinstance(d, dict):
            for key, value in d.items():
                full_key = f"{prefix}/{key}"
                _flatten_dict(value, full_key)
        else:
            # Leaf value - save as array
            data_to_save[prefix] = d

    # Flatten params
    _flatten_dict(params.as_dict(), "params")

    # Flatten metrics
    _flatten_dict(metrics, "results/metrics")

    # Flatten observables
    if observables:
        _flatten_dict(observables, "results/observables")

    # Flatten records
    _flatten_dict(records_buffer, "results/records")

    np.savez(output_file, **data_to_save, allow_pickle=False)


def load_results(result_file: str | Path) -> RunResults:
    """
    Loads simulation results from a .npz file with hierarchical structure.

    Numpy arrays are set to read-only to prevent accidental modification.

    Example:
        >>> run_results = load_results("run.npz")
        >>> params = run_results["params"]
        >>> metrics = run_results["results"]["metrics"]
        >>> xyz_average = run_results["results"]["records"]["xyz_average"]
        >>> m1_prof = run_results["results"]["records"]["x_profiles"]["m1"]

    Args:
        result_file: path to the .npz file.

    Returns:
        Dictionary containing the loaded data with hierarchical structure.
    """
    npz_data: np.lib.npyio.NpzFile = np.load(result_file, allow_pickle=False)

    def _unflatten_dict(prefix: str) -> dict[str, Any]:
        """Reconstruct nested dict from flat NPZ keys with given prefix."""
        result: dict[str, Any] = {}

        for key in npz_data.files:
            if not key.startswith(prefix + "/"):
                continue

            # Extract relative path after prefix
            rel_path = key[len(prefix) + 1 :]
            parts = rel_path.split("/")

            # Navigate/create nested structure
            current = result
            for part in parts[:-1]:
                if part not in current:
                    current[part] = {}
                current = current[part]

            # Store the value
            arr = npz_data[key]
            arr.flags.writeable = False
            current[parts[-1]] = arr

        return result

    # Load params
    params_dict = _unflatten_dict("params")
    params = RunParameters(**params_dict)

    # Load metrics
    metrics: Metrics = cast(Metrics, _unflatten_dict("results/metrics"))

    # Load observables
    observables: Observables | None = None
    observables_dict = _unflatten_dict("results/observables")
    if observables_dict:
        observables = cast(Observables, observables_dict)

    # Load records
    records: Records = cast(Records, _unflatten_dict("results/records"))

    # Build the results dict
    simulation_results: SimulationResults = {"metrics": metrics}
    if observables:
        simulation_results["observables"] = observables
    if records:
        simulation_results["records"] = records

    run_results = RunResults(
        params=params, results=simulation_results, file=result_file
    )

    return run_results


def format_profiling_table(
    profiling_dict: dict[str, dict[str, float | int]], total_time: float | None = None
) -> str:
    """
    Format profiling statistics as a table.

    Args:
        profiling_dict: Dictionary with profiling stats, where each entry has
            'time' (float) and 'calls' (int) keys
        total_time: Total simulation time for percentage calculation.
            If None, the percentage column is omitted.

    Returns:
        A formatted table string
    """
    if not profiling_dict:
        return "(empty)"

    # Determine the maximum width for alignment
    col1_width = max(len(name) for name in profiling_dict.keys())

    # Header
    if total_time is not None:
        s = (
            f"{'Function':<{col1_width}} | {'Calls':>5} | {'total_time (s)':>14} "
            f"| {'%':>6} | {'Avg Time (s)':>13}\n"
        )
    else:
        s = (
            f"{'Function':<{col1_width}} | {'Calls':>5} | {'total_time (s)':>14} "
            f"| {'Avg Time (s)':>13}\n"
        )
    s += "-" * (len(s) - 1) + "\n"

    # Sort by total time descending
    sorted_profiling = dict(
        sorted(profiling_dict.items(), key=lambda item: item[1]["time"], reverse=True)
    )

    # Data rows
    for name, stats in sorted_profiling.items():
        func_time = stats["time"]
        calls = stats["calls"]
        avg_time = func_time / calls if calls > 0 else 0.0

        if total_time is not None:
            percent = (func_time / total_time * 100) if total_time > 0 else 0.0
            s += (
                f"{name:<{col1_width}} | {calls:>5} | {func_time:>14.6f} "
                f"| {percent:>5.1f}% | {avg_time:>13.6f}\n"
            )
        else:
            s += (
                f"{name:<{col1_width}} | {calls:>5} | {func_time:>14.6f} "
                f"| {avg_time:>13.6f}\n"
            )

    return s
