from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, Iterable, List

import xarray as xr
import yaml


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


def load_member_specs(path: str | Path = 'config/member_registry.yaml') -> Dict[str, Dict[str, Any]]:
    data = load_yaml(path)
    specs = {}
    for m in data.get('members', []):
        specs[str(m['name'])] = dict(m)
    return specs


def infer_member_spec(member: str, registry_path: str | Path = 'config/member_registry.yaml') -> Dict[str, Any]:
    """Resolve member name to method/predictor_set/spatial_mode.

    Registry is authoritative.  The fallback is for manually supplied names such
    as `ridge_P2`, `random_forest_P4`, or `ridge_P4_regional`.
    """
    specs = load_member_specs(registry_path)
    if member in specs:
        return specs[member]

    spatial = 'regional' if str(member).endswith('_regional') else 'national'
    base = str(member)[:-len('_regional')] if spatial == 'regional' else str(member)
    bits = base.split('_')
    if len(bits) >= 2 and bits[-1].upper().startswith('P'):
        predictor_set = bits[-1]
        method = '_'.join(bits[:-1])
    else:
        predictor_set = ''
        method = base
    return {'name': member, 'method': method, 'predictor_set': predictor_set, 'spatial_mode': spatial}


def bc_full_label(source: str, lead: int, target: str, member: str, registry_path: str | Path = 'config/member_registry.yaml') -> str:
    spec = infer_member_spec(member, registry_path)
    return f"{source}_lead{int(lead)}_{target}_{spec['spatial_mode']}_{spec['predictor_set']}_{spec['method']}"


def bc_full_path(source: str, lead: int, target: str, member: str,
                 corrected_full_dir: str | Path = 'outputs/corrected_full',
                 registry_path: str | Path = 'config/member_registry.yaml') -> Path:
    return Path(corrected_full_dir) / f"bc_full_{bc_full_label(source, lead, target, member, registry_path)}.nc"


def find_bc_full_file(source: str, lead: int, target: str, member: str,
                      corrected_full_dir: str | Path = 'outputs/corrected_full',
                      registry_path: str | Path = 'config/member_registry.yaml') -> Path:
    """Find a BC full-period file robustly.

    The exact expected filename is tried first.  Then several filename-pattern
    fallbacks are tried.  Finally, attributes are scanned as a last resort.  This
    prevents the previous mismatch between task member names like
    `random_forest_P4` and filenames like `national_P4_random_forest`.
    """
    cdir = Path(corrected_full_dir)
    exact = bc_full_path(source, lead, target, member, cdir, registry_path)
    if exact.exists():
        return exact

    spec = infer_member_spec(member, registry_path)
    patterns = [
        f"bc_full_{source}_lead{int(lead)}_{target}_{spec['spatial_mode']}_{spec['predictor_set']}_{spec['method']}.nc",
        f"bc_full_{source}_lead{int(lead)}_{target}_*_{spec['predictor_set']}_{spec['method']}.nc",
        f"bc_full_{source}_lead{int(lead)}_{target}_*{spec['method']}*.nc",
    ]
    hits: List[Path] = []
    for pat in patterns:
        hits.extend(sorted(cdir.glob(pat)))
    if hits:
        return hits[0]

    for f in sorted(cdir.glob(f'bc_full_{source}_lead{int(lead)}_{target}_*.nc')):
        try:
            ds = xr.open_dataset(f)
            a = ds.attrs
            ok = (
                str(a.get('source_model', '')) == str(source)
                and int(a.get('lead', -999)) == int(lead)
                and str(a.get('target', '')) == str(target)
                and str(a.get('member', '')) == str(member)
            )
            ds.close()
            if ok:
                return f
        except Exception:
            try:
                ds.close()
            except Exception:
                pass
            continue

    raise FileNotFoundError(
        f'Cannot find BC full file for {source} lead{lead} {target} member={member}. '
        f'Tried exact path: {exact}'
    )


def read_selected_member_names(output_root: str | Path = 'outputs', registry_path: str | Path = 'config/member_registry.yaml') -> List[str]:
    f = Path(output_root) / 'selected_members' / 'selected_member_names.txt'
    if f.exists():
        names = [x.strip() for x in f.read_text().splitlines() if x.strip()]
        if names:
            return names
    specs = load_yaml(registry_path).get('members', [])
    return [str(m['name']) for m in specs if bool(m.get('use_for_operational_candidate', True))]


def validate_bc_inputs_for_members(sources: Iterable[str], leads: Iterable[int], targets: Iterable[str], members: Iterable[str],
                                   corrected_full_dir: str | Path = 'outputs/corrected_full',
                                   registry_path: str | Path = 'config/member_registry.yaml',
                                   max_report: int = 30) -> None:
    missing = []
    for source in sources:
        for lead in leads:
            for target in targets:
                for member in members:
                    try:
                        find_bc_full_file(source, int(lead), target, member, corrected_full_dir, registry_path)
                    except FileNotFoundError as exc:
                        missing.append(str(exc))
    if missing:
        msg = '\n'.join(missing[:max_report])
        if len(missing) > max_report:
            msg += f'\n... plus {len(missing) - max_report} more missing combinations.'
        raise FileNotFoundError('Missing BC full-period inputs; run/complete Stage 1 first.\n' + msg)