"""
Zone budget (ZBudget) control file parser for IWFM post-processing.

Reads a zbudget control file that specifies which HDF5 zone-budget
output files to process, zone definition files, unit conversion factors,
and which zones to export.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path

from pyiwfm.io.iwfm_reader import next_data_or_empty as _next_data_or_empty
from pyiwfm.io.iwfm_reader import next_data_value as _next_data_value
from pyiwfm.io.iwfm_reader import resolve_path as _resolve_path


@dataclass
class ZBudgetOutputSpec:
    """Specification for a single zone-budget output in the control file."""

    zone_def_file: Path | None = None
    """Zone definition file path (``None`` if blank)."""

    hdf_file: Path = field(default_factory=Path)
    """HDF5 zone-budget input file."""

    output_file: Path = field(default_factory=Path)
    """Output file path (.xlsx or .txt)."""

    output_interval: str | None = None
    """Print interval (e.g. ``"1MON"``).  ``None`` = same as data."""

    zone_ids: list[int] = field(default_factory=list)
    """1-based zone IDs.  ``[-1]`` = all, ``[]`` = none."""


@dataclass
class ZBudgetControlConfig:
    """Parsed zone-budget control file."""

    area_factor: float = 1.0
    """FACTAROU — area unit conversion factor."""

    area_unit: str = "AC"
    """UNITAROU — output area unit string."""

    volume_factor: float = 1.0
    """FACTVLOU — volume unit conversion factor."""

    volume_unit: str = "AC.FT."
    """UNITVLOU — output volume unit string."""

    cache_size: int = 100
    """NTIME — cache/block size for reading."""

    begin_date: str | None = None
    """Begin date (``MM/DD/YYYY_HH:MM``) or ``None`` for start-of-data."""

    end_date: str | None = None
    """End date (``MM/DD/YYYY_HH:MM``) or ``None`` for end-of-data."""

    zbudgets: list[ZBudgetOutputSpec] = field(default_factory=list)
    """List of zone-budget output specifications."""


def read_zbudget_control(filepath: Path | str) -> ZBudgetControlConfig:
    """Parse an IWFM zone-budget control file.

    Parameters
    ----------
    filepath : Path or str
        Path to the zbudget control / input file.

    Returns
    -------
    ZBudgetControlConfig
        Parsed configuration.
    """
    filepath = Path(filepath)
    base_dir = filepath.parent
    config = ZBudgetControlConfig()

    with open(filepath, encoding="ascii", errors="replace") as f:
        # FACTAROU
        config.area_factor = float(_next_data_value(f))

        # UNITAROU
        config.area_unit = _next_data_value(f)

        # FACTVLOU
        config.volume_factor = float(_next_data_value(f))

        # UNITVLOU
        config.volume_unit = _next_data_value(f)

        # NTIME (cache size)
        config.cache_size = int(_next_data_value(f))

        # NZBUDGET — number of zbudget file groups
        n_zbudget = int(_next_data_value(f))

        # Begin/end dates
        begin_raw = _next_data_value(f)
        config.begin_date = None if begin_raw.strip() in ("*", "") else begin_raw.strip()

        end_raw = _next_data_value(f)
        config.end_date = None if end_raw.strip() in ("*", "") else end_raw.strip()

        # Read each zbudget specification
        for _ in range(n_zbudget):
            spec = ZBudgetOutputSpec()

            # Zone definition file (optional — may be blank)
            zdef_raw = _next_data_or_empty(f)
            if zdef_raw.strip():
                spec.zone_def_file = Path(_resolve_path(base_dir, zdef_raw))
            else:
                spec.zone_def_file = None

            # HDF input file
            hdf_raw = _next_data_value(f)
            spec.hdf_file = Path(_resolve_path(base_dir, hdf_raw))

            # Output file
            out_raw = _next_data_value(f)
            spec.output_file = Path(_resolve_path(base_dir, out_raw))

            # Output print interval
            interval_raw = _next_data_value(f)
            spec.output_interval = (
                None if interval_raw.strip() in ("*", "") else interval_raw.strip()
            )

            # NZPRNT — number of zones to print
            nzprnt = int(_next_data_value(f))

            if nzprnt < 0:
                spec.zone_ids = [-1]
            elif nzprnt == 0:
                spec.zone_ids = []
            else:
                ids: list[int] = []
                for _ in range(nzprnt):
                    ids.append(int(_next_data_value(f)))
                spec.zone_ids = ids

            config.zbudgets.append(spec)

    return config
