#!/usr/bin/env python3
#
# disp_lab.py --- Enhanced Graphics Labels
#                 See HYSPLIT Tutorial Sec. 10.3
#
# Change history:
# 27 Oct 2025 (BB) - standard-aligned, cross-platform version
# 12 Nov 2025 (BB) - guard missing shapefiles: only write existing; fallback to ARLMAP if none
# 18 Jun 2026 (SZ) - standardize code for run_cmd(), convert_html_to_img(),
#                    etc. across Python scripts
#                  - ensure encoding is specified when opening files

import os
import sys
import platform
from pathlib import Path
import subprocess
import webbrowser

# -------------- OS / args --------------
if sys.version_info < (3, 8):
    print("Please use Python 3.8 or higher.")
    sys.exit(1)

system = platform.system().lower()
if system == "windows":
    PGM = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("C:/hysplit")
    OUT = Path(sys.argv[2]) if len(sys.argv) > 2 else PGM / "working"
    DSP = sys.argv[3] if len(sys.argv) > 3 else "YES"
    TTR = Path(os.environ.get("TTR", "C:/Tutorial"))
    exe_suffix = ".exe"
else:
    PGM = Path(sys.argv[1]) if len(sys.argv) > 1 else Path.home() / "hysplit"
    OUT = Path(sys.argv[2]) if len(sys.argv) > 2 else PGM / "working"
    DSP = sys.argv[3] if len(sys.argv) > 3 else "YES"
    TTR = Path(os.environ.get("TTR", str(Path.home() / "tutorial")))
    exe_suffix = ""

DSP = str(DSP).upper()

if not PGM.exists():
    script = Path(__file__).name if "__file__" in globals() else "disp_lab.py"
    print(f"ERROR: HYSPLIT not found at {PGM}")
    print(f"""
Please specify the path to your HYSPLIT installation using a command-line argument:

python {script} YOUR_HYSPLIT_DIRECTORY [OUT] [DSP (YES|NO)]

Example:
    python {script} C:\\hysplit C:\\hysplit\\working YES
""")
    sys.exit(1)

if not TTR.exists():
    print(f"ERROR: HYSPLIT tutorial not found at {TTR}")
    print("""
Please set the TTR environment variable to point to your HYSPLIT tutorial directory.

For Windows users (Command Prompt):
   SET TTR=YOUR_HYSPLIT_TUTORIAL_DIRECTORY

For Linux users (bash shell):
   export TTR='YOUR_HYSPLIT_TOTURIAL_DIRECTORY'

For Linux users (csh or tcsh shell):
   setenv TTR 'YOUR_HYSPLIT_TOTURIAL_DIRECTORY'""")
    sys.exit(1)

os.chdir(OUT)

# ------------------------------------------------------------------
# Helper functions
# ------------------------------------------------------------------

def run_cmd(cmd, desc="", allow_fail=False, stdout_file=None, append=False):
    if desc:
        print(f"{desc}...")

    exitcode = 0
    with subprocess.Popen(cmd,
                         stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT,
                         text=True) as p:
        if stdout_file:
            mode="a" if append else "w"
            with open(stdout_file, mode, encoding="utf-8") as f_out:
                for ln in p.stdout:
                    print(ln, end="")
                    f_out.write(ln)
        else:
            with p.stdout as f:
                for ln in iter(f.readline, b'\n'):
                    if len(ln) == 0:
                        break
                    print(ln.rstrip())

        done = False
        while not done:
            try:
                exitcode = p.wait(timeout=1)
                done = True
            except subprocess.TimeoutExpired:
                print("WARN: subprocess timed out.")

    if exitcode != 0:
        if allow_fail:
            print(f"WARN: command failed with exit code {exitcode} -> {' '.join(cmd)}")
        else:
            print(f"ERROR: command failed with exit code {exitcode} -> {' '.join(cmd)}")
            sys.exit(1)


