#!/usr/bin/env python3
"""
═══════════════════════════════════════════════════════════════
  Демоэкзамен 2026 — Автодеплой Модуля 1
  Разворачивает конфигурацию на всех VM через Proxmox Guest Agent
═══════════════════════════════════════════════════════════════

  Порядок: ISP → HQ-RTR → HQ-SRV → BR-RTR → BR-SRV → HQ-CLI

  Для каждой VM:
    1. (если нужно) Добавляем vmbr0 адаптер через Proxmox API
    2. Получаем DHCP на внешнем интерфейсе
    3. curl setup.sh с demo.digit-shop.ru
    4. Запускаем (если ошибка timezone — перезапуск)
    5. Ждём завершения
"""

import requests
import urllib3
import json
import time
import re
import sys
import os
import getpass
import argparse

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

CONFIG_FILE = os.path.expanduser("~/.demo_checker.json")
SETUP_BASE = "https://demo.digit-shop.ru/generated"
GA_TIMEOUT = 120  # setup scripts can take a while
GA_POLL = 2

# VM deployment order and config
DEPLOY_ORDER = [
    {
        "name": "ISP",
        "vmid": 100,
        "slug": "isp",
        "needs_vmbr0": False,   # already has vmbr0 on net0
        "dhcp_cmd": "dhcpcd ens19 2>/dev/null || dhclient ens19 2>/dev/null; sleep 3",
        "wait_after": 5,
    },
    {
        "name": "HQ-RTR",
        "vmid": 101,
        "slug": "hq-rtr",
        "needs_vmbr0": True,    # add temp vmbr0 for internet
        "dhcp_cmd": None,       # will be set dynamically after vmbr0 added
        "wait_after": 5,
    },
    {
        "name": "HQ-SRV",
        "vmid": 103,
        "slug": "hq-srv",
        "needs_vmbr0": True,    # has net1[vmbr0] but may need DHCP
        "dhcp_cmd": None,
        "wait_after": 5,
    },
    {
        "name": "BR-RTR",
        "vmid": 102,
        "slug": "br-rtr",
        "needs_vmbr0": True,
        "dhcp_cmd": None,
        "wait_after": 5,
    },
    {
        "name": "BR-SRV",
        "vmid": 104,
        "slug": "br-srv",
        "needs_vmbr0": True,
        "dhcp_cmd": None,
        "wait_after": 5,
    },
    {
        "name": "HQ-CLI",
        "vmid": 105,
        "slug": "hq-cli",
        "needs_vmbr0": True,
        "dhcp_cmd": None,
        "wait_after": 5,
    },
]


