#!/usr/bin/env python3
"""Cross-compile the SMOD imx6 atsc_app on the Ubuntu VM and deploy it to the board.

    build         package app source -> VM (~/atsc_head_build) -> make -> pull imx_atsc_app back
    deploy        push a local imx_atsc_app to 192.168.1.143, md5-verify, then swap it in
    build-deploy  both
    log           tail /var/module_log/imx_atsc_app.log from the device

Two device quirks drive the design here:

  * The imx6 runs an old dropbear. Modern paramiko can't negotiate with it
    ("unknown cipher" / "no acceptable host key"), so this script vendors
    paramiko 2.12 into ~/.cache/smod_atsc_dev and re-execs itself under it.
  * That dropbear has no SFTP subsystem, so files go over an exec channel with
    ``cat > file`` rather than sftp.put().

Deploy writes imx_atsc_app_head first and only overwrites the active imx_atsc_app once
the md5 matches -- a half-transferred binary should never become the thing rc.local runs.

After deploying you still need a cold boot (power_off -> USB_Blast_off -> power_on ->
USB_Blast_on) for rc.local to pick up the new app. This script deliberately does not
touch the relays; that sequencing lives in the vivado-mcp tools.
"""

import argparse
import hashlib
import os
import re
import subprocess
import sys
import tarfile
import tempfile


def _env_lookup(key, default=None):
    """Machine-local values (paths, IPs, credentials) live in ~/skills_env.md, never in git.
    An OS environment variable of the same name overrides the file."""
    v = os.environ.get(key)
    if v:
        return v
    path = os.path.join(os.path.expanduser("~"), "skills_env.md")
    try:
        with open(path, encoding="utf-8") as f:
            for line in f:
                m = re.match(rf"\s*{key}\s*=\s*(\S+)", line)
                if m:
                    return m.group(1)
    except OSError:
        pass
    if default is None:
        sys.exit(f"missing {key}: add it to ~/skills_env.md or set it as an env var")
    return default


REPO = r"E:\Xilinx_design\smod\SMOD_ATSC\my_exciter"
APP_SUBDIR = "imx6/src/atsc_app"

VM_HOST = _env_lookup("CROSS_VM_HOST", "192.168.223.128")
VM_USER = _env_lookup("CROSS_VM_USER", "bowei")
VM_PASS = _env_lookup("CROSS_VM_PASS")
VM_BUILD_DIR = "atsc_head_build"  # relative to the VM home dir
TOOLCHAIN = "/opt/fsl-imx-x11/4.1.15-2.0.0/environment-setup-cortexa7hf-neon-poky-linux-gnueabi"

IMX_HOST = _env_lookup("SMOD_BOARD_IP", "192.168.1.143")
IMX_USER = _env_lookup("SMOD_SSH_USER", "root")
IMX_PASS = _env_lookup("SMOD_SSH_PASS")
IMX_APP_DIR = "/aw/anywave/app"
IMX_STAGING = f"{IMX_APP_DIR}/imx_atsc_app_head"
IMX_ACTIVE = f"{IMX_APP_DIR}/imx_atsc_app"
IMX_LOG = "/var/module_log/imx_atsc_app.log"

BINARY = "imx_atsc_app"
VENDOR_MARKER = "SMOD_PARAMIKO_VENDORED"
# paramiko 2.12 reaches for algorithms.TripleDES, which cryptography >= 43 removed.
VENDOR_PKGS = ["paramiko==2.12.0", "cryptography<43"]


# --------------------------------------------------------------------------- vendoring

def ensure_legacy_paramiko():
    """Re-exec under a vendored paramiko 2.12 unless we're already running under it."""
    if os.environ.get(VENDOR_MARKER):
        return
    cache = os.path.join(os.path.expanduser("~"), ".cache", "smod_atsc_dev", "paramiko212")
    if not os.path.isdir(os.path.join(cache, "paramiko")):
        print(f"[setup] installing legacy paramiko into {cache} (one time)", flush=True)
        os.makedirs(cache, exist_ok=True)
        # --no-warn-conflicts: pip diffs --target installs against the ambient env and
        # reports conflicts that don't apply here, which reads like a failure.
        subprocess.run([sys.executable, "-m", "pip", "install", "--quiet", "--no-warn-conflicts",
                        "--target", cache] + VENDOR_PKGS, check=True)
    env = dict(os.environ, **{VENDOR_MARKER: "1"})
    env["PYTHONPATH"] = cache + os.pathsep + env.get("PYTHONPATH", "")
    env["PYTHONWARNINGS"] = "ignore"
    sys.exit(subprocess.run([sys.executable, os.path.abspath(__file__)] + sys.argv[1:], env=env).returncode)


