from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List, Tuple
import json
import yaml
import pandas as pd


def read_yaml(path: str | Path) -> Dict[str, Any]:
    with Path(path).open('r', encoding='utf-8') as f:
        return yaml.safe_load(f) or {}


def ensure_dirs(output_root: str | Path) -> Dict[str, Path]:
    out = Path(output_root)
    dirs = {
        'output': out,
        'tasks': out / 'tasks',
        'logs': out / 'logs',
        'status': out / 'status',
        'standardized': out / 'standardized',
        'bc_members': out / 'bc_members',
        'bc_stack': out / 'bc_stack',
        'dp_members': out / 'dp_members',
        'spei_members': out / 'spei_members',
        'spei_best': out / 'spei_best',
        'probability': out / 'probability',
        'issue_outputs': out,
        'monthly_final': out / 'monthly_final',
        'maps': out / 'maps',
        'report': out / 'report',
        'done': out / 'done',
        'model_cache': out / 'model_cache',
    }
    for p in dirs.values():
        p.mkdir(parents=True, exist_ok=True)
    return dirs


def issue_date(year: int, month: int, suffix: str, issue_days: Dict[str, int]) -> pd.Timestamp:
    day = int(issue_days[str(suffix).zfill(2)])
    return pd.Timestamp(int(year), int(month), int(day))


def ym_iter(start_ym: str, end_ym: str):
    t = pd.Timestamp(start_ym + '-01')
    e = pd.Timestamp(end_ym + '-01')
    while t <= e:
        yield int(t.year), int(t.month)
        t = t + pd.DateOffset(months=1)


def latest_ym_under(root: str | Path) -> str | None:
    root = Path(root)
    found = []
    for ydir in root.glob('[0-9][0-9][0-9][0-9]'):
        for mdir in ydir.glob('[0-9][0-9]'):
            if any(mdir.glob('*.nc')):
                found.append(f'{ydir.name}-{mdir.name}')
    return max(found) if found else None


def write_status(path: Path, obj: Dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(obj, indent=2, default=str), encoding='utf-8')