class PVE:
    def __init__(self, host, port, token_id, token_secret, node):
        self.base = f"https://{host}:{port}/api2/json"
        self.node = node
        self.s = requests.Session()
        self.s.headers["Authorization"] = f"PVEAPIToken={token_id}={token_secret}"
        self.s.verify = False

    def get(self, path):
        r = self.s.get(f"{self.base}{path}")
        r.raise_for_status()
        return r.json().get("data")

    def post(self, path, **kwargs):
        r = self.s.post(f"{self.base}{path}", data=kwargs)
        r.raise_for_status()
        return r.json().get("data")

    def put(self, path, **kwargs):
        r = self.s.put(f"{self.base}{path}", data=kwargs)
        r.raise_for_status()
        return r.json().get("data")

    def vm_config(self, vmid):
        return self.get(f"/nodes/{self.node}/qemu/{vmid}/config")

    def agent_ping(self, vmid):
        try:
            self.post(f"/nodes/{self.node}/qemu/{vmid}/agent/ping")
            return True
        except:
            return False

    def exec_cmd(self, vmid, cmd, timeout=GA_TIMEOUT):
        """Execute command via guest agent with stdin method"""
        try:
            r = self.s.post(
                f"{self.base}/nodes/{self.node}/qemu/{vmid}/agent/exec",
                data={"command": "/bin/sh", "input-data": cmd + "\n"}
            )
            r.raise_for_status()
            result = r.json().get("data", {})
            pid = result.get("pid")
            if pid is None:
                return (None, "", "no pid")

            elapsed = 0
            while elapsed < timeout:
                time.sleep(GA_POLL)
                elapsed += GA_POLL
                try:
                    st = self.get(
                        f"/nodes/{self.node}/qemu/{vmid}/agent/exec-status?pid={pid}"
                    )
                    if st.get("exited"):
                        return (st.get("exitcode", -1),
                                st.get("out-data", ""),
                                st.get("err-data", ""))
                except:
                    pass

            return (None, "", "timeout")
        except Exception as e:
            return (None, "", str(e))

    def add_net(self, vmid, net_key, bridge, model="virtio"):
        """Add network adapter to VM"""
        try:
            self.put(
                f"/nodes/{self.node}/qemu/{vmid}/config",
                **{net_key: f"{model},bridge={bridge}"}
            )
            return True
        except Exception as e:
            print(f"    ⚠️  Не удалось добавить {net_key}: {e}")
            return False

    def remove_net(self, vmid, net_key):
        """Remove network adapter from VM"""
        try:
            self.put(
                f"/nodes/{self.node}/qemu/{vmid}/config",
                delete=net_key
            )
            return True
        except:
            return False

    def find_free_net(self, vmid):
        """Find first unused netX slot"""
        cfg = self.vm_config(vmid)
        for i in range(10):
            if f"net{i}" not in cfg:
                return f"net{i}"
        return None

    def has_vmbr0(self, vmid):
        """Check if VM already has vmbr0"""
        cfg = self.vm_config(vmid)
        for key, val in cfg.items():
            if key.startswith("net") and isinstance(val, str) and "vmbr0" in val:
                return True
        return False


def load_config():
    """Load PVE config (same as checker)"""
    if not os.path.exists(CONFIG_FILE):
        return None
    try:
        with open(CONFIG_FILE, "r") as f:
            data = json.load(f)
        if "host" in data:
            return data
        # Multi-PVE format
        items = list(data.values())
        if len(items) == 1:
            return items[0]
        # Select
        print("\n  Сохранённые PVE серверы:")
        for i, cfg in enumerate(items, 1):
            print(f"    {i}. {cfg.get('name', cfg.get('host'))}")
        ch = input(f"  Выбор [1]: ").strip() or "1"
        return items[int(ch)-1] if ch.isdigit() and int(ch) <= len(items) else items[0]
    except:
        return None


