#!/usr/bin/env python3
#
# disp_ge.py --- Creating KML/KMZ for Google Earth
#                See HYSPLIT Tutorial Sec. 10.4
#
# Change history:
# 01 Nov 2025 (BB) - standard-aligned, cross-platform version
# 12 Nov 2025 (BB) - fix: avoid Windows system32 convert.exe; prefer 'magick convert ...'
# 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 webbrowser
from shutil import copy
import zipfile

# ---------------- 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 = ""

if not PGM.exists():
    script = Path(__file__).name if "__file__" in globals() else "disp_ge.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 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 sep required
        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")
        # trailing sep critical
        f.write(f"{str(met_dir)}{os.sep}\n")
        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")
        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")


def find_imagemagick():
    # Locate ImageMagick for cross-platform conversion
    convert_exe = None
    v7_or_later = False
    if system == "windows":
        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
                    if cmd == 'magick':
                        v7_or_later = True
                    break
    else:
        convert_exe = "convert"

    return convert_exe, v7_or_later


# ---------------- Main flow ----------------
def main():
    arlmap = str(PGM / "graphics" / "arlmap")  # -j background
    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, DAT = captex2_wrf27uw.bin
    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 (exactly like batch)
    write_control(syr, smo, sda, shr, lat, lon, lvl, run_hours, top,
                  met_dir, dat_name, inp)
    write_setup()

    # Fresh run 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")

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

    clean(Path("concplot.html"))
    write_labels()

    cmd = [
        str(concplot),
        f"-i{inp}",
        "-a3",
        "+g1",
        "-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, "Creating KML 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}")

    # -------------- KMZ packaging --------------
    kml = Path("HYSPLIT_html.kml")
    if kml.exists():
        kmz = Path("HYSPLITconc.kmz")
        if kmz.exists():
            kmz.unlink()

        # Copy assets beside KML (like batch)
        for src, dst in [
            (PGM / "guicode" / "logocon.gif", Path("logocon.gif")),
            (PGM / "guicode" / "noaa_google.gif", Path("noaa_google.gif")),
            (PGM / "graphics" / "icon63.png", Path("icon63.png")),
        ]:
            if src.exists() and not dst.exists():
                copy(str(src), str(dst))

        files = [
            f for f in [
                "HYSPLIT_html.kml", "icon63.png", "logocon.gif", 
                "noaa_google.gif"
            ]
            if Path(f).exists()
        ]
        with zipfile.ZipFile("HYSPLITconc.kmz", "w", compression=zipfile.ZIP_DEFLATED) as zf:
            for f in files:
                zf.write(f, Path(f).name)
        print("Created KMZ.")

        try:
            kml.unlink()
        except OSError:
            pass

    # -------------- GELABEL overlays --------------
    gelabel_txt = Path("GELABEL_html.txt")
    if gelabel_txt.exists():
        gelabel = PGM / "exec" / f"gelabel{exe_suffix}"
        run_cmd([str(gelabel), "+g1", "-phtml"], "Running gelabel (+g1 -phtml)")

        kmz = Path("HYSPLITconc.kmz")
        splitsvg = PGM / "exec" / f"splitsvg{exe_suffix}"

        im_prog, v7_or_later = find_imagemagick()
        if not im_prog:
            print("NOTE: ImageMagick not found. Skipping GIF label conversion.")
        else:
            for num in range(1, 100):
                tag = f"{num:02d}"  # 01..99
                html_in = Path(f"GELABEL_{tag}_html.html")
                if not html_in.exists():
                    break

                # HTML -> SVG
                run_cmd([str(splitsvg), f"-i{html_in}", f"-oGELABEL_{tag}_html.svg"],
                         f"Split SVG for label {tag}")

                # First frame name usually F01-*
                svg_frame = Path(f"F01-GELABEL_{tag}_html.svg")

                gif_out = Path(f"GELABEL_{tag}_html.gif")

                if v7_or_later:
                    im_cmd = [
                        im_prog,
                        "-background", "white",
                        str(svg_frame),
                        "-flatten",
                        "+adjoin",
                        "+page",
                        "-trim",
                        "+repage",
                        str(gif_out),
                    ]
                else:
                    # Build ImageMagick command exactly like the batch:
                    # magick convert -background white -flatten +adjoin -trim +page
                    #                +repage input.svg output.gif
                    im_cmd = [
                        im_prog,
                        "-background", "white",
                        "-flatten",
                        "+adjoin",
                        "-trim",
                        "+page",
                        "+repage",
                        str(svg_frame),
                        str(gif_out),
                    ]
                run_cmd(im_cmd, f"Convert SVG->GIF for label {tag}")

                # Zip/update into KMZ
                with zipfile.ZipFile("HYSPLITconc.kmz", "a",
                                     compression=zipfile.ZIP_DEFLATED) as zf:
                    zf.write(str(gif_out), gif_out.name)

                # Clean per batch
                try:
                    gif_out.unlink()
                    html_in.unlink()
                except OSError:
                    pass

    if os.path.exists("HYSPLITconc.kmz"):
        print("Created KMZ file: HYSPLITconc.kmz")
    else:
        print("Failed to create: HYSPLITconc.kmz")


if __name__ == "__main__":
    main()