from __future__ import annotations

import json
import re
import subprocess
import tempfile
from pathlib import Path
import xml.etree.ElementTree as ET

import numpy as np
import rasterio

# ---------------- CONFIG ----------------
BASE_DIR = Path.cwd()          # run from project cwd
RAW_SUBFOLDER = "raw_nc"       # structure: cwd/<parameter>/raw_nc/*.nc

# If set, force a specific analysed variable name (e.g. "t_an").
# Otherwise, auto-detect exactly one "*_an" per NetCDF (recommended).
SUBDATASET_VAR: str | None = None

# Depth sets
DEPTHS_SEASONAL_M = [0, 50, 100, 200, 500]
DEPTHS_ANNUAL_M   = [0, 50, 100, 200, 500, 1000, 2000]

# Crop south of 50S in lon/lat (force lon,lat ordering)
CROP_TE_4326 = (-180.0, -90.0, 180.0, -50.0)
CROP_SRS_4326 = "OGC:CRS84"    # lon,lat order

TARGET_SRS = "EPSG:3031"

# Nodata conversion
DST_NODATA = -9999.0

# Robust handling: treat any absurdly large values as fill
# (Float32 variants of 9.96921e+36 won’t always match exactly)
FILL_THRESHOLD = 1e35

# Skip outputs that already exist
SKIP_IF_OUTPUT_EXISTS = True

# GDAL location
GDAL_BIN = r"C:\Program Files\QGIS 3.40.4\bin"  # set "" if GDAL already on PATH
# ----------------------------------------


def exe(name: str) -> str:
    return str(Path(GDAL_BIN) / name) if GDAL_BIN else name


GDALINFO = exe("gdalinfo.exe")
GDALTRANSLATE = exe("gdal_translate.exe")
GDALWARP = exe("gdalwarp.exe")


def run(cmd: list[str]) -> None:
    print(" ".join(cmd), flush=True)
    subprocess.run(cmd, check=True)


def capture(cmd: list[str]) -> str:
    p = subprocess.run(cmd, check=True, capture_output=True, text=True)
    return p.stdout


def cleanup_ovr_tmp(folder: Path) -> None:
    """Remove leftover GDAL overview temp files (*.ovr.tmp)."""
    for f in folder.glob("*.ovr.tmp"):
        try:
            print(f"[CLEAN] removing {f}", flush=True)
            f.unlink()
        except Exception:
            pass


def gdalinfo_json_for_var(nc_path: Path, var: str) -> dict:
    src = f'NETCDF:"{nc_path}":{var}'
    return json.loads(capture([GDALINFO, "-json", src]))


def flatten_metadata(info: dict) -> dict[str, str]:
    md: dict[str, str] = {}
    if "metadata" in info:
        for _, v in info["metadata"].items():
            if isinstance(v, dict):
                for kk, vv in v.items():
                    md[kk] = vv
    if "metadata" in info and "" in info["metadata"] and isinstance(info["metadata"][""], dict):
        for kk, vv in info["metadata"][""].items():
            md[kk] = vv
    return md


def detect_analysis_var(nc_path: Path) -> str:
    """
    Auto-detect the analysed mean variable '*_an' by parsing the *text* output of gdalinfo.
    Robust even when gdalinfo -json <file.nc> doesn't expose subdatasets.
    """
    txt = capture([GDALINFO, str(nc_path)])

    # SUBDATASET_*_NAME=NETCDF:"...":<var>
    vars_found = re.findall(r'SUBDATASET_\d+_NAME=NETCDF:"[^"]+":([A-Za-z0-9_]+)', txt)
    candidates = sorted({v for v in vars_found if v.endswith("_an")})

    if not candidates:
        snippet = "\n".join([line for line in txt.splitlines() if "SUBDATASET_" in line][:60])
        raise RuntimeError(
            f"[{nc_path.name}] Could not find any '*_an' subdataset variables.\n"
            f"First SUBDATASET lines seen:\n{snippet or '(none)'}"
        )

    if len(candidates) > 1:
        raise RuntimeError(
            f"[{nc_path.name}] Multiple '*_an' variables found: {candidates}. "
            f"Set SUBDATASET_VAR explicitly to the one you want."
        )

    return candidates[0]


def get_global_title(nc_path: Path, var: str) -> str:
    info = gdalinfo_json_for_var(nc_path, var)
    md = flatten_metadata(info)
    return str(md.get("NC_GLOBAL#title", md.get("title", ""))).strip()