def wait_agent(api, vmid, name, timeout=60):
    """Wait for guest agent to become available"""
    print(f"    ⏳ Жду agent на {name}...", end="", flush=True)
    for i in range(timeout // 3):
        if api.agent_ping(vmid):
            print(" ✅")
            return True
        time.sleep(3)
        print(".", end="", flush=True)
    print(" ❌ timeout")
    return False


def find_new_iface(api, vmid):
    """Find the interface connected to vmbr0 by looking for one without IP"""
    ec, out, _ = api.exec_cmd(vmid, "ip -br addr show | grep -v lo | grep -v '@'")
    if not out:
        return None
    # Find interfaces, pick the last one (newest) or one without IP
    lines = [l.strip() for l in out.split("\n") if l.strip()]
    for line in reversed(lines):
        parts = line.split()
        if len(parts) >= 1:
            iface = parts[0]
            # If it has no IP (only name + state), it's likely the new one
            if len(parts) <= 2 or not any("." in p for p in parts[2:]):
                return iface
    # Fallback: return last interface
    if lines:
        return lines[-1].split()[0]
    return None


def deploy_vm(api, vm_cfg, vms_config, keep_vmbr0=False):
    """Deploy setup script on a single VM"""
    name = vm_cfg["name"]
    vmid = vms_config.get(name, vm_cfg["vmid"])
    slug = vm_cfg["slug"]

    print(f"\n{'═'*60}")
    print(f"  🚀 {name} (VMID {vmid})")
    print(f"{'═'*60}")

    # 1. Check agent
    if not wait_agent(api, vmid, name):
        print(f"    ❌ Agent недоступен, пропускаю {name}")
        return False

    # 2. Add vmbr0 if needed
    added_net = None
    new_iface = None
    if vm_cfg["needs_vmbr0"]:
        if api.has_vmbr0(vmid):
            print(f"    ℹ️  vmbr0 уже есть")
        else:
            net_key = api.find_free_net(vmid)
            if net_key:
                # Запоминаем интерфейсы ДО добавления
                ec, before, _ = api.exec_cmd(vmid,
                    "ip -br link show | grep -v lo | awk '{print $1}'")
                ifaces_before = set((before or "").split())

                print(f"    📡 Добавляю {net_key}=vmbr0...")
                if api.add_net(vmid, net_key, "vmbr0"):
                    added_net = net_key
                    print(f"    ✅ {net_key} добавлен, жду появления интерфейса...")

                    # Ждём появления нового интерфейса в ОС (до 30 сек)
                    for attempt in range(15):
                        time.sleep(2)
                        ec, after, _ = api.exec_cmd(vmid,
                            "ip -br link show | grep -v lo | awk '{print $1}'")
                        ifaces_after = set((after or "").split())
                        new_ifaces = ifaces_after - ifaces_before
                        if new_ifaces:
                            new_iface = new_ifaces.pop()
                            print(f"    ✅ Новый интерфейс: {new_iface}")
                            break
                        print(f"    .", end="", flush=True)
                    else:
                        print(f"\n    ⚠️  Интерфейс не появился за 30 сек")
            else:
                print(f"    ⚠️  Нет свободных net-слотов")

    # 3. Get DHCP on the new interface (or all)
    if new_iface:
        print(f"    📡 DHCP на {new_iface}...")
        ec, out, _ = api.exec_cmd(vmid,
            f"ip link set {new_iface} up && "
            f"dhcpcd {new_iface} 2>/dev/null || dhclient {new_iface} 2>/dev/null; "
            f"sleep 3",
            timeout=30)
    elif vm_cfg.get("dhcp_cmd"):
        print(f"    📡 DHCP: {vm_cfg['dhcp_cmd'][:50]}...")
        ec, out, err = api.exec_cmd(vmid, vm_cfg["dhcp_cmd"], timeout=30)
    else:
        # Fallback: try all interfaces
        print(f"    📡 DHCP на всех интерфейсах...")
        ec, out, _ = api.exec_cmd(vmid,
            "for iface in $(ip -br link show | grep -v lo | awk '{print $1}'); do "
            "  ip link set $iface up 2>/dev/null; "
            "  dhcpcd $iface 2>/dev/null || dhclient $iface 2>/dev/null; "
            "done; sleep 3",
            timeout=30)

    # 4. Verify internet
    print(f"    🌐 Проверка интернета...")
    ec, out, _ = api.exec_cmd(vmid, "ping -c1 -W3 77.88.8.8 2>/dev/null && echo INET_OK", timeout=15)
    if "INET_OK" not in (out or ""):
        # Try again with explicit DHCP
        print(f"    ⚠️  Нет интернета, пробую DHCP на всех интерфейсах...")
        ec, out, _ = api.exec_cmd(vmid,
            "for iface in $(ip -br link show | grep -v lo | awk '{print $1}'); do "
            "  ip link set $iface up 2>/dev/null; "
            "  dhcpcd $iface 2>/dev/null || dhclient $iface 2>/dev/null; "
            "done; sleep 5; ping -c1 -W3 77.88.8.8 2>/dev/null && echo INET_OK",
            timeout=30)
        if "INET_OK" not in (out or ""):
            print(f"    ❌ Нет интернета на {name}! Пробую продолжить...")

    # 5. Download setup script
    url = f"{SETUP_BASE}/{slug}/setup.sh"
    print(f"    📥 Скачиваю: {url}")
    ec, out, err = api.exec_cmd(vmid,
        f"cd /root && curl -sSfO {url} && chmod +x setup.sh && echo DOWNLOAD_OK",
        timeout=30)
    if "DOWNLOAD_OK" not in (out or ""):
        print(f"    ❌ Не удалось скачать: {(err or out or '')[:100]}")
        return False
    print(f"    ✅ Скрипт скачан")

    # 6. Run setup script (first attempt)
    print(f"    ▶️  Запуск setup.sh (попытка 1)...")
    ec, out, err = api.exec_cmd(vmid, "cd /root && ./setup.sh 2>&1", timeout=GA_TIMEOUT)
    combined = (out or "") + (err or "")

    # 7. Check for timezone error → re-run
    if "Failed to set time zone" in combined or "No such file" in combined or ec != 0:
        print(f"    ⚠️  Ошибка timezone или другая, перезапуск...")
        time.sleep(3)
        print(f"    ▶️  Запуск setup.sh (попытка 2)...")
        ec, out, err = api.exec_cmd(vmid, "cd /root && ./setup.sh 2>&1", timeout=GA_TIMEOUT)
        combined = (out or "") + (err or "")

        # Third attempt if still failing
        if ec != 0:
            print(f"    ⚠️  Ещё ошибка (ec={ec}), попытка 3...")
            time.sleep(3)
            ec, out, err = api.exec_cmd(vmid, "cd /root && ./setup.sh 2>&1", timeout=GA_TIMEOUT)
            combined = (out or "") + (err or "")

    if ec == 0:
        print(f"    ✅ {name} — setup.sh завершён успешно!")
    else:
        print(f"    ⚠️  {name} — setup.sh ec={ec}")
        # Show last lines of output
        lines = combined.strip().split("\n")
        for line in lines[-5:]:
            clean = re.sub(r'[^\x20-\x7E]', '', line).strip()
            if clean:
                print(f"       {clean[:80]}")

    # 8. Wait a bit for services to settle
    if vm_cfg.get("wait_after"):
        time.sleep(vm_cfg["wait_after"])

    # 9. Убираем временный vmbr0
    if added_net and not keep_vmbr0:
        print(f"    🗑️  Удаляю {added_net} (vmbr0) из Proxmox...")
        if api.remove_net(vmid, added_net):
            print(f"    ✅ {added_net} удалён")
        else:
            print(f"    ⚠️  Не удалось удалить {added_net}")

        # Перезагрузка сети только на HQ-CLI (DHCP клиент, нужно обновить)
        if name == "HQ-CLI":
            print(f"    🔄 Перезагрузка сети на HQ-CLI...")
            api.exec_cmd(vmid,
                "systemctl restart network 2>/dev/null; "
                "systemctl restart networking 2>/dev/null; "
                "systemctl restart NetworkManager 2>/dev/null; "
                "sleep 3",
                timeout=30)
            print(f"    ✅ Сеть перезагружена")

    return True


def main():
    parser = argparse.ArgumentParser(description="Автодеплой Модуля 1")
    parser.add_argument("--vm", help="Только конкретные VM (ISP,HQ-RTR,...)")
    parser.add_argument("--skip", help="Пропустить VM (ISP,HQ-CLI,...)")
    parser.add_argument("--no-vmbr0", action="store_true", help="Не добавлять vmbr0")
    parser.add_argument("--keep-vmbr0", action="store_true", help="Не удалять vmbr0 после деплоя")
    parser.add_argument("--clean-vmbr0", action="store_true", help="Удалить ВСЕ vmbr0 (кроме ISP) без деплоя")
    parser.add_argument("--dry-run", action="store_true", help="Только показать план")
    parser.add_argument("--timeout", type=int, default=120, help="Таймаут exec (сек)")
    args = parser.parse_args()

    global GA_TIMEOUT
    GA_TIMEOUT = args.timeout

    print()
    print("╔══════════════════════════════════════════════════════╗")
    print("║     Демоэкзамен 2026 — Автодеплой Модуля 1         ║")
    print("╚══════════════════════════════════════════════════════╝")

    # Load config
    config = load_config()
    if not config:
        print("  ❌ Нет сохранённого конфига. Сначала запустите checker_v6.py")
        return

    if not config.get("token_secret"):
        config["token_secret"] = getpass.getpass("  Token Secret: ")

    vms_config = config.get("vms", {})
    print(f"\n  🖥️  PVE: {config.get('name', config['host'])} ({config['host']})")

    # Filter VMs
    deploy_list = DEPLOY_ORDER[:]
    if args.vm:
        selected = [v.strip().upper() for v in args.vm.split(",")]
        deploy_list = [v for v in deploy_list if v["name"] in selected]
    if args.skip:
        skipped = [v.strip().upper() for v in args.skip.split(",")]
        deploy_list = [v for v in deploy_list if v["name"] not in skipped]

    if args.no_vmbr0:
        for v in deploy_list:
            v["needs_vmbr0"] = False

    # Show plan
    if not args.clean_vmbr0:
        print(f"\n  📋 План деплоя ({len(deploy_list)} VM):")
        for v in deploy_list:
            vmbr = " +vmbr0" if v["needs_vmbr0"] else ""
            vmid = vms_config.get(v["name"], v["vmid"])
            print(f"    {v['name']:10s} (VMID {vmid}){vmbr}")
            print(f"      → {SETUP_BASE}/{v['slug']}/setup.sh")

    if args.dry_run:
        print("\n  (dry-run, выход)")
        return

    if not args.clean_vmbr0:
        confirm = input(f"\n  Начать деплой? [Y/n]: ").strip().lower()
        if confirm not in ("", "y", "yes", "д", "да"):
            print("  Отмена")
            return

    # Connect
    api = PVE(
        config["host"], config.get("port", 8006),
        config.get("token_id", "root@pam!checker"),
        config["token_secret"], config.get("node", "pve")
    )

    # Mode: clean-vmbr0 only
    if args.clean_vmbr0:
        print(f"\n  🧹 Удаляю vmbr0 со всех VM (кроме ISP)...")
        for vm_cfg in deploy_list:
            name = vm_cfg["name"]
            if name == "ISP":
                continue
            vmid = vms_config.get(name, vm_cfg["vmid"])
            cfg = api.vm_config(vmid)
            removed = False
            for key, val in cfg.items():
                if key.startswith("net") and isinstance(val, str) and "vmbr0" in val:
                    print(f"    {name}: удаляю {key} (vmbr0)...")
                    api.remove_net(vmid, key)
                    removed = True
            if removed:
                # Перезагрузка сети только на HQ-CLI
                if name == "HQ-CLI" and api.agent_ping(vmid):
                    api.exec_cmd(vmid,
                        "systemctl restart network 2>/dev/null; "
                        "systemctl restart networking 2>/dev/null; "
                        "systemctl restart NetworkManager 2>/dev/null",
                        timeout=30)
                    print(f"    {name}: ✅ vmbr0 удалён, сеть перезагружена")
                else:
                    print(f"    {name}: ✅ vmbr0 удалён")
            else:
                print(f"    {name}: нет vmbr0")
        return

    # Deploy each VM
    results = {}
    for vm_cfg in deploy_list:
        try:
            ok = deploy_vm(api, vm_cfg, vms_config, keep_vmbr0=args.keep_vmbr0)
            results[vm_cfg["name"]] = ok
        except Exception as e:
            print(f"    ❌ Ошибка: {e}")
            results[vm_cfg["name"]] = False

    # Summary
    print(f"\n{'═'*60}")
    print(f"  РЕЗУЛЬТАТЫ ДЕПЛОЯ")
    print(f"{'═'*60}")
    for name, ok in results.items():
        ico = "✅" if ok else "❌"
        print(f"  {ico} {name}")

    ok_count = sum(1 for v in results.values() if v)
    total = len(results)
    print(f"\n  Готово: {ok_count}/{total}")

    if ok_count == total:
        print(f"\n  💡 Теперь запустите проверку:")
        print(f"     python3 checker_v6.py --module 1 --all")


if __name__ == "__main__":
    main()