#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
2520 交互式安装向导
=========================
首次使用2520时的引导工具，全中文交互。
功能:
  1. 配置Jetson连接信息 (IP, 用户名, 密码)
  2. 测试SSH连接
  3. 部署jetson_agent.py到Jetson设备
  4. 在Jetson上安装依赖
  5. 启动PC端和Jetson端服务
  6. 打开浏览器
"""

import os
import sys
import json
import time
import getpass
import subprocess
import threading
import webbrowser

# ==================== 常量 ====================

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
CONFIG_FILE = os.path.join(SCRIPT_DIR, "server", "config.json")

DEPLOY_FILES = [
    "jetson_agent.py",
    "jetson_fuzz_modules.py",
    "install_jetson_deps.sh",
]

REMOTE_DIR = "/home/nvidia/2520Fuzz"

PC_PORT = 7885
AGENT_PORT = 7886

# ==================== 工具函数 ====================


def banner():
    """显示启动横幅"""
    print()
    print("=" * 56)
    print("     2520  -  交互式安装向导")
    print("     2520 Interactive Setup Wizard")
    print("=" * 56)
    print()
    print("  本向导将引导您完成以下步骤:")
    print("    [1] 配置 Jetson 连接信息")
    print("    [2] 测试 SSH 连接")
    print("    [3] 部署 Agent 到 Jetson")
    print("    [4] 安装 Jetson 依赖")
    print("    [5] 启动服务并打开浏览器")
    print()


def step_header(step_num, text):
    """显示步骤标题"""
    print()
    print("=" * 52)
    print(f"  步骤 {step_num} : {text}")
    print("=" * 52)


def progress(text, status="running"):
    """显示带状态标志的进度信息"""
    symbols = {
        "running": " >",
        "success": " OK",
        "fail":    " FAIL",
        "skip":    " SKIP",
        "info":    " *",
    }
    sym = symbols.get(status, "  ")
    print(f" {sym} {text}")


def prompt(text, default=None):
    """带默认值的交互式输入"""
    if default:
        msg = f"  {text} [{default}]: "
    else:
        msg = f"  {text}: "
    value = input(msg).strip()
    if not value and default:
        return default
    return value


def prompt_yes_no(text, default="Y"):
    """是/否选择器"""
    default = default.upper()
    opts = {"Y": "[Y/n]", "N": "[y/N]"}
    msg = f"  {text} {opts.get(default, '[Y/n]')}: "
    value = input(msg).strip().lower()
    if not value:
        return default == "Y"
    return value in ("y", "yes", "是")


def ensure_paramiko():
    """确保paramiko可用，否则尝试安装"""
    try:
        import paramiko
        return True
    except ImportError:
        progress("paramiko 未安装，正在安装...", "info")
        ret = os.system(f"{sys.executable} -m pip install paramiko -q")
        if ret != 0:
            progress("pip 安装失败，尝试使用 pip3...", "info")
            ret = os.system(f"pip3 install paramiko -q")
        if ret == 0:
            try:
                import paramiko
                progress("paramiko 安装成功", "success")
                return True
            except ImportError:
                pass
        progress("paramiko 安装失败！", "fail")
        print()
        print("  请手动安装:")
        print("    pip install paramiko")
        print()
        return False


def load_config():
    """加载现有配置"""
    if os.path.exists(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            pass
    return {}


def save_config(cfg):
    """保存配置到config.json"""
    existing = load_config()
    existing.update(cfg)
    with open(CONFIG_FILE, "w", encoding="utf-8") as f:
        json.dump(existing, f, ensure_ascii=False, indent=2)
    progress(f"配置已保存到 {CONFIG_FILE}", "success")


def test_ssh(host, username, password):
    """测试SSH连接，返回 (成功, 错误信息)"""
    try:
        import paramiko
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(host, username=username, password=password, timeout=10)
        stdin, stdout, stderr = ssh.exec_command("uname -a")
        result = stdout.read().decode().strip()
        ssh.close()
        return True, result
    except Exception as e:
        return False, str(e)


def ssh_exec(ssh, command, timeout=30):
    """在SSH连接上执行命令并返回输出"""
    stdin, stdout, stderr = ssh.exec_command(command, timeout=timeout)
    out = stdout.read().decode().strip()
    err = stderr.read().decode().strip()
    return out, err


def create_remote_dirs(ssh):
    """在Jetson上创建远程目录结构"""
    dirs = [
        REMOTE_DIR,
        f"{REMOTE_DIR}/fuzz_workspace",
        f"{REMOTE_DIR}/crash_records",
    ]
    for d in dirs:
        ssh.exec_command(f"mkdir -p {d}")
    progress("远程目录结构已创建", "success")


def upload_file(ssh, local_name):
    """上传单个文件到Jetson"""
    import paramiko
    local_path = os.path.join(SCRIPT_DIR, local_name)
    if not os.path.exists(local_path):
        local_path = os.path.join(SCRIPT_DIR, "jetson", local_name)
    if not os.path.exists(local_path):
        progress(f"本地文件不存在: {local_name}", "skip")
        return False
    remote_path = f"{REMOTE_DIR}/{local_name}"
    try:
        sftp = ssh.open_sftp()
        sftp.put(local_path, remote_path)
        sftp.close()
        ssh.exec_command(f"chmod +x {remote_path}")
        progress(f"已上传: {local_name}", "success")
        return True
    except Exception as e:
        progress(f"上传失败 {local_name}: {e}", "fail")
        return False


def upload_files(ssh):
    """上传所有部署文件到Jetson"""
    progress("开始上传部署文件...", "info")
    all_ok = True
    for fname in DEPLOY_FILES:
        if not upload_file(ssh, fname):
            all_ok = False
    return all_ok


def run_remote_script(ssh, script_name):
    """在Jetson上远程运行脚本"""
    remote_path = f"{REMOTE_DIR}/{script_name}"
    cmd = f"cd {REMOTE_DIR} && bash {script_name}"
    progress(f"在 Jetson 上执行: {script_name}", "info")
    out, err = ssh_exec(ssh, cmd, timeout=120)
    for line in out.split("\n"):
        if line.strip():
            print(f"     {line}")
    return out, err


def start_jetson_agent(ssh):
    """在Jetson上启动Agent服务"""
    progress("正在启动 Jetson Agent...", "info")

    ssh.exec_command(f"pkill -f jetson_agent.py 2>/dev/null")
    time.sleep(1)

    ssh.exec_command(
        f"cd {REMOTE_DIR} && nohup python3 jetson_agent.py > agent.log 2>&1 &"
    )
    time.sleep(2)

    out, err = ssh_exec(ssh, "curl -s http://localhost:7886/api/status")
    if out:
        progress("Jetson Agent 已成功启动!", "success")
        progress(f"   Agent 地址: http://{host}:7886/", "info")
        progress(f"   查看日志: ssh {username}@{host} 'tail -f {REMOTE_DIR}/agent.log'", "info")
        return True
    else:
        time.sleep(3)
        out, err = ssh_exec(ssh, "curl -s http://localhost:7886/api/status")
        if out:
            progress("Jetson Agent 已成功启动!", "success")
            return True
        progress("Agent 可能未正常启动，请稍后检查日志", "fail")
        progress(f"   运行: ssh {username}@{host} 'tail -f {REMOTE_DIR}/agent.log'", "info")
        return False


def start_pc_server():
    """启动PC端服务"""
    server_script = os.path.join(SCRIPT_DIR, "server", "Run_2520.py")
    if not os.path.exists(server_script):
        progress(f"找不到 {server_script}", "fail")
        return None

    progress("正在启动 PC 端服务...", "info")

    startupinfo = None
    creationflags = 0
    if sys.platform == "win32":
        creationflags = subprocess.CREATE_NO_WINDOW

    try:
        proc = subprocess.Popen(
            [sys.executable, server_script],
            cwd=SCRIPT_DIR,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            creationflags=creationflags,
        )
        time.sleep(2)
        if proc.poll() is None:
            progress(f"PC 端服务已启动 (PID: {proc.pid})", "success")
            progress(f"   Web 界面: http://127.0.0.1:{PC_PORT}/", "info")
            return proc
        else:
            progress("PC 端服务启动后立即退出，请检查错误", "fail")
            return None
    except Exception as e:
        progress(f"启动 PC 端服务失败: {e}", "fail")
        return None


def open_browser():
    """打开浏览器"""
    url = f"http://127.0.0.1:{PC_PORT}/"
    progress(f"正在打开浏览器: {url}", "info")
    try:
        webbrowser.open(url)
        progress("浏览器已打开", "success")
    except Exception as e:
        progress(f"自动打开浏览器失败: {e}", "fail")
        progress(f"请手动访问: {url}", "info")


def show_summary(host, username, agent_deployed, pc_started):
    """显示安装总结"""
    print()
    print("=" * 56)
    print("         安装向导完成")
    print("=" * 56)
    print()
    print("  [*] Jetson 连接信息:")
    print(f"      IP 地址:   {host}")
    print(f"      用户名:    {username}")
    print(f"      SSH 端口:  22")
    print()
    print("  [*] 服务状态:")

    if agent_deployed:
        print(f"      Jetson Agent: 已部署")
        print(f"      地址: http://{host}:{AGENT_PORT}/")
    else:
        print(f"      Jetson Agent: 已跳过")

    if pc_started:
        print(f"      PC 服务:      运行中")
        print(f"      Web 界面: http://127.0.0.1:{PC_PORT}/")
    else:
        print(f"      PC 服务:      未启动")

    print()
    print("  [*] 常用命令:")
    print(f"      启动 PC 服务:   cd 2520Fuzz && python server\\Run_2520.py")
    print(f"      查看 Agent 日志: ssh {username}@{host} 'tail -f {REMOTE_DIR}/agent.log'")
    print(f"      停止 Agent:     ssh {username}@{host} 'pkill -f jetson_agent.py'")
    print(f"      启动向导:       python setup_wizard.py")
    print()

    if pc_started:
        print("  [提示]: 浏览器已自动打开，如果未弹出请手动访问:")
        print(f"    http://127.0.0.1:{PC_PORT}/")
    else:
        print("  [提示]: 稍后您可以手动启动 PC 服务:")
        print(f"    cd 2520Fuzz")
        print(f"    python server\\Run_2520.py")
        print(f"    然后访问 http://127.0.0.1:{PC_PORT}/")

    print()
    print("=" * 56)
    print("      感谢使用 2520!")
    print("=" * 56)
    print()


# ==================== 主流程 ====================


def main():
    global host, username

    banner()

    if not ensure_paramiko():
        progress("无法继续，缺少 paramiko 模块", "fail")
        sys.exit(1)

    import paramiko

    existing_cfg = load_config()
    default_host = existing_cfg.get("jetson_host", "192.168.1.108")
    default_user = existing_cfg.get("jetson_user", "nvidia")
    default_pass = "nvidia"

    step_header("1", "配置 Jetson 连接信息")
    print()
    print("  说明: 请输入 Jetson 设备的连接信息")
    print("  如果您不确定，请保留默认值直接回车")
    print()

    host = prompt("Jetson IP 地址", default_host)
    username = prompt("SSH 用户名", default_user)
    password = getpass.getpass(f"  SSH 密码 [******]: ")
    if not password:
        password = default_pass

    print()
    progress(f"连接信息: {username}@{host}", "info")

    step_header("2", "测试 SSH 连接")
    print()

    progress(f"正在连接 {username}@{host} ...", "running")
    success, result = test_ssh(host, username, password)

    if success:
        progress("SSH 连接成功!", "success")
        print(f"    {result}")
    else:
        progress("SSH 连接失败!", "fail")
        print(f"    {result}")
        print()
        print("  请检查:")
        print("    1. Jetson 设备已开机并联网")
        print(f"    2. IP 地址是否正确: {host}")
        print("    3. 用户名和密码是否正确")
        print("    4. SSH 服务是否已启用 (sudo systemctl enable ssh)")
        print()
        retry = prompt_yes_no("是否重新输入连接信息?", "Y")
        if retry:
            return main()
        else:
            progress("跳过连接测试，将手动部署", "skip")
            save_config({
                "jetson_host": host,
                "jetson_user": username,
            })
            print()
            print("  手动部署步骤:")
            print(f"    ssh {username}@{host}")
            print(f"    mkdir -p {REMOTE_DIR}/fuzz_workspace")
            print(f"    mkdir -p {REMOTE_DIR}/crash_records")
            print(f"    然后将 {SCRIPT_DIR}\\jetson\\ 下的文件通过 scp 上传到 {REMOTE_DIR}/")
            print(f"    最后运行: cd {REMOTE_DIR} && python3 jetson_agent.py &")
            print()
            input("  按回车键继续...")
            agent_deployed = False
            step_header("5", "启动 PC 端服务")
            print()
            start_pc = prompt_yes_no("是否现在启动 PC 端服务?", "Y")
            pc_started = False
            if start_pc:
                proc = start_pc_server()
                if proc:
                    pc_started = True
                    open_browser()
            show_summary(host, username, agent_deployed, pc_started)
            return

    save_config({
        "jetson_host": host,
        "jetson_user": username,
    })

    step_header("3", "部署 Agent 到 Jetson")
    print()

    deploy_now = prompt_yes_no("是否现在部署 Agent 到 Jetson?", "Y")
    agent_deployed = False

    if deploy_now:
        progress("建立 SSH 连接...", "running")
        try:
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(host, username=username, password=password, timeout=10)
            progress("SSH 连接已建立", "success")
        except Exception as e:
            progress(f"SSH 连接失败: {e}", "fail")
            deploy_now = False

    if deploy_now:
        progress("创建远程目录...", "running")
        create_remote_dirs(ssh)
        progress("上传部署文件...", "info")
        upload_ok = upload_files(ssh)
        if not upload_ok:
            progress("部分文件上传失败，请检查后重试", "fail")
            deploy_now = False

    if deploy_now:
        step_header("4", "在 Jetson 上安装依赖")
        print()
        install_deps = prompt_yes_no("是否现在在 Jetson 上安装依赖?", "Y")

        if install_deps:
            progress("开始安装依赖（可能需要几分钟）...", "info")
            run_remote_script(ssh, "install_jetson_deps.sh")
        else:
            progress("跳过依赖安装", "skip")
            print("  稍后可手动安装:")
            print(f"    ssh {username}@{host} 'cd {REMOTE_DIR} && bash install_jetson_deps.sh'")
            print()

        progress("部署 Agent 服务...", "info")
        agent_deployed = start_jetson_agent(ssh)
        ssh.close()
    else:
        progress("跳过 Agent 部署", "skip")
        print()
        print("  您可以稍后手动部署:")
        print(f"    cd {SCRIPT_DIR}")
        print(f"    python deploy\\deploy_windows.py")
        print()

    step_header("5", "启动 PC 端服务")
    print()

    start_pc = prompt_yes_no("是否现在启动 PC 端服务?", "Y")
    pc_started = False

    if start_pc:
        proc = start_pc_server()
        if proc:
            pc_started = True
            time.sleep(1)
            open_browser()

    show_summary(host, username, agent_deployed, pc_started)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print()
        print("\n  用户中断")
        print("  安装向导已取消")
        sys.exit(0)
    except Exception as e:
        print(f"\n  发生未预期的错误: {e}")
        sys.exit(1)
