#!/usr/bin/env python3
#
# conc_util.py ---  Air Concentration Utilities
#                   See HYSPLIT Tutorial Sec. 8.3
#
# Change history:
# 30 Sep 2025 (BB) - standard-aligned, cross-platform version
# 19 Nov 2025 (SZ) - create image files when DSP=NO
# 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


# Standard Header
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", Path.home() / "tutorial"))
    exe_suffix = ""


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

For Linux users (csh or tcsh shell):
   setenv TTR 'YOUR_HYSPLIT_TUTORIAL_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")


# Input concentration file from previous example
inp = OUT / "hysplit2.bin"
if not inp.exists():
    print(f"ERROR: Missing input file {inp}")
    print("Re-run conc_cexp.py to create output file!")
    sys.exit(1)


# Step 1: conread
if DSP.upper() == "YES":
    input("Press ENTER to continue with conread...")

with open("input.txt", "w", encoding="ascii") as f:
    f.write(f"{inp.name}\n")
    f.write("0\n")
    f.write("/n\n")
    f.write("EOD\n")

print("Running conread < input.txt...")
with open("input.txt", "r", encoding="ascii") as fin:
    result = subprocess.run(
        [str(PGM / "exec" / f"conread{exe_suffix}")],
        stdin=fin,
        capture_output=True,
        text=True
    )
    if result.returncode != 0:
        print(f"ERROR: {result.stderr}")
        sys.exit(1)
    else:
        print(result.stdout)

os.remove("input.txt")


# Step 2: con2asc
if DSP.upper() == "YES":
    input("Press ENTER to continue with con2asc...")

run_cmd([
    str(PGM / "exec" / f"con2asc{exe_suffix}"),
    f"-i{inp.name}", "-s", "-u1.0E+12"
], "Running con2asc")

txt_out = inp.with_suffix(".bin.txt")
if txt_out.exists():
    print(txt_out.read_text())


# Step 3: con2stn
if DSP.upper() == "YES":
    input("Press ENTER to continue with con2stn...")

with open("samplers.txt", "w", encoding="ascii") as f:
    f.write("C510 42.25 -78.80\n")

run_cmd([
    str(PGM / "exec" / f"con2stn{exe_suffix}"),
    f"-i{inp.name}", "-ocon2stn.txt",
    "-r1", "-c1.0E+12", "-ssamplers.txt"
], "Running con2stn")

if Path("con2stn.txt").exists():
    print(Path("con2stn.txt").read_text(encoding="ascii"))


# Step 4: timeplot
if DSP.upper() == "YES":
    input("Press ENTER to continue with timeplot...")

with open("LABELS.CFG", "w", encoding="ascii") as f:
    f.write("'TITLE&','### conc_util.py ### &'\n")

run_cmd([
    str(PGM / "exec" / f"timeplot{exe_suffix}"),
    "+g1", "-icon2stn.txt"
], "Generating timeplot (no obs)")

html = OUT / "timeplot.html"

if html.exists():
    if DSP.upper() == "YES":
        webbrowser.open(html.as_uri())
    else:
        convert_html_to_img(str(html), [("util008.png", 1)])
else:
    print("Warning: timeplot.html not found")


# With observations (CAPTEX M510 file)
if DSP.upper() == "YES":
    input("Press ENTER to continue with measured data...")

obs = TTR / "captex" / "captex2_m510.txt"
if obs.exists():
    run_cmd([
        str(PGM / "exec" / f"timeplot{exe_suffix}"),
        "+g1", "-icon2stn.txt", f"-s{obs}"
    ], "Generating timeplot (with obs)")

    html = OUT / "timeplot.html"
    if html.exists():
        if DSP.upper() == "YES":
            webbrowser.open(html.as_uri())
        else:
            convert_html_to_img(str(html), [("util009.png", 1)])
    else:
        print("Warning: timeplot.html not found")
else:
    print(f"WARNING: Obs file not found: {obs}")