import requests, gzip, shutil, subprocess, sys
from datetime import datetime, timezone
from pathlib import Path

BASE_DIR = Path("/var/www/html/mrms/output")
HEARTBEAT_FILE = Path("/var/www/html/mrms/output/hourly.txt")
CROP_BOX = ["-95.0", "38.0", "-75.0", "24.0"]

URL_MAP = {
    "3h_qpe": "https://mrms.ncep.noaa.gov/2D/RadarOnly_QPE_03H/MRMS_RadarOnly_QPE_03H.latest.grib2.gz",
    "6h_qpe": "https://mrms.ncep.noaa.gov/2D/RadarOnly_QPE_06H/MRMS_RadarOnly_QPE_06H.latest.grib2.gz",
    "12h_qpe": "https://mrms.ncep.noaa.gov/2D/MultiSensor_QPE_12H_Pass1/MRMS_MultiSensor_QPE_12H_Pass1.latest.grib2.gz",
    "24h_qpe": "https://mrms.ncep.noaa.gov/2D/MultiSensor_QPE_24H_Pass1/MRMS_MultiSensor_QPE_24H_Pass1.latest.grib2.gz"
}

def update_heartbeat(name, timestamp):
    data = {}
    if HEARTBEAT_FILE.exists():
        for line in HEARTBEAT_FILE.read_text().splitlines():
            if ',' in line:
                k, v = line.split(',')
                data[k] = v
    data[name] = timestamp
    with open(HEARTBEAT_FILE, "w") as f:
        for k in sorted(data.keys()):
            f.write(f"{k},{data[k]}\n")

def process(name):
    url = URL_MAP.get(name)
    gz, grib = BASE_DIR / f"{name}.gz", BASE_DIR / f"{name}.grib2"
    tif, warp = BASE_DIR / f"{name}_t.tif", BASE_DIR / f"{name}_w.tif"
    cog = BASE_DIR / f"mrms_{name}_southeast_cog.tif"
    try:
        r = requests.get(url, stream=True, timeout=60)
        r.raise_for_status()
        with open(gz, "wb") as f: shutil.copyfileobj(r.raw, f)
        with gzip.open(gz, "rb") as f_in:
            with open(grib, "wb") as f_out: shutil.copyfileobj(f_in, f_out)
        subprocess.run(["gdal_translate", "-of", "GTiff", "-ot", "Float32", "-projwin", *CROP_BOX, "-a_srs", "EPSG:4326", str(grib), str(tif)], check=True)
        subprocess.run(["gdalwarp", "-overwrite", "-s_srs", "EPSG:4326", "-t_srs", "EPSG:3857", "-r", "near", "-dstnodata", "-9999", str(tif), str(warp)], check=True)
        if cog.exists(): cog.unlink()
        subprocess.run(["gdal_translate", "-of", "COG", "-co", "COMPRESS=DEFLATE", "-co", "PREDICTOR=2", str(warp), str(cog)], check=True)
        ts = datetime.now(timezone.utc).isoformat(timespec='seconds').replace("+00:00", "Z")
        update_heartbeat(name, ts)
    finally:
        for f in [gz, grib, tif, warp]:
            if f.exists(): f.unlink()

if __name__ == "__main__":
    if len(sys.argv) > 1: process(sys.argv[1])