# ------------------------------------------------------------------------------- ssh

def connect(host, user, password):
    import paramiko
    cli = paramiko.SSHClient()
    cli.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    cli.connect(host, username=user, password=password, timeout=20, look_for_keys=False, allow_agent=False)
    return cli


def run(cli, cmd, check=True, quiet=False):
    if not quiet:
        print(f"  [ssh] {cmd}")
    _, out, err = cli.exec_command(cmd, timeout=1800)
    rc = out.channel.recv_exit_status()
    stdout, stderr = out.read().decode(errors="replace"), err.read().decode(errors="replace")
    if check and rc != 0:
        sys.exit(f"[FAIL] rc={rc}\n{stdout}\n{stderr}")
    return rc, stdout, stderr


def md5_of(path):
    h = hashlib.md5()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


# ----------------------------------------------------------------------------- build

def package_source(repo, source_mode):
    tar_path = os.path.join(tempfile.mkdtemp(prefix="smod_app_"), "app.tar")
    if source_mode == "git":
        subprocess.run(["git", "-C", repo, "archive", "--format=tar", "-o", tar_path, f"HEAD:{APP_SUBDIR}"],
                       check=True)
        rev = subprocess.run(["git", "-C", repo, "rev-parse", "--short", "HEAD"],
                             capture_output=True, text=True, check=True).stdout.strip()
        print(f"[build] packaged {APP_SUBDIR} from git HEAD ({rev})")
    else:
        src = os.path.join(repo, *APP_SUBDIR.split("/"))
        skip = {".o", ".bak", ".zip", ".tar"}
        with tarfile.open(tar_path, "w") as tf:
            for root, _, files in os.walk(src):
                for name in files:
                    if os.path.splitext(name)[1] in skip or name == BINARY:
                        continue
                    full = os.path.join(root, name)
                    tf.add(full, arcname=os.path.relpath(full, src))
        print(f"[build] packaged {APP_SUBDIR} from the working tree (stale .o excluded)")
    return tar_path


def cmd_build(args):
    tar_path = package_source(args.repo, args.source)
    out_path = os.path.abspath(args.out)

    vm = connect(VM_HOST, VM_USER, VM_PASS)
    try:
        sftp = vm.open_sftp()
        print(f"[build] uploading {os.path.getsize(tar_path)/1e6:.1f} MB to VM")
        sftp.put(tar_path, f"/home/{VM_USER}/app_src.tar")

        run(vm, f"rm -rf ~/{VM_BUILD_DIR} && mkdir -p ~/{VM_BUILD_DIR} && "
                f"tar xf ~/app_src.tar -C ~/{VM_BUILD_DIR}")

        print("[build] cross-compiling (this takes a few minutes)")
        rc, out, err = run(vm, f"bash -lc 'source {TOOLCHAIN} && cd ~/{VM_BUILD_DIR} && make {args.make_args}'",
                           check=False)
        tail = "\n".join((out + err).splitlines()[-40:])
        if rc != 0:
            sys.exit(f"[FAIL] make exited {rc}. Last output:\n{tail}")
        print(tail)

        remote_bin = f"/home/{VM_USER}/{VM_BUILD_DIR}/{BINARY}"
        rc, _, _ = run(vm, f"test -f {remote_bin}", check=False)
        if rc != 0:
            sys.exit(f"[FAIL] make succeeded but {remote_bin} is missing")

        os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
        sftp.get(remote_bin, out_path)
    finally:
        vm.close()

    print(f"\n[OK] {out_path}  ({os.path.getsize(out_path)/1e6:.2f} MB, md5 {md5_of(out_path)})")
    return out_path


