"""HDF5 I/O utilities with automatic provenance tracking.

This module provides helpers for saving simulation data to HDF5 files
with automatic metadata, versioning, and provenance information.

Author: Generated by AgentBible
"""

import datetime
import os
import sys
from pathlib import Path
from typing import Any, Optional, Union

import h5py
import numpy as np
from numpy.typing import NDArray


class SimulationOutput:
    """Context manager for writing simulation data with provenance.

    Automatically records:
    - Creation timestamp
    - Python version
    - Package versions
    - Git commit (if available)
    - User and hostname

    Example:
        >>> with SimulationOutput("output.h5") as out:
        ...     out.save_parameters(dt=0.001, n_steps=1000)
        ...     out.save_snapshot(t=1.0, density=rho)
    """

    def __init__(
        self,
        filepath: Union[str, Path],
        mode: str = "w",
        compression: str = "gzip",
        compression_opts: int = 4,
    ) -> None:
        """Initialize simulation output file.

        Args:
            filepath: Path to HDF5 output file.
            mode: File mode ('w' for write, 'a' for append).
            compression: Compression algorithm ('gzip', 'lzf', or None).
            compression_opts: Compression level (1-9 for gzip).
        """
        self.filepath = Path(filepath)
        self.mode = mode
        self.compression = compression
        self.compression_opts = compression_opts
        self._file: Optional[h5py.File] = None
        self._snapshot_count = 0

    def __enter__(self) -> "SimulationOutput":
        """Open file and write provenance."""
        self._file = h5py.File(self.filepath, self.mode)
        self._write_provenance()
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Close file."""
        if self._file is not None:
            self._file.close()
            self._file = None

    def _write_provenance(self) -> None:
        """Write provenance metadata to file."""
        assert self._file is not None

        prov = self._file.require_group("provenance")
        prov.attrs["created_at"] = datetime.datetime.now().isoformat()
        prov.attrs["python_version"] = sys.version
        prov.attrs["numpy_version"] = np.__version__
        prov.attrs["h5py_version"] = h5py.__version__
        prov.attrs["user"] = os.environ.get("USER", "unknown")
        prov.attrs["hostname"] = os.uname().nodename

        # Try to get git commit
        try:
            import subprocess

            result = subprocess.run(
                ["git", "rev-parse", "HEAD"],
                capture_output=True,
                text=True,
                timeout=5,
            )
            if result.returncode == 0:
                prov.attrs["git_commit"] = result.stdout.strip()
        except Exception:
            pass

    def save_parameters(self, **params: Any) -> None:
        """Save simulation parameters as attributes.

        Args:
            **params: Keyword arguments of parameter names and values.

        Example:
            >>> out.save_parameters(dt=0.001, dx=0.1, viscosity=1e-3)
        """
        assert self._file is not None

        param_group = self._file.require_group("parameters")
        for key, value in params.items():
            if isinstance(value, np.ndarray):
                param_group.create_dataset(key, data=value)
            else:
                param_group.attrs[key] = value

    def save_snapshot(
        self,
        t: float,
        **fields: NDArray,
    ) -> None:
        """Save a simulation snapshot at time t.

        Args:
            t: Simulation time.
            **fields: Keyword arguments of field names and arrays.

        Example:
            >>> out.save_snapshot(t=1.0, density=rho, velocity=v)
        """
        assert self._file is not None

        snapshots = self._file.require_group("snapshots")
        snap_name = f"snap_{self._snapshot_count:06d}"
        snap = snapshots.create_group(snap_name)
        snap.attrs["time"] = t
        snap.attrs["index"] = self._snapshot_count

        for name, data in fields.items():
            snap.create_dataset(
                name,
                data=data,
                compression=self.compression,
                compression_opts=self.compression_opts,
            )

        self._snapshot_count += 1

    def save_timeseries(
        self,
        name: str,
        times: NDArray[np.floating],
        values: NDArray,
    ) -> None:
        """Save a time series (e.g., global quantities over time).

        Args:
            name: Name of the time series.
            times: Array of time values.
            values: Array of values at each time.
        """
        assert self._file is not None

        timeseries = self._file.require_group("timeseries")
        grp = timeseries.create_group(name)
        grp.create_dataset("times", data=times)
        grp.create_dataset("values", data=values)


def load_simulation(filepath: Union[str, Path]) -> dict[str, Any]:
    """Load simulation data from HDF5 file.

    Args:
        filepath: Path to HDF5 file.

    Returns:
        Dictionary with keys: 'parameters', 'provenance', 'snapshots', 'timeseries'.
    """
    filepath = Path(filepath)
    result: dict[str, Any] = {
        "parameters": {},
        "provenance": {},
        "snapshots": [],
        "timeseries": {},
    }

    with h5py.File(filepath, "r") as f:
        # Load provenance
        if "provenance" in f:
            for key, value in f["provenance"].attrs.items():
                result["provenance"][key] = value

        # Load parameters
        if "parameters" in f:
            param_grp = f["parameters"]
            for key, value in param_grp.attrs.items():
                result["parameters"][key] = value
            for key in param_grp:
                result["parameters"][key] = param_grp[key][:]

        # Load snapshots
        if "snapshots" in f:
            snap_grp = f["snapshots"]
            for snap_name in sorted(snap_grp.keys()):
                snap = snap_grp[snap_name]
                snap_data = {"time": snap.attrs["time"]}
                for field_name in snap:
                    snap_data[field_name] = snap[field_name][:]
                result["snapshots"].append(snap_data)

        # Load timeseries
        if "timeseries" in f:
            ts_grp = f["timeseries"]
            for ts_name in ts_grp:
                ts = ts_grp[ts_name]
                result["timeseries"][ts_name] = {
                    "times": ts["times"][:],
                    "values": ts["values"][:],
                }

    return result