def is_annual_file(nc_path: Path, title: str) -> bool:
    """
    Annual can be recognized two ways:
      1) Title contains 'Annual'
      2) Filename contains '_00_' in the time-period slot (e.g. t00, n00, i00)
    """
    if "annual" in (title or "").lower():
        return True
    # filename examples: woa23_decav91C0_t00_04.nc, woa23_all_n00_01.nc
    return bool(re.search(r"_[a-z]\d{2}_", nc_path.name.lower()) and re.search(r"_[a-z]00_", nc_path.name.lower()))


def season_or_annual_from_title(title: str, annual: bool) -> str:
    """
    If annual, return 'ANNUAL' (caps).
    Else Southern Hemisphere seasons based on month range in parentheses:
      (Jan-Mar) -> summer
      (Apr-Jun) -> autumn
      (Jul-Sep) -> winter
      (Oct-Dec) -> spring
    """
    if annual:
        return "ANNUAL"

    if not title:
        return "season_unknown"

    t = title.lower()
    m = re.search(r"\(([^)]+)\)", t)
    if m:
        rng = m.group(1).replace("–", "-").replace("—", "-").strip()
        if any(x in rng for x in ("jan", "feb", "mar")):
            return "summer"
        if any(x in rng for x in ("apr", "may", "jun")):
            return "autumn"
        if any(x in rng for x in ("jul", "aug", "sep")):
            return "winter"
        if any(x in rng for x in ("oct", "nov", "dec")):
            return "spring"

    # fallback: sometimes titles contain season words
    if "summer" in t:
        return "summer"
    if "autumn" in t or "fall" in t:
        return "autumn"
    if "winter" in t:
        return "winter"
    if "spring" in t:
        return "spring"

    return "season_unknown"


def daterange_from_title(title: str) -> str:
    """
    Extract YYYY-YYYY from title, e.g. '1965-2022', '1971-2000', '1991-2020'.
    If not found, returns 'daterange_unknown'.
    """
    m = re.search(r"(\d{4}\s*-\s*\d{4})", title)
    if not m:
        return "daterange_unknown"
    return m.group(1).replace(" ", "")


def res_from_title(title: str) -> str:
    """
    Extracts '<number> degree' from NC_GLOBAL#title and converts to:
      0.25 degree -> 0p25deg
      1.00 degree -> 1deg   (also 1.0, 1.000 etc)
    """
    m = re.search(r"(\d+(?:\.\d+)?)\s*degree", (title or "").lower())
    if not m:
        return "deg_unknown"

    val = m.group(1)
    try:
        f = float(val)
        # if it's basically an integer, render as Ndeg
        if abs(f - round(f)) < 1e-9:
            return f"{int(round(f))}deg"
        # otherwise keep decimals as p
        s = f"{f}".rstrip("0").rstrip(".")  # 0.250 -> 0.25
        return s.replace(".", "p") + "deg"
    except Exception:
        # fallback string handling
        if "." in val:
            # compress 1.00 -> 1
            val2 = val.rstrip("0").rstrip(".")
            if val2.isdigit():
                return val2 + "deg"
            return val2.replace(".", "p") + "deg"
        return val + "deg"


def get_depth_band_map(nc_path: Path, var: str) -> dict[float, int]:
    info = gdalinfo_json_for_var(nc_path, var)
    md = flatten_metadata(info)

    depth_values = None
    for key in ("NETCDF_DIM_depth_VALUES", "NETCDF_DIM_depth_VALUES_m", "depth#values", "depth_values"):
        if key in md:
            depth_values = md[key]
            break

    if depth_values is None:
        raise RuntimeError(f"[{nc_path.name}] Could not find depth dimension values for {var}.")

    s = str(depth_values).strip().strip("{}")
    parts = [p.strip() for p in s.split(",") if p.strip()]
    depths: list[float] = []
    for ptxt in parts:
        try:
            depths.append(float(ptxt))
        except ValueError:
            pass

    if not depths:
        raise RuntimeError(f"[{nc_path.name}] Parsed zero depth values; unexpected format.")

    return {d: i + 1 for i, d in enumerate(depths)}  # 1-based


def filter_depths_available(depth_band_map: dict[float, int], wanted: list[int]) -> list[int]:
    """
    Keep only requested depths that exist in the file.
    Warn if some are missing; error if none exist.
    """
    avail = set(depth_band_map.keys())
    ok = [d for d in wanted if float(d) in avail]
    if not ok:
        mins = min(avail)
        maxs = max(avail)
        raise RuntimeError(
            f"No requested depths found. Requested={wanted}. Available depth range={mins}..{maxs}"
        )
    missing = [d for d in wanted if float(d) not in avail]
    if missing:
        print(f"[WARN] Missing depths (skipping): {missing}", flush=True)
    return ok