# ---------------------------------------------------------------------------- deploy

def push_file(cli, local, remote):
    """Send a file over an exec channel: the device's dropbear has no SFTP subsystem."""
    stdin, stdout, _ = cli.exec_command(f"cat > {remote}", timeout=600)
    with open(local, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 16), b""):
            stdin.write(chunk)
    stdin.channel.shutdown_write()
    rc = stdout.channel.recv_exit_status()
    if rc != 0:
        sys.exit(f"[FAIL] `cat > {remote}` exited {rc}")


def cmd_deploy(args):
    local = os.path.abspath(args.binary)
    if not os.path.exists(local):
        sys.exit(f"[FAIL] not found: {local}")
    want = md5_of(local)
    print(f"[deploy] {local} (md5 {want})")

    imx = connect(IMX_HOST, IMX_USER, IMX_PASS)
    try:
        print(f"[deploy] staging -> {IMX_STAGING}")
        push_file(imx, local, IMX_STAGING)
        run(imx, f"chmod 755 {IMX_STAGING}")

        _, out, _ = run(imx, f"md5sum {IMX_STAGING}")
        got = out.split()[0]
        if got != want:
            sys.exit(f"[FAIL] md5 mismatch after transfer: device {got} != local {want}. "
                     f"Active app untouched.")
        print("[deploy] md5 verified")

        run(imx, f"cp -a {IMX_STAGING} {IMX_ACTIVE}")
        _, out, _ = run(imx, f"md5sum {IMX_ACTIVE}")
        if out.split()[0] != want:
            sys.exit("[FAIL] active binary md5 mismatch after swap")
    finally:
        imx.close()

    print(f"\n[OK] deployed to {IMX_ACTIVE}")
    print("Now cold-boot the board so rc.local starts the new app:")
    print("  vivado-mcp: power_off -> USB_Blast_off -> power_on -> (wait 'sys init over') -> USB_Blast_on")
    print(f"Then check the log:  python {os.path.basename(__file__)} log")


# ------------------------------------------------------------------------------- log

def cmd_log(args):
    imx = connect(IMX_HOST, IMX_USER, IMX_PASS)
    try:
        cmd = f"tail -n {args.lines} {IMX_LOG}"
        if args.grep:
            cmd += f" | grep -E '{args.grep}'"
        _, out, _ = run(imx, cmd, check=False, quiet=True)
        print(out)
    finally:
        imx.close()
    print("\nHealthy markers: TALISE_radioOn / 'fb init state ok' / fpga[0xBB]=0x03")


# ------------------------------------------------------------------------------ main

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    def add_build_args(p):
        p.add_argument("--repo", default=REPO)
        p.add_argument("--source", choices=["git", "worktree"], default="git",
                       help="git (default, reproducible) packages HEAD; worktree packages uncommitted edits")
        p.add_argument("--make-args", default="", help="extra args passed to make on the VM")
        p.add_argument("--out", default=os.path.join(tempfile.gettempdir(), BINARY))

    b = sub.add_parser("build", help="cross-compile on the VM and pull the binary back")
    add_build_args(b)

    d = sub.add_parser("deploy", help="push a local binary to the imx6 and swap it in")
    d.add_argument("binary", nargs="?", default=os.path.join(tempfile.gettempdir(), BINARY))

    bd = sub.add_parser("build-deploy", help="build then deploy")
    add_build_args(bd)

    lg = sub.add_parser("log", help="tail the app log on the device")
    lg.add_argument("-n", "--lines", type=int, default=80)
    lg.add_argument("-g", "--grep", help="extended regex to filter with")

    args = ap.parse_args()
    ensure_legacy_paramiko()  # after parsing, so --help never triggers a pip install

    if args.cmd == "build":
        cmd_build(args)
    elif args.cmd == "deploy":
        cmd_deploy(args)
    elif args.cmd == "build-deploy":
        args.binary = cmd_build(args)
        cmd_deploy(args)
    elif args.cmd == "log":
        cmd_log(args)


if __name__ == "__main__":
    main()