#!/home/hapt/miniconda3/bin/python3
"""Generate legacy Dr_SFC web maps for one completed operational cycle."""

from __future__ import annotations

import json
import os
import fcntl
import subprocess
import sys
import tempfile
import time
from pathlib import Path

PROJECT = Path("/home/tanpv/Progs/CLI/Dr")
PLOT_ROOT = Path("/home/tanpv/Progs/CLI/dro")
OUTPUT_ROOT = PLOT_ROOT / "PDF"
ALLOWED_DAYS = {"01", "06", "11", "16", "21", "26"}
EXPECTED = {
    "Pr_forecast.png", "T2_forecast.png",
    "Dr_ens_1month.png", "Dr_ens_3month.png", "Dr_ens_6month.png",
    *(f"Dr_Prob_{level}_{scale}month.png"
      for level in ("mild", "mord", "seve", "extr") for scale in (1, 3, 6)),
}


def write_status(path: Path, **payload: object) -> None:
    payload["updated_at"] = time.time()
    with tempfile.NamedTemporaryFile("w", dir=path.parent, delete=False) as temp:
        json.dump(payload, temp)
        temp.write("\n")
        name = temp.name
    os.replace(name, path)


def main() -> int:
    if len(sys.argv) != 2 or len(sys.argv[1]) != 8 or not sys.argv[1].isdigit():
        return 2
    tag = sys.argv[1]
    year, month, day = tag[:4], tag[4:6], tag[6:]
    if day not in ALLOWED_DAYS:
        return 2

    source = PROJECT / "outputs" / "Ope" / year / month / day
    target = OUTPUT_ROOT / tag
    status_path = target / ".generation.json"
    log_path = target / "generation.log"
    execution_lock = PLOT_ROOT / ".dr_sfc_generate_worker.lock"
    target.mkdir(parents=True, exist_ok=True)
    required = (
        source / "SUCCESS",
        source / "Ope_Pr_T2_and_Ano.nc",
        source / "Dr_Prob.nc",
    )
    if not all(path.is_file() for path in required):
        display_date = f"{day}-{month}-{year}"
        write_status(status_path, state="unavailable", date=tag, error=f"Ngày {display_date} chưa có số liệu.")
        return 3

    environment = {
        "HOME": "/home/tanpv",
        "PATH": "/home/hapt/miniconda3/bin:/usr/local/bin:/usr/bin:/bin",
        "LANG": "C.UTF-8",
        "LC_ALL": "C.UTF-8",
    }
    try:
        with execution_lock.open("w") as lock_file:
            fcntl.flock(lock_file, fcntl.LOCK_EX)
            write_status(status_path, state="running", date=tag, pid=os.getpid())
            with log_path.open("a", encoding="utf-8") as log:
                log.write(f"{time.strftime('%FT%T%z')} generating {tag}\n")
                for script in ("p6.sh", "p7.sh", "p8.sh"):
                    subprocess.run(
                        ["/bin/bash", str(PLOT_ROOT / script), year, month, day],
                        cwd=PLOT_ROOT,
                        env=environment,
                        stdout=log,
                        stderr=subprocess.STDOUT,
                        check=True,
                        timeout=900,
                    )
        missing = sorted(name for name in EXPECTED if not (target / name).is_file())
        if missing:
            write_status(status_path, state="failed", date=tag, error="missing generated files", missing=missing)
            return 4
        write_status(status_path, state="complete", date=tag, files=len(EXPECTED))
        return 0
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
        write_status(status_path, state="failed", date=tag, error=str(exc))
        return 1
    except Exception as exc:
        write_status(status_path, state="failed", date=tag, error=str(exc))
        return 1


if __name__ == "__main__":
    raise SystemExit(main())