def pick_band_indices(depth_band_map: dict[float, int], wanted_depths: list[int]) -> list[int]:
    out: list[int] = []
    for d in wanted_depths:
        df = float(d)
        if df not in depth_band_map:
            available = sorted(depth_band_map.keys())
            raise RuntimeError(
                f"[depths] Requested {d}m not in file. Available: {available[:10]} ... {available[-10:]}"
            )
        out.append(depth_band_map[df])
    return out


def add_band_descriptions_to_vrt(vrt_path: Path, depths_m: list[int]) -> None:
    """Set <Description> per band as '200m' (no spaces)."""
    tree = ET.parse(vrt_path)
    root = tree.getroot()
    bands = root.findall("VRTRasterBand")
    if len(bands) < len(depths_m):
        raise RuntimeError(f"VRT has {len(bands)} bands but expected {len(depths_m)}.")
    for i, depth in enumerate(depths_m):
        band_el = bands[i]
        desc_el = band_el.find("Description")
        if desc_el is None:
            desc_el = ET.SubElement(band_el, "Description")
        desc_el.text = f"{depth}m"
    tree.write(vrt_path, encoding="UTF-8", xml_declaration=True)


def strip_metadata_with_rasterio(src_tif: Path, dst_tif: Path, band_desc: list[str]) -> None:
    """
    Rewrite to a brand-new GeoTIFF with ONLY:
      - pixels
      - CRS/transform
      - nodata
      - band descriptions

    AND force conversion of any fill-like values to DST_NODATA.
    """
    with rasterio.open(src_tif) as src:
        profile = src.profile.copy()
        profile.update(
            driver="GTiff",
            dtype="float32",
            nodata=float(DST_NODATA),
            tiled=True,
            compress="DEFLATE",
            bigtiff="IF_SAFER",
            interleave="band",
        )

        with rasterio.open(dst_tif, "w", **profile) as dst:
            for b in range(1, src.count + 1):
                data = src.read(b).astype(np.float32, copy=False)
                mask = np.isnan(data) | (data > np.float32(FILL_THRESHOLD))
                if mask.any():
                    data[mask] = np.float32(DST_NODATA)
                    print(f"        fill->nodata: band {b} ({int(mask.sum())} px)", flush=True)

                dst.write(data, b)
                dst.set_band_description(b, band_desc[b - 1])


