from __future__ import annotations
from pathlib import Path
from dataclasses import dataclass
import numpy as np
import xarray as xr
from .time_utils import make_time_axis

MODEL_VARS = ['pr','t2','tx','tm','h500','h850','olr','q2','rh2','slp','u200','u850','w700']
OBS_VARS = ['pr_o','t2_o','tx_o','tm_o']
ALL_VARS = OBS_VARS + MODEL_VARS
TARGET_TO_OBS = {'pr':'pr_o', 't2':'t2_o'}
TARGET_TO_MODEL = {'pr':'pr', 't2':'t2'}

@dataclass
class GridData:
    ds: xr.Dataset
    mask: xr.Dataset
    time: object
    years: np.ndarray
    months: np.ndarray
    valid_mask: np.ndarray
    reg: np.ndarray
    lat2d: np.ndarray
    lon2d: np.ndarray


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


def open_lead_dataset(cfg: dict, source_model: str, lead: int) -> xr.Dataset:
    data_dir = Path(cfg['paths']['data_dir'])
    pattern = cfg['paths'].get('file_pattern', 'lead{lead}_{source_model}_01_BL_Test_VN.nc')
    path = data_dir / pattern.format(lead=int(lead), source_model=source_model)
    if not path.exists():
        raise FileNotFoundError(path)
    ds = xr.open_dataset(path, decode_times=False)
    for name in list(ds.data_vars):
        if {'mon','lat','lon'}.issubset(ds[name].dims):
            ds[name] = ds[name].transpose('mon','lat','lon')
    return _replace_missing(ds, float(cfg['grid'].get('missing_value', -99.0)))


def open_mask(cfg: dict) -> xr.Dataset:
    path = Path(cfg['paths']['mask_file'])
    if not path.exists():
        raise FileNotFoundError(path)
    ds = xr.open_dataset(path, decode_times=False)
    for name in list(ds.data_vars):
        if {'lat','lon'}.issubset(ds[name].dims):
            ds[name] = ds[name].transpose('lat','lon')
    return _replace_missing(ds, float(cfg['grid'].get('missing_value', -99.0)))


def load_grid(cfg: dict, source_model: str, lead: int) -> GridData:
    ds = open_lead_dataset(cfg, source_model, lead)
    mask = open_mask(cfg)
    if ds.sizes['lat'] != mask.sizes['lat'] or ds.sizes['lon'] != mask.sizes['lon']:
        raise ValueError('Input and mask grid sizes differ')
    if not np.allclose(ds['lat'].values, mask['lat'].values) or not np.allclose(ds['lon'].values, mask['lon'].values):
        raise ValueError('Input and mask lat/lon coordinates differ')
    time = make_time_axis(cfg['grid'].get('year_start', 1982), cfg['grid'].get('n_months', ds.sizes['mon']))
    reg = mask['reg'].values.astype(float)
    valid_mask = np.isfinite(reg) & (reg > 0)
    lon2d, lat2d = np.meshgrid(ds['lon'].values, ds['lat'].values)
    return GridData(ds=ds, mask=mask, time=time, years=time.year.to_numpy(), months=time.month.to_numpy(), valid_mask=valid_mask, reg=reg, lat2d=lat2d, lon2d=lon2d)


def write_dataset_safe(ds: xr.Dataset, path: str | Path) -> None:
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    tmp = p.with_suffix(p.suffix + '.tmp')
    if tmp.exists():
        tmp.unlink()
    enc = {v: {'zlib': True, 'complevel': 1} for v in ds.data_vars if ds[v].ndim >= 2}
    try:
        ds.to_netcdf(tmp, encoding=enc)
    except Exception:
        ds.to_netcdf(tmp)
    tmp.replace(p)