def convert_html_to_img(html_file, img_name_frame_nums):
    # Split SVG into Linux-consistent naming: F00-sec71_001.svg, etc.
    temp_img_name = "temp"
    run_cmd([str(PGM / "exec" / f"splitsvg{exe_suffix}"),
             f"-i{html_file}", f"-o{temp_img_name}.svg"])

    # Locate ImageMagick for cross-platform conversion
    if system == "windows":
        convert_exe = None
        imgk_dir = next(
            (k for k in os.environ.get('PATH', '').split(os.pathsep)
             if "ImageMagick" in k),
            None
        )
        if imgk_dir:
            # Use magick.exe for ImageMagick 7.x and later. Use convert.exe for earlier versions.
            for cmd in ('magick', 'convert'):
                candidate = os.path.join(imgk_dir, f'{cmd}{exe_suffix}')
                if os.path.exists(candidate):
                    convert_exe = candidate
                    break
    else:
        convert_exe = "convert"

    if convert_exe:
        # Match Linux naming for DSP=NO mode exactly
        for img_name, frame_num in img_name_frame_nums:
            svg = Path(f"F{frame_num:02d}-{temp_img_name}.svg")
            run_cmd([convert_exe, str(svg), img_name])
            print(f"Converted {svg} -> {img_name}")
            svg.unlink()
    else:
        print("ImageMagick not found — skipping .svg conversion")


def clean(path: Path):
    path.unlink(missing_ok=True)


def write_ascdata(pgm: Path):
    clean(Path("ASCDATA.CFG"))
    with open("ASCDATA.CFG", "w", encoding="ascii") as f:
        f.write("-90.0   -180.0  lat/lon of lower left corner\n")
        f.write("1.0     1.0     lat/lon spacing in degrees\n")
        f.write("180     360     lat/lon number of data points\n")
        f.write("2               default land use category\n")
        f.write("0.2             default roughness length (m)\n")
        f.write(f"'{pgm /'bdyfiles'}{os.sep}'  directory of files\n")  # trailing sep matters


def write_control(
    syr, smo, sda, shr,
    lat, lon, lvl,
    run_hours, top,
    met_dir: Path, dat_name: str,
    inp_name: str
):
    with open("CONTROL", "w", encoding="ascii") as f:
        f.write(f"{syr:02d} {smo:02d} {sda:02d} {shr:02d}\n")
        f.write("1\n")
        f.write(f"{lat:.2f} {lon:.2f} {lvl:.1f}\n")
        f.write(f"{run_hours}\n")
        f.write("0\n")
        f.write(f"{top:.1f}\n")
        f.write("1\n")
        f.write(f"{str(met_dir)}{os.sep}\n")  # trailing sep critical
        f.write(f"{dat_name}\n")
        f.write("1\n")
        f.write("PMCH\n")
        f.write("67000.0\n")
        f.write("3.0\n")
        f.write("00 00 00 00 00\n")
        f.write("1\n")
        f.write("42.0 -78.0\n")
        f.write("0.05 0.05\n")
        f.write("15.0 25.0\n")
        f.write(f".{os.sep}\n")
        f.write(f"{inp_name}\n")
        f.write("2\n")
        f.write("800 1000\n")
        f.write("83 09 26 03 00\n")
        f.write("83 09 28 15 00\n")
        f.write("00 01 00\n")
        f.write("1\n")
        f.write("0.0 0.0 0.0\n")
        f.write("0.0 0.0 0.0 0.0 0.0\n")
        f.write("0.0 0.0 0.0\n")
        f.write("0.0\n")
        f.write("0.0\n")


def write_setup():
    clean(Path("SETUP.CFG"))
    with open("SETUP.CFG", "w", encoding="ascii") as f:
        f.write("&SETUP\n")
        # Order matches the provided batch (numpar, maxpar, initd, kmsl)
        f.write("numpar = 5000,\n")
        f.write("maxpar = 100000,\n")
        f.write("initd = 3,\n")
        f.write("kmsl = 1,\n")
        f.write("/\n")


def write_labels():
    with open("LABELS.CFG", "w", encoding="ascii") as f:
        fn = Path(__file__).name if '__file__' in globals() else 'disp_lab.py'
        f.write(f"'TITLE&','### {fn} ### &'\n")


def run_concplot(map_arg: str, inp_name: str, qfile: Path, label: str,
                 img_name_frame_nums=None):
    clean(Path("concplot.html"))
    write_labels()
    concplot = PGM / "exec" / f"concplot{exe_suffix}"
    if not concplot.exists():
        print(f"ERROR: concplot not found at {concplot}")
        sys.exit(1)
    cmd = [
        str(concplot),
        "+g1",
        f"-i{inp_name}",
        "-z100",
        f"-j{map_arg}",
        "-x1.0E+12",
        "-upg",
        "-h41.0:-82.0",
        "-g0:200",
        "-c4",
        "-v50000+20000+10000+5000+2000+1000+500",
        "-b1000",
        "-t1000",
        f"-q{qfile}",
    ]
    run_cmd(cmd, f"Create concplot ({label})")
    html = Path("concplot.html")
    if html.exists():
        if DSP == "YES":
            try:
                webbrowser.open(html.resolve().as_uri())
            except Exception as e:
                print(f"WARNING: could not open concplot.html in browser: {e}")
        else:
            convert_html_to_img(str(html), img_name_frame_nums)