def process_nc(param_dir: Path, nc_path: Path) -> None:
    param_name = param_dir.name

    var = SUBDATASET_VAR or detect_analysis_var(nc_path)

    title = get_global_title(nc_path, var)
    annual = is_annual_file(nc_path, title)

    season_or_annual = season_or_annual_from_title(title, annual=annual)
    daterange = daterange_from_title(title)
    res = res_from_title(title)

    out_path = param_dir / f"WOA2023_{param_name}_{season_or_annual}_{daterange}_{res}.tif"

    if SKIP_IF_OUTPUT_EXISTS and out_path.exists():
        print(f"[SKIP] exists: {out_path}", flush=True)
        return

    # Choose depth set based on annual vs seasonal, then filter to what exists
    depth_band_map = get_depth_band_map(nc_path, var)
    depths_m = DEPTHS_ANNUAL_M if annual else DEPTHS_SEASONAL_M
    depths_m = filter_depths_available(depth_band_map, depths_m)

    src = f'NETCDF:"{nc_path}":{var}'
    band_idxs = pick_band_indices(depth_band_map, depths_m)

    print("\n" + "=" * 90, flush=True)
    print(f"PARAM:   {param_name}", flush=True)
    print(f"NC:      {nc_path}", flush=True)
    print(f"Var:     {var}", flush=True)
    print(f"Title:   {title or '(no title found)'}", flush=True)
    print(f"Type:    {'ANNUAL' if annual else 'SEASONAL'}", flush=True)
    print(f"Token:   {season_or_annual}", flush=True)
    print(f"Range:   {daterange}", flush=True)
    print(f"Res:     {res}", flush=True)
    print(f"Output:  {out_path}", flush=True)
    print("Depth -> band:", flush=True)
    for d, b in zip(depths_m, band_idxs):
        print(f"  {d}m -> band {b}", flush=True)

    band_desc = [f"{d}m" for d in depths_m]

    with tempfile.TemporaryDirectory() as td:
        td = Path(td)
        vrt = td / "subset.vrt"
        ll_tif = td / "subset_ll.tif"
        ll_crop = td / "subset_ll_S50S.tif"
        warped = td / "subset_3031_S50S.tif"
        stripped = td / "subset_3031_S50S_STRIPPED.tif"

        # 1) VRT subset (selected depth bands)
        cmd = [GDALTRANSLATE, "-of", "VRT"]
        for b in band_idxs:
            cmd += ["-b", str(b)]
        cmd += [src, str(vrt)]
        run(cmd)

        # Set band descriptions BEFORE any COG creation (carried through)
        add_band_descriptions_to_vrt(vrt, depths_m)

        # 2) Materialize lon/lat GeoTIFF and force CRS84
        run(
            [
                GDALTRANSLATE,
                "-of",
                "GTiff",
                "-a_srs",
                CROP_SRS_4326,
                "-co",
                "TILED=YES",
                "-co",
                "COMPRESS=DEFLATE",
                "-co",
                "BIGTIFF=IF_SAFER",
                str(vrt),
                str(ll_tif),
            ]
        )

        # 3) Crop south of 50S in lon/lat (best-effort nodata assignment)
        run(
            [
                GDALWARP,
                "-s_srs",
                CROP_SRS_4326,
                "-t_srs",
                CROP_SRS_4326,
                "-te",
                *map(str, CROP_TE_4326),
                "-te_srs",
                CROP_SRS_4326,
                "-r",
                "bilinear",
                "-multi",
                "-wo",
                "NUM_THREADS=ALL_CPUS",
                "-of",
                "GTiff",
                "-co",
                "TILED=YES",
                "-co",
                "COMPRESS=DEFLATE",
                "-co",
                "BIGTIFF=IF_SAFER",
                "-dstnodata",
                f"{DST_NODATA}",
                str(ll_tif),
                str(ll_crop),
            ]
        )

        # 4) Reproject to EPSG:3031
        run(
            [
                GDALWARP,
                "-s_srs",
                CROP_SRS_4326,
                "-t_srs",
                TARGET_SRS,
                "-r",
                "bilinear",
                "-multi",
                "-wo",
                "NUM_THREADS=ALL_CPUS",
                "-of",
                "GTiff",
                "-co",
                "TILED=YES",
                "-co",
                "COMPRESS=DEFLATE",
                "-co",
                "BIGTIFF=IF_SAFER",
                "-dstnodata",
                f"{DST_NODATA}",
                str(ll_crop),
                str(warped),
            ]
        )

        # 4b) HARD fix + metadata strip (guaranteed fill->-9999 here)
        print("[STRIP] rewriting GeoTIFF to remove metadata + force fill->-9999", flush=True)
        strip_metadata_with_rasterio(warped, stripped, band_desc)

        # 5) Final COG (no post-editing)
        run(
            [
                GDALTRANSLATE,
                "-of",
                "COG",
                "-co",
                "COMPRESS=DEFLATE",
                "-co",
                "BIGTIFF=IF_SAFER",
                "-co",
                "OVERVIEW_RESAMPLING=BILINEAR",
                str(stripped),
                str(out_path),
            ]
        )

    cleanup_ovr_tmp(param_dir)
    cleanup_ovr_tmp(nc_path.parent)
    print(f"[OK] {out_path.name}", flush=True)


def main() -> None:
    print(f"[INFO] BASE_DIR = {BASE_DIR}", flush=True)
    print(f"[INFO] RAW_SUBFOLDER = {RAW_SUBFOLDER}", flush=True)
    print(f"[INFO] SUBDATASET_VAR = {SUBDATASET_VAR or '(auto detect *_an via gdalinfo text)'}", flush=True)
    print(f"[INFO] SKIP_IF_OUTPUT_EXISTS = {SKIP_IF_OUTPUT_EXISTS}", flush=True)
    print(f"[INFO] DEPTHS_SEASONAL_M = {DEPTHS_SEASONAL_M}", flush=True)
    print(f"[INFO] DEPTHS_ANNUAL_M   = {DEPTHS_ANNUAL_M}", flush=True)

    param_dirs = [p for p in BASE_DIR.iterdir() if p.is_dir()]
    print(f"[INFO] Found parameter dirs: {[p.name for p in param_dirs]}", flush=True)

    did_any = False

    for param_dir in sorted(param_dirs):
        raw_dir = param_dir / RAW_SUBFOLDER
        if not raw_dir.exists():
            print(f"[INFO] {param_dir.name}: no '{RAW_SUBFOLDER}' folder, skipping", flush=True)
            continue

        nc_files = sorted(raw_dir.glob("*.nc"))
        print(f"[INFO] {param_dir.name}: found {len(nc_files)} .nc files in {raw_dir}", flush=True)

        for nc in nc_files:
            did_any = True
            process_nc(param_dir, nc)

    if not did_any:
        print(
            f"[WARN] No NetCDFs processed. Expected: {BASE_DIR}\\<parameter>\\{RAW_SUBFOLDER}\\*.nc",
            flush=True,
        )


if __name__ == "__main__":
    main()