#!/usr/bin/env python3
"""Run the fixed full-catalog dummy diagnostic once without touching images."""

import fcntl
import hashlib
import os
from pathlib import Path
import resource
import shlex
import signal
import subprocess
import time

ROOT = Path(__file__).resolve().parents[2]
LOG = Path(__file__).with_name("full_catalog.log")
BINARY = ROOT / "build/Release/schwarzschild_sky_dummy"
LENS = ROOT / "output/lens/schwarzschild_galactic_center_R100_45deg_16_4_4k_.grlens"
PNG = ROOT / "output/imgs/2mass_galactic_center_schwarzschild_R100_4k_1e13_beta4.5_16_4-full-gpu.png"
FITS = ROOT / "output/imgs/2mass_galactic_center_schwarzschild_R100_4k_1e13_beta4.5_16_4-full-gpu_HDR.fits"


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main():
    if LOG.exists():
        raise RuntimeError(f"refusing to overwrite {LOG}")
    for path in (BINARY, LENS, PNG, FITS):
        if not path.is_file():
            raise RuntimeError(f"required file is missing: {path}")
    protected_before = {
        path: (digest(path), path.stat().st_size, path.stat().st_mtime_ns)
        for path in (PNG, FITS)
    }
    command = [
        str(BINARY), "--verbose",
        "--psf-relative-tail", "1e-8", "--psf-min-y", "0",
        "--max-cache-psf-flux", "1e8",
        "--all-sky-catalog", "assets/2mass/processed/all_sky",
        "--exposure", "1e13", "--psf-fwhm-pixels", "2.7",
        "--psf-moffat-beta", "4.5", "--catalog-load-workers", "4",
        "--hdr-output", "--lens-map-input",
        "output/lens/schwarzschild_galactic_center_R100_45deg_16_4_4k_.grlens",
        "--output",
        "output/imgs/2mass_galactic_center_schwarzschild_R100_4k_1e13_beta4.5_16_4-full-gpu.png",
    ]
    guard = open("/tmp/gr_dummy_psf_chunks.lock", "w")
    fcntl.flock(guard, fcntl.LOCK_EX | fcntl.LOCK_NB)
    env = dict(os.environ, OMP_NUM_THREADS="16", OMP_DYNAMIC="FALSE")
    with LOG.open("x") as out:
        out.write("COMMAND OMP_NUM_THREADS=16 OMP_DYNAMIC=FALSE " +
                  shlex.join(command) + "\n")
        out.write("GIT " + subprocess.check_output(
            ["git", "rev-parse", "HEAD"], cwd=ROOT, text=True))
        out.write("BINARY_SHA256 " + digest(BINARY) + "\n")
        out.write("LENS_SHA256 " + digest(LENS) + "\n")
        for path, state in protected_before.items():
            out.write(
                f"PROTECTED_BEFORE {path.relative_to(ROOT)} sha256={state[0]} "
                f"size={state[1]} mtime_ns={state[2]}\n")
        out.flush()
        start = time.monotonic()
        usage_start = resource.getrusage(resource.RUSAGE_CHILDREN)
        process = subprocess.Popen(command, cwd=ROOT, env=env, stdout=out,
                                   stderr=subprocess.STDOUT,
                                   start_new_session=True)
        try:
            returncode = process.wait(timeout=600)
        except subprocess.TimeoutExpired:
            os.killpg(process.pid, signal.SIGKILL)
            process.wait()
            out.write("TIMEOUT killed and reaped process group\n")
            raise
        out.write(
            f"PROCESS_WALL {time.monotonic() - start:.9f} EXIT {returncode}\n")
        usage_end = resource.getrusage(resource.RUSAGE_CHILDREN)
        out.write(
            f"PROCESS_USER {usage_end.ru_utime - usage_start.ru_utime:.9f} "
            f"PROCESS_SYSTEM {usage_end.ru_stime - usage_start.ru_stime:.9f} "
            f"MAX_RSS_KIB {usage_end.ru_maxrss}\n")
    protected_after = {
        path: (digest(path), path.stat().st_size, path.stat().st_mtime_ns)
        for path in (PNG, FITS)
    }
    if protected_after != protected_before:
        raise RuntimeError("a protected render output changed")
    with LOG.open("a") as out:
        for path, state in protected_after.items():
            out.write(
                f"PROTECTED_AFTER {path.relative_to(ROOT)} sha256={state[0]} "
                f"size={state[1]} mtime_ns={state[2]} unchanged=yes\n")
    if returncode:
        raise RuntimeError(f"dummy diagnostic exited {returncode}")
    print(LOG)


if __name__ == "__main__":
    main()