from __future__ import annotations

import concurrent.futures as cf
import gzip
import json
import pickle
from pathlib import Path

import numpy as np
import pandas as pd


def build_patch_inputs(
    *,
    patch_def_path: Path,
    out_dir: Path,
    image_dict: dict,
    input_catalog: pd.DataFrame,
    include_wcs_in_payload: bool,
    save_patch_inputs: bool,
    skip_empty_patch: bool,
    max_workers: int,
    gzip_compresslevel: int,
    keep_payloads_in_memory: bool,
) -> list[dict]:
    out_dir.mkdir(parents=True, exist_ok=True)
    assert patch_def_path.exists(), f"Missing patch defs: {patch_def_path}"

    patch_defs = json.loads(patch_def_path.read_text())
    print("Loaded patch defs:", len(patch_defs))
    print("INCLUDE_WCS_IN_PATCH_INPUTS:", include_wcs_in_payload)
    print("SAVE_PATCH_INPUTS:", save_patch_inputs, "| MAX_WORKERS_PATCH_INPUTS:", max_workers, "| GZIP_COMPRESSLEVEL:", gzip_compresslevel)

    N_TOTAL = len(patch_defs)
    N_DONE = 0
    N_KEEP = 0
    N_SKIP = 0
    PROGRESS_EVERY = max(1, N_TOTAL // 20)

    cols = list(input_catalog.columns)
    col_map = {str(c).strip().lower(): str(c) for c in cols}
    _ra_col = col_map.get("ra")
    _dec_col = col_map.get("dec")
    if _ra_col is None or _dec_col is None:
        raise ValueError("input_catalog must have RA/DEC columns")
    if _ra_col != "RA" or _dec_col != "DEC":
        print(f"[WARN] input_catalog uses {_ra_col}/{_dec_col} columns; interpreting as Right Ascension / Declination. Expected format is RA/DEC.")

    ra = pd.to_numeric(input_catalog[_ra_col], errors="coerce").to_numpy(dtype=float)
    dec = pd.to_numeric(input_catalog[_dec_col], errors="coerce").to_numpy(dtype=float)

    wcs_white = next(iter(image_dict.values()))["wcs"]

    ok = np.isfinite(ra) & np.isfinite(dec)
    xpix = np.full(len(ra), np.nan, dtype=float)
    ypix = np.full(len(dec), np.nan, dtype=float)
    if np.any(ok):
        xw, yw = wcs_white.all_world2pix(ra[ok], dec[ok], 0)
        xpix[ok] = xw
        ypix[ok] = yw

    cat_global = input_catalog.copy()
    cat_global["x_pix_white"] = xpix
    cat_global["y_pix_white"] = ypix

    patch_inputs = []
    patch_payloads = [] if keep_payloads_in_memory else None

    xg_all = cat_global["x_pix_white"].to_numpy(dtype=float)
    yg_all = cat_global["y_pix_white"].to_numpy(dtype=float)

    _band_items = [(str(band), d) for band, d in image_dict.items()]

    def _build_payload_for_patch(pdef: dict):
        x0b, x1b, y0b, y1b = pdef["x0_base"], pdef["x1_base"], pdef["y0_base"], pdef["y1_base"]
        x0r, x1r, y0r, y1r = pdef["x0_roi"], pdef["x1_roi"], pdef["y0_roi"], pdef["y1_roi"]

        in_base = (
            np.isfinite(xg_all)
            & np.isfinite(yg_all)
            & (xg_all >= x0b)
            & (xg_all < x1b)
            & (yg_all >= y0b)
            & (yg_all < y1b)
        )
        cat_patch = cat_global.loc[in_base].copy().reset_index(drop=True)
        if skip_empty_patch and len(cat_patch) == 0:
            return None

        cat_patch["x_pix_patch"] = cat_patch["x_pix_white"] - float(x0r)
        cat_patch["y_pix_patch"] = cat_patch["y_pix_white"] - float(y0r)

        per_filter_patch = {}
        for band, d in _band_items:
            img_full = d["img_scaled"]
            sig_total = d["sigma_total_scaled"]
            sig_sky = d["sigma_sky_scaled"]
            bad_full = d.get("bad", None)
            if bad_full is None:
                bad_full = ~np.isfinite(img_full)

            entry = dict(
                img_scaled=img_full[y0r:y1r, x0r:x1r],
                sigma_total_scaled=sig_total[y0r:y1r, x0r:x1r],
                sigma_sky_scaled=sig_sky[y0r:y1r, x0r:x1r],
                bad=np.asarray(bad_full[y0r:y1r, x0r:x1r], dtype=bool),
                scale=float(d.get("scale", 1.0)),
                fp=str(d.get("fp", "")),
                x0w=int(x0r),
                x1w=int(x1r),
                y0w=int(y0r),
                y1w=int(y1r),
            )

            if include_wcs_in_payload:
                wcs_patch = d["wcs"].slice((slice(y0r, y1r), slice(x0r, x1r)))
                entry["wcs_patch"] = wcs_patch
                entry["wcs_full"] = d["wcs"]

            per_filter_patch[band] = entry

        meta = dict(
            patch_tag=pdef["patch_tag"],
            epsf_tag=pdef["epsf_tag"],
            epsf_r=int(pdef["epsf_r"]),
            epsf_c=int(pdef["epsf_c"]),
            patch_r=int(pdef["patch_r"]),
            patch_c=int(pdef["patch_c"]),
            halo_pix=int(pdef["halo_pix"]),
            base_bbox=(int(x0b), int(x1b), int(y0b), int(y1b)),
            roi_bbox=(int(x0r), int(x1r), int(y0r), int(y1r)),
            n_sources=int(len(cat_patch)),
        )

        payload = dict(per_filter_patch=per_filter_patch, cat_patch=cat_patch, meta=meta)
        return meta, payload

    def _write_payload(meta: dict, payload: dict):
        outpath = out_dir / f"{meta['patch_tag']}.pkl.gz"
        with gzip.open(outpath, "wb", compresslevel=int(gzip_compresslevel)) as f:
            pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
        return str(outpath)

    if save_patch_inputs and int(max_workers) > 1:
        futs = {}
        with cf.ThreadPoolExecutor(max_workers=int(max_workers)) as ex:
            for pdef in patch_defs:
                N_DONE += 1
                built = _build_payload_for_patch(pdef)
                if built is None:
                    N_SKIP += 1
                    if (N_DONE % PROGRESS_EVERY) == 0 or N_DONE == N_TOTAL:
                        print(f"[patch inputs] {N_DONE}/{N_TOTAL} processed | keep={N_KEEP} skip={N_SKIP}")
                    continue
                meta, payload = built
                patch_inputs.append(meta)
                if keep_payloads_in_memory:
                    patch_payloads.append(payload)
                N_KEEP += 1
                futs[ex.submit(_write_payload, meta, payload)] = meta["patch_tag"]

                if (N_DONE % PROGRESS_EVERY) == 0 or N_DONE == N_TOTAL:
                    print(f"[patch inputs] {N_DONE}/{N_TOTAL} processed | keep={N_KEEP} skip={N_SKIP} (writing async)")

            ndone_w = 0
            for fut in cf.as_completed(futs):
                _ = fut.result()
                ndone_w += 1
                if ndone_w % max(1, len(futs) // 10) == 0 or ndone_w == len(futs):
                    print(f"[patch inputs] wrote {ndone_w}/{len(futs)} files")
    else:
        for pdef in patch_defs:
            N_DONE += 1
            built = _build_payload_for_patch(pdef)
            if built is None:
                N_SKIP += 1
                if (N_DONE % PROGRESS_EVERY) == 0 or N_DONE == N_TOTAL:
                    print(f"[patch inputs] {N_DONE}/{N_TOTAL} processed | keep={N_KEEP} skip={N_SKIP}")
                continue

            meta, payload = built
            patch_inputs.append(meta)
            if keep_payloads_in_memory:
                patch_payloads.append(payload)
            N_KEEP += 1

            if save_patch_inputs:
                _ = _write_payload(meta, payload)

            if (N_DONE % PROGRESS_EVERY) == 0 or N_DONE == N_TOTAL:
                print(f"[patch inputs] {N_DONE}/{N_TOTAL} processed | keep={N_KEEP} skip={N_SKIP}")

    print("Patch inputs created:", len(patch_inputs))
    if save_patch_inputs:
        print("Saved to:", out_dir)
    else:
        print("Saved to: (skipped)")

    return patch_inputs
