from __future__ import annotations

from pathlib import Path
from typing import Dict, Tuple

import numpy as np
import pandas as pd
import xarray as xr

from .dates import days_in_month

OPERATIONAL_TO_HINDCAST_VAR = {
    'tas': 't2',
    'tasmax': 'tx',
    'tasmin': 'tm',
    'pr': 'pr',
    'h500': 'h500',
    'h850': 'h850',
    'olr': 'olr',
    'q2': 'q2',
    'rh2': 'rh2',
    'slp': 'slp',
    'u200': 'u200',
    'u850': 'u850',
    'w700': 'w700',
}


def _replace_missing(ds: xr.Dataset, missing_value: float = -32767.0) -> xr.Dataset:
    out = ds.copy()
    for v in list(out.data_vars):
        arr = out[v]
        fill_values = [missing_value]
        for att in ['_FillValue', 'missing_value']:
            if att in arr.attrs:
                try:
                    fill_values.append(float(arr.attrs[att]))
                except Exception:
                    pass
        mask = xr.zeros_like(arr, dtype=bool)
        for mv in set(fill_values):
            mask = mask | np.isclose(arr, mv)
        out[v] = arr.where(~mask)
    return out


def read_operational_rcm_file(path: str | Path, init_year: int, init_month: int, missing_value: float = -32767.0) -> xr.Dataset:
    """Read operational RCM file and standardize variable names/units.

    Input dimension `time=6` is interpreted as LT=1..6. Output dimension is `lead`.
    Temperature K is converted to degC. Precipitation flux kg m-2 s-1 is converted to mm/month.
    """
    path = Path(path)
    ds = xr.open_dataset(path, decode_times=False)
    ds = _replace_missing(ds, missing_value)
    if 'time' not in ds.dims or int(ds.sizes['time']) != 6:
        raise ValueError(f'Operational file must have time=6 for LT=1..6: {path}')
    leads = np.arange(1, 7, dtype=np.int16)
    target_year = []
    target_month = []
    target_dates = []
    for lead in leads:
        t = pd.Timestamp(int(init_year), int(init_month), 1) + pd.DateOffset(months=int(lead))
        target_year.append(int(t.year))
        target_month.append(int(t.month))
        target_dates.append(t)

    data_vars = {}
    for in_name, out_name in OPERATIONAL_TO_HINDCAST_VAR.items():
        if in_name not in ds:
            continue
        arr = ds[in_name].transpose('time', 'lat', 'lon').values.astype(np.float32)
        if in_name in {'tas', 'tasmax', 'tasmin'}:
            arr = arr - 273.15
            units = 'degC'
        elif in_name == 'pr':
            # Convert flux to monthly total for each target month.
            for k, (yy, mm) in enumerate(zip(target_year, target_month)):
                arr[k] = arr[k] * 86400.0 * days_in_month(yy, mm)
            units = 'mm/month'
        else:
            units = str(ds[in_name].attrs.get('units', ''))
        data_vars[out_name] = (('lead', 'lat', 'lon'), arr, {'units': units, 'source_variable': in_name})

    out = xr.Dataset(
        coords={
            'lead': leads,
            'target_time': ('lead', np.array(target_dates, dtype='datetime64[ns]')),
            'lat': ds['lat'].values,
            'lon': ds['lon'].values,
        },
        data_vars=data_vars,
        attrs={'init_year': int(init_year), 'init_month': int(init_month), 'source_file': str(path)},
    )
    return out