from __future__ import annotations

import hashlib
import io
import json
from pathlib import Path
import tarfile
import tempfile
from urllib.request import Request, urlopen

from astropy.coordinates import SkyCoord
from astropy.io import fits
import astropy.units as u
import numpy as np

HERE = Path(__file__).resolve().parent
OUTPUT = HERE / "m31_cgmsum_xleap_m31_local_context_audit.json"
CATALOG_URL = (
    "https://content.cld.iop.org/journals/0067-0049/271/2/62/"
    "revision1/apjsad2ea0_table3.tar.gz"
)
M31_ICRS = SkyCoord(10.6847083 * u.deg, 41.26875 * u.deg, frame="icrs")
RINGS_DEG = ((0.0, 4.4), (4.4, 10.0), (10.0, 20.0), (20.0, 30.0), (30.0, 40.0))


def download_catalog() -> bytes:
    request = Request(CATALOG_URL, headers={"User-Agent": "M31CGM-literature-audit/1.0"})
    with urlopen(request, timeout=90) as response:
        return response.read()


def extract_fits(payload: bytes, directory: Path) -> Path:
    with tarfile.open(fileobj=io.BytesIO(payload), mode="r:gz") as archive:
        members = archive.getmembers()
        if len(members) != 1 or Path(members[0].name).name != members[0].name:
            raise RuntimeError("unexpected X-LEAP archive contents")
        if not members[0].name.endswith(".fits"):
            raise RuntimeError("X-LEAP archive does not contain the expected FITS table")
        archive.extract(members[0], path=directory, filter="data")
        return directory / members[0].name


def detection_summary(values: np.ndarray, measured: np.ndarray) -> dict[str, object]:
    selected = np.asarray(values[measured], dtype=float)
    if selected.size == 0:
        return {"detection_count": 0, "median_lu": None, "p16_lu": None, "p84_lu": None}
    return {
        "detection_count": int(selected.size),
        "median_lu": float(np.median(selected)),
        "p16_lu": float(np.quantile(selected, 0.16)),
        "p84_lu": float(np.quantile(selected, 0.84)),
    }


def main() -> None:
    payload = download_catalog()
    with tempfile.TemporaryDirectory(prefix="xleap_m31_audit_") as temporary:
        fits_path = extract_fits(payload, Path(temporary))
        with fits.open(fits_path, memmap=False) as hdul:
            data = hdul[1].data
            sky = SkyCoord(
                np.asarray(data["GLON"], dtype=float) * u.deg,
                np.asarray(data["GLAT"], dtype=float) * u.deg,
                frame="galactic",
            )
            separation = sky.separation(M31_ICRS.galactic).deg
            flags = np.char.strip(np.asarray(data["Flag"]).astype(str))
            clean = flags == "0"
            rings: list[dict[str, object]] = []
            for low, high in RINGS_DEG:
                within = (separation >= low) & (separation < high)
                selected = within & clean
                rings.append(
                    {
                        "low_deg": low,
                        "high_deg": high,
                        "all_row_count": int(np.sum(within)),
                        "clean_row_count": int(np.sum(selected)),
                        "clean_unique_obsid_count": int(
                            np.unique(np.asarray(data["ObsID"])[selected]).size
                        ),
                        "o_vii": detection_summary(
                            np.asarray(data["I_O7"])[selected],
                            np.asarray(data["I_O7_MorU"], dtype=bool)[selected],
                        ),
                        "o_viii": detection_summary(
                            np.asarray(data["I_O8"])[selected],
                            np.asarray(data["I_O8_MorU"], dtype=bool)[selected],
                        ),
                        "fe_l": detection_summary(
                            np.asarray(data["I_FeL"])[selected],
                            np.asarray(data["I_FeL_MorU"], dtype=bool)[selected],
                        ),
                    }
                )

    m31_galactic = M31_ICRS.galactic
    theta_gc_deg = float(
        m31_galactic.separation(SkyCoord(0.0 * u.deg, 0.0 * u.deg, frame="galactic")).deg
    )
    exponential = float(np.exp(-theta_gc_deg / 180.0))
    xleap2_log_temperature = 0.9 * exponential + 5.4
    xleap2_log_temperature_sigma = float(np.hypot(0.1 * exponential, 0.2))
    xleap2_emission_measure = (21.0 * exponential - 12.0) * 1.0e-3
    xleap2_emission_measure_sigma = float(
        np.hypot(13.0 * exponential, 11.0) * 1.0e-3
    )

    result = {
        "schema": "m31-cgmsum-xleap-local-context-audit-v1",
        "catalog": {
            "paper": "Pan et al. 2024, X-LEAP I",
            "doi": "10.3847/1538-4365/ad2ea0",
            "url": CATALOG_URL,
            "download_sha256": hashlib.sha256(payload).hexdigest(),
            "row_count": 5470,
            "quantity": "observed O VII, O VIII, and Fe-L line intensities",
            "correction_boundary": (
                "catalog intensities are not the paper's SWCX- and absorption-corrected map values"
            ),
        },
        "m31": {
            "ra_deg": float(M31_ICRS.ra.deg),
            "dec_deg": float(M31_ICRS.dec.deg),
            "galactic_l_deg": float(m31_galactic.l.deg),
            "galactic_b_deg": float(m31_galactic.b.deg),
            "theta_from_galactic_center_deg": theta_gc_deg,
        },
        "rings": rings,
        "xleap2_empirical_profile_at_m31": {
            "paper": "Qu et al. 2024, X-LEAP II",
            "doi": "10.3847/1538-4357/ad31a0",
            "log10_temperature_K": xleap2_log_temperature,
            "log10_temperature_sigma_independent": xleap2_log_temperature_sigma,
            "emission_measure_pc_cm-6": xleap2_emission_measure,
            "emission_measure_sigma_independent_pc_cm-6": xleap2_emission_measure_sigma,
            "interpretation": (
                "the nominal empirical radial EM is non-positive at the M31 angular distance "
                "and has a large uncertainty; X-LEAP II states that a 3D density decomposition "
                "requires improved SWCX modeling and is future work"
            ),
        },
        "figure3_decision": {
            "use_as_primary_mw_broadband_prior": False,
            "reasons": [
                "X-LEAP I flags every catalog row within 4.4 deg of M31, so there is no clean on-footprint MW-only measurement.",
                "X-LEAP I reports an M31-associated X-ray enhancement extending roughly 10--20 deg, making the nearest annulus potentially circular for an M31/MW decomposition.",
                "The public catalog contains observed line intensities, not a component-separated absorbed 0.4--1.25 keV MW broadband flux.",
                "X-LEAP II constrains temperature structure but does not publish the 3D density posterior needed for a direction-specific broadband projection.",
            ],
            "allowed_role": "local line-intensity and temperature-structure validation only",
        },
    }
    OUTPUT.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    print(OUTPUT)


if __name__ == "__main__":
    main()