"""Install scripts/comfy-host-agent.mjs on the sidecar image-Comfy PC over SSH.

Reads AIGEN_IMAGE_SSH_HOST, AIGEN_IMAGE_SSH_USER, and AIGEN_IMAGE_SSH_KEY
(default: %USERPROFILE%\\.ssh\\aigen_image_comfy) or AIGEN_IMAGE_SSH_PASS.
Installs the host agent only. Does not start Comfy or load models.
"""
from __future__ import annotations

import os
import sys
import time
from pathlib import Path

import paramiko

ROOT = Path(__file__).resolve().parents[1]
AGENT = ROOT / "scripts" / "comfy-host-agent.mjs"
REMOTE_DIR = r"C:\Users\comfyapi\aigen-agent"
REMOTE_AGENT = REMOTE_DIR + r"\comfy-host-agent.mjs"
REMOTE_START = REMOTE_DIR + r"\start.cmd"
NODE_ZIP = "https://nodejs.org/dist/v22.18.0/node-v22.18.0-win-x64.zip"


def connect():
    host = os.environ.get("AIGEN_IMAGE_SSH_HOST", "192.168.77.101")
    user = os.environ.get("AIGEN_IMAGE_SSH_USER", "comfyapi")
    password = os.environ.get("AIGEN_IMAGE_SSH_PASS")
    key_path = os.environ.get("AIGEN_IMAGE_SSH_KEY") or str(Path.home() / ".ssh" / "aigen_image_comfy")
    if not Path(key_path).exists():
        key_path = os.environ.get("AIGEN_IMAGE_SSH_KEY")
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    kwargs = {
        "hostname": host,
        "username": user,
        "timeout": 15,
        "allow_agent": False,
        "look_for_keys": False,
    }
    if key_path:
        kwargs["key_filename"] = key_path
    elif password:
        kwargs["password"] = password
    else:
        raise SystemExit("Set AIGEN_IMAGE_SSH_PASS or AIGEN_IMAGE_SSH_KEY")
    client.connect(**kwargs)
    return client, host, user


def run(client, cmd, timeout=120):
    stdin, stdout, stderr = client.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode("utf-8", "replace")
    err = stderr.read().decode("utf-8", "replace")
    code = stdout.channel.recv_exit_status()
    return code, out, err


def main():
    if not AGENT.exists():
        raise SystemExit(f"missing {AGENT}")
    token = os.environ.get("IMAGE_COMFY_CONTROL_TOKEN") or os.urandom(24).hex()
    token_file = Path(os.environ.get("TEMP", ".")) / "image_comfy_control_token"
    token_file.write_text(token, encoding="ascii")
    pubkey = Path.home() / ".ssh" / "aigen_image_comfy.pub"

    start_cmd = "\r\n".join([
        "@echo off",
        f"cd /d {REMOTE_DIR}",
        "set COMFY_CONTROL_PORT=8199",
        "set COMFY_PROXY_PORT=8198",
        "set COMFY_HOST=http://127.0.0.1:8188",
        f"set COMFY_CONTROL_TOKEN={token}",
        r"set COMFY_LAUNCH_CWD=C:\ComfyUI-QwenEdit\ComfyUI",
        r"set COMFY_INPUT_DIR=C:\ComfyUI-QwenEdit\ComfyUI\input",
        r"set COMFY_OUTPUT_DIR=C:\ComfyUI-QwenEdit\ComfyUI\output",
        r"set COMFY_LAUNCH_CMD=C:\ComfyUI-QwenEdit\python\python.exe C:\ComfyUI-QwenEdit\ComfyUI\main.py --listen 0.0.0.0 --port 8188 --disable-auto-launch --lowvram --fp8_e4m3fn-unet --disable-cuda-malloc",
        r'"C:\Users\comfyapi\aigen-agent\node.exe" comfy-host-agent.mjs',
        "",
    ])

    client, host, user = connect()
    print("connected", user, host)
    code, out, err = run(client, f'mkdir "{REMOTE_DIR}" 2>nul & echo dir-ok')
    print("mkdir", code, out.strip() or err.strip())

    sftp = client.open_sftp()
    with sftp.file(REMOTE_AGENT, "wb") as handle:
        handle.write(AGENT.read_bytes())
    with sftp.file(REMOTE_START, "w") as handle:
        handle.write(start_cmd)
    if pubkey.exists():
        run(client, r'if not exist "%USERPROFILE%\.ssh" mkdir "%USERPROFILE%\.ssh"')
        existing = ""
        try:
            with sftp.file(r"C:\Users\comfyapi\.ssh\authorized_keys", "r") as handle:
                existing = handle.read().decode("utf-8", "replace")
        except FileNotFoundError:
            existing = ""
        key_text = pubkey.read_text(encoding="ascii").strip()
        if key_text not in existing:
            with sftp.file(r"C:\Users\comfyapi\.ssh\authorized_keys", "a") as handle:
                handle.write(("\n" if existing and not existing.endswith("\n") else "") + key_text + "\n")
            print("authorized_keys updated")
    sftp.close()

    code, out, err = run(client, f'if exist "{REMOTE_DIR}\\node.exe" echo node-present')
    if "node-present" not in out:
        print("downloading node")
        code, out, err = run(
            client,
            f'curl.exe -L --retry 3 -o "{REMOTE_DIR}\\node.zip" {NODE_ZIP}',
            timeout=180,
        )
        print("curl", code, (err or out)[-400:])
        code, out, err = run(
            client,
            f'"C:\\Program Files\\7-Zip\\7z.exe" x -y -o"{REMOTE_DIR}\\node-dist" "{REMOTE_DIR}\\node.zip"',
            timeout=120,
        )
        print("unzip", code, out[-400:] or err[-400:])
        run(client, f'copy /Y "{REMOTE_DIR}\\node-dist\\node-v22.18.0-win-x64\\node.exe" "{REMOTE_DIR}\\node.exe"')

    run(client, r'powershell -NoProfile -Command "Get-Process node -ErrorAction SilentlyContinue | Where-Object { $_.Path -like \'*aigen-agent*\' } | Stop-Process -Force"')
    code, out, err = run(
        client,
        r'cmd /c start "" /MIN C:\Users\comfyapi\aigen-agent\start.cmd',
    )
    print("start", code, out.strip(), err.strip())
    run(
        client,
        r'schtasks /Create /TN AigenImageComfyAgent /TR "C:\Users\comfyapi\aigen-agent\start.cmd" /SC ONLOGON /F',
    )
    client.close()
    print("token-file", str(token_file))
    print("wait for agent")
    time.sleep(2)


if __name__ == "__main__":
    main()