# -------------- Main flow --------------
def main():
    inp = "hysplit2.bin"

    # ASCDATA
    write_ascdata(PGM)

    # CONTROL params (from batch)
    syr, smo, sda, shr = 83, 9, 25, 17
    lat, lon, lvl = 39.90, -84.22, 10.0
    run_hours = 11
    top = 10000.0

    # MET defaults to TTR\captex
    met_dir = TTR / "captex"
    if not met_dir.exists():
        print(f"WARNING: meteorology directory not found at TTR: {met_dir}")
    dat_name = "captex2_wrf27uw.bin"
    if not (met_dir / dat_name).exists():
        print(f"WARNING: met file not found: {met_dir / dat_name}")

    # Write CONTROL & SETUP
    write_control(syr, smo, sda, shr, lat, lon, lvl, run_hours, top,
                  met_dir, dat_name, inp)
    write_setup()

    # Fresh run (delete inp then hycs_std)
    clean(Path(inp))
    hycs = PGM / "exec" / f"hycs_std{exe_suffix}"
    if not hycs.exists():
        print(f"ERROR: hycs_std not found at {hycs}")
        sys.exit(1)
    run_cmd([str(hycs)], "Running hycs_std")

    # Build shapefiles list from TTR\maps, but only include existing ones.
    maps_dir = TTR / "maps"
    entries = [
        (maps_dir / "roads.shp",  "0 0.005 0.8 0.7 0.5"),
        (maps_dir / "county.shp", "0 0.010 0.4 0.6 0.8"),
        (maps_dir / "states.shp", "0 0.015 0.2 0.2 0.5"),
    ]
    existing = []
    for shp, style in entries:
        if shp.exists():
            existing.append((shp, style))
        else:
            print(f"WARNING: missing shapefile: {shp}")

    # data_case.txt for -q
    qfile = met_dir / "data_case.txt" \
            if (met_dir / "data_case.txt").exists() \
            else (TTR / "captex" / "data_case.txt")
    if not qfile.exists():
        print(f"WARNING: data_case.txt not found in {met_dir} "
              f"or {TTR / 'captex'}; concplot will likely miss obs.")

    # Write MAPTEXT.CFG (exact lines/order from batch)
    with open("MAPTEXT.CFG", "w", encoding="ascii") as f:
        f.write("--------------------------------------------\n")
        f.write("Hysplit Dispersion Calculation\n")
        f.write("--------------------------------------------\n")
        f.write("Source Location: Dayton, Ohio\n")
        f.write("Start Month/Day: 25 September 1983\n")
        f.write("Start Time(UTC): 1700\n")
        f.write("Meteorology Data Source: NARR\n")
        f.write("Trajectory Computation Heights:\n")
        f.write("Pollutant Emission Rate: 201 Kg over 3 hours\n")
        f.write("Initial Pollutant Distribution: 10 m release\n")
        f.write("Deposition Options Enabled: NONE\n")
        f.write("Notes:\n")
        f.write("--------------------------------------------\n")
        f.write("Issued:\n")
        f.write("--------------------------------------------\n")

    # Decide map background: shapefiles.txt if any exist; otherwise ARLMAP.
    if existing:
        with open("shapefiles.txt", "w", encoding="ascii") as f:
            for shp, style in existing:
                f.write(f"'{shp}' {style}\n")
        map_arg = "shapefiles.txt"
        label = "shapefiles_with_maptext"
    else:
        # Fallback avoids concplot error 899
        map_arg = str(PGM / "graphics" / "arlmap")
        label = "arlmap_with_maptext"
        print("INFO: No shapefiles present; falling back to ARLMAP.")

    # concplot with chosen background and MAPTEXT overlay
    run_concplot(map_arg, inp, qfile, label,
                 img_name_frame_nums=[("lab003.png", 1)])

    # Clean up MAPTEXT.CFG as in the batch
    clean(Path("MAPTEXT.CFG"))

if __name__ == "__main__":
    main()