#!/usr/bin/env python3
#
# disp_shp.py ---  Creating HYSPLIT Shapefiles
#                  See HYSPLIT Tutorial Sec. 10.5
#
# Change history:
# 28 Oct 2025 (BB) - standard-aligned, cross-platform version
# 20 Nov 2025 (SZ) - code refactoring and cleanup
# 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 glob
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"  # not used here; kept for parity
    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 = ""

if not PGM.exists():
    script = Path(__file__).name if "__file__" in globals() else "disp_shp.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 function
# ------------------------------------------------------------------

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 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")
        # trailing separator is important
        f.write(f"'{pgm / 'bdyfiles'}{os.sep}'  directory of files\n")


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 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_ge.py'
        f.write(f"'TITLE&','### {fn} ### &'\n")


# --------------- Main flow ---------------
def main():
    arlmap = str(PGM / "graphics" / "arlmap")
    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 = TTR\captex (keeps you from 900s)
    met_dir = TTR / "captex"
    if not met_dir.exists():
        print(f"WARNING: meteorology directory not found: {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}")

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

    # Reuse the previous run output file
    if Path(inp).exists():
        print("Using previous simulation output file")
    else:
        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")

    # ---------------- concplot (GIS dump) ----------------
    concplot = PGM / "exec" / f"concplot{exe_suffix}"
    if not concplot.exists():
        print(f"ERROR: concplot not found at {concplot}")
        sys.exit(1)

    write_labels()

    # Equivalent to:
    # concplot +g1 -i%inp% -a1 -z80 -j%arlmap% -x1.0E+12 -upg -h41.0:-82.5
    #          -c4 -b1000 -t1000 -v50000+20000+10000+5000+2000+1000
    cmd = [
        str(concplot),
        "+g1",
        f"-i{inp}",
        "-a1",
        "-z100",
        f"-j{arlmap}",
        "-x1.0E+12",
        "-upg",
        "-h41.0:-82.0",
        "-g0:200",
        "-c4",
        "-b1000",
        "-t1000",
        "-v50000+20000+10000+5000+2000+1000",
    ]
    run_cmd(cmd, "Create GIS text artifacts via concplot")

    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}")

    # ---------------- Shapefile build ----------------
    # Batch checks for GIS_01000_html_01.txt
    gis_txt = Path("GIS_01000_html_01.txt")
    if gis_txt.exists():
        # Delete any prior concplot shapefile components (concplot.???)
        for leftover in glob.glob("concplot.???"):
            try:
                Path(leftover).unlink()
            except OSError:
                pass

        ascii2shp = PGM / "exec" / f"ascii2shp{exe_suffix}"
        txt2dbf  = PGM / "exec" / f"txt2dbf{exe_suffix}"

        if not ascii2shp.exists():
            print(f"ERROR: ascii2shp not found at {ascii2shp}")
            sys.exit(1)
        if not txt2dbf.exists():
            print(f"ERROR: txt2dbf not found at {txt2dbf}")
            sys.exit(1)

        # ascii2shp -d concplot polygons < GIS_01000_html_01.txt
        # We’ll pass the file via stdin like the batch’s redirection.
        with open(gis_txt, "r", encoding="ascii") as fin:
            print("Running ascii2shp (polygons -> concplot.shp)...")
            r = subprocess.run([str(ascii2shp), "-d", "concplot", "polygons"],
                               stdin=fin, capture_output=True, text=True)
            if r.returncode != 0:
                msg = r.stderr.strip() or r.stdout.strip() or "ascii2shp failed"
                print(f"ERROR: {msg}")
                sys.exit(1)

        # txt2dbf -C11 -C5 -C9 -C5 -C6../ -C6 -d,  GIS_01000_html_01.att concplot.dbf
        # Reproduce the exact arguments/spacing from the batch (including the quirky -C6../ token).
        txt_args = [
            str(txt2dbf),
            "-C11", "-C5", "-C9", "-C5", "-C6../", "-C6", "-d,",
            "GIS_01000_html_01.att", "concplot.dbf",
        ]
        run_cmd(txt_args, "Build DBF attributes via txt2dbf")

        print("Shapefile export complete: concplot.shp / .shx / .dbf (and related) created.")
    else:
        print("ERROR: Failed to create shapefiles ... GIS_01000_html_01.txt is missing!")


if __name__ == "__main__":
    main()