"""up / stop / destroy: GPU lifetime; SwarmUI bootstrap after SSH."""

from __future__ import annotations

from collections.abc import Callable

from gpu_rent.cloud import (
    create_gpu_server,
    delete_floating_ip,
    delete_server,
    ensure_boot_volume,
    ensure_data_volume,
    ensure_floating_ip,
    ensure_keypair,
    ensure_network,
    ensure_security_group,
    guess_operator_cidr,
    pick_existing_server,
    server_status,
    unshelve,
    wait_volume,
)
from gpu_rent.bootstrap import run_bootstrap
from gpu_rent.provision import (
    ensure_swarm_comfy_installed,
    provision_vm,
    seed_swarmui_api_keys,
    tune_swarm_perf,
)
from gpu_rent.ready import verify_gpu_env, verify_stack_on_vm, wait_backend_idle
from gpu_rent.snapshot import ensure_boot_snapshot
from gpu_rent.notify import notify_ready
from gpu_rent.config import Config
from gpu_rent.errors import CloudError, GpuRentError
from gpu_rent.inventory import (
    gpu_quota_from_compute,
    pick_boot_image,
    pick_volume_type,
    resolve_flavor,
)
from gpu_rent.ux import print_up_preview, prompt_server_plan
from gpu_rent.pools import best_offer, format_pool_scan, scan_pools
from gpu_rent.lock import SessionLock
from gpu_rent.os_client import (
    KEYPAIR_NAME,
    SG_NAME,
    compute_quotas,
    connect,
    iter_flavors,
    iter_images,
    iter_volume_types,
)
from gpu_rent.ssh_keys import ensure_ed25519
from gpu_rent.ssh_ops import probe_ssh, run_ssh, wait_ssh
from gpu_rent.state import SessionState, load_state, save_state, utc_now
from gpu_rent.timing import PhaseTimes

Log = Callable[[str], None]


def _log_default(msg: str) -> None:
    if msg.startswith("\r"):
        print("\r" + msg[1:], end="", flush=True)
        return
    print(msg)


def _require_gpu_quota(conn) -> None:
    quota = compute_quotas(conn)
    limit = gpu_quota_from_compute(quota)
    if limit is not None and limit <= 0:
        raise CloudError(
            "квота GPU = 0. Напиши в поддержку Selectel (текст в docs/setup.md). "
            "CLI не создаст сервер."
        )


def _bind_access(
    conn,
    server,
    state: SessionState,
    cfg: Config,
    log: Log,
    *,
    update: bool = True,
    phases: PhaseTimes | None = None,
) -> SessionState:
    clock = phases or PhaseTimes()
    ip, fip_id = ensure_floating_ip(
        conn, server, state.floating_ip_id, state.floating_ip, log
    )
    state.floating_ip = ip
    if fip_id:
        state.floating_ip_id = fip_id
    save_state(state)
    wait_ssh(cfg, ip, log=log)
    clock.mark("SSH", log)
    log(f"SSH {cfg.ssh_user}@{ip}")
    state.phase = "bootstrapping"
    save_state(state)
    swarm = bool(getattr(cfg, "enable_swarmui", True))
    if swarm and update:
        active = run_ssh(cfg, ip, "systemctl is-active swarmui 2>/dev/null || true", check=False).strip()
        if active == "active":
            log("systemctl stop swarmui перед git update")
            run_ssh(cfg, ip, "sudo -n systemctl stop swarmui", timeout=120, check=False)

    # Detect existing markers so light/full choice is explicit (swarm ↔ llm-only).
    probe = run_ssh(
        cfg,
        ip,
        "echo swarm=$(test -f /opt/swarmui/.gpu-rent-bootstrapped && echo yes || echo no); "
        "echo data=$(test -f /mnt/swarm_data/.gpu-rent-ready && echo yes || echo no); "
        "echo llm=$(test -f /mnt/swarm_data/.gpu-rent-llm-only && echo yes || echo no)",
        check=False,
    )
    flags: dict[str, str] = {}
    for line in probe.splitlines():
        if "=" in line:
            k, v = line.strip().split("=", 1)
            flags[k] = v
    has_swarm = flags.get("swarm") == "yes"
    has_data = flags.get("data") == "yes"
    was_llm_only = flags.get("llm") == "yes"

    if swarm:
        light = has_swarm
        if was_llm_only and not has_swarm:
            log("bootstrap: был llm-only → полный проход (ставим SwarmUI)")
        elif light:
            why = (
                "локальный bootstrapped был сброшен, маркер на VM есть"
                if not state.bootstrapped
                else "маркер /opt/swarmui/.gpu-rent-bootstrapped"
            )
            log(f"bootstrap: LIGHT (без apt) — {why}")
        else:
            log("bootstrap: FULL (apt + SwarmUI) — маркера bootstrapped нет")
    else:
        light = has_data
        if has_swarm and not was_llm_only:
            log("bootstrap: llm-only на диске со SwarmUI — LIGHT data, Swarm unit stop")
        if light:
            log("bootstrap: LIGHT llm-only (без apt) — есть .gpu-rent-ready")
        else:
            log("bootstrap: FULL llm-only — маркера data ready нет")

    run_bootstrap(cfg, ip, log, update=update and swarm, light=light)
    clock.mark("bootstrap", log)
    provision_vm(
        cfg,
        ip,
        log,
        conn=conn,
        server_id=getattr(server, "id", None) or state.server_id,
        update=update,
    )
    clock.mark("provision", log)
    if swarm:
        try:
            ensure_swarm_comfy_installed(cfg, ip, log)
        except Exception as exc:
            log(f"SwarmUI Comfy install: {exc}")
            try:
                from gpu_rent.ready import collect_swarm_diagnostics

                collect_swarm_diagnostics(cfg, ip, log)
            except Exception as diag_exc:
                log(f"diag: {diag_exc}")
            raise
        clock.mark("comfy-install", log)
        try:
            wait_backend_idle(cfg, ip, log)
        except CloudError as exc:
            # wait_backend_idle already collected diag on errored/timeout
            log(f"ready: {exc}")
            raise
        clock.mark("Idle", log)
        try:
            if tune_swarm_perf(cfg, ip, log):
                # Stop before remount — otherwise umount fails with target busy.
                run_ssh(
                    cfg,
                    ip,
                    "sudo -n systemctl stop swarmui",
                    check=False,
                    timeout=120,
                )
                run_ssh(
                    cfg,
                    ip,
                    "sudo -n bash -c '"
                    "for p in dlbackend Data Models Output; do "
                    "src=/mnt/swarm_data/$p; dst=/opt/swarmui/$p; "
                    "mkdir -p \"$src\" \"$dst\"; "
                    "cur=$(findmnt -n -o SOURCE --target \"$dst\" 2>/dev/null || true); "
                    "if [[ \"$cur\" != \"$src\" ]]; then "
                    "umount -l \"$dst\" 2>/dev/null || umount \"$dst\" 2>/dev/null || true; "
                    "mount --bind \"$src\" \"$dst\" || true; "
                    "echo remounted $dst; "
                    "fi; "
                    "done'",
                    check=False,
                    timeout=60,
                )
                log("systemctl restart swarmui (perf ExtraArgs)")
                run_ssh(cfg, ip, "sudo -n systemctl restart swarmui", timeout=120)
                try:
                    wait_backend_idle(cfg, ip, log)
                except CloudError as exc:
                    log(f"ready after perf: {exc} — recover backends")
                    try:
                        ensure_swarm_comfy_installed(cfg, ip, log)
                        wait_backend_idle(cfg, ip, log)
                    except Exception as exc2:
                        log(f"recover after perf: {exc2}")
                        raise exc from exc2
            else:
                log("perf tune: restart не нужен")
        except Exception as exc:
            log(f"perf tune: {exc}")
        clock.mark("perf", log)
        try:
            seed_swarmui_api_keys(cfg, ip, log)
        except Exception as exc:
            log(f"SwarmUI API keys: {exc}")
        clock.mark("api-keys", log)
    else:
        log("ready: llm-only (без ожидания SwarmUI Idle)")

    try:
        checks = verify_stack_on_vm(cfg, ip, log, timeout=300.0)
        state.notes = dict(state.notes or {})
        state.notes["stack_vm"] = [
            {"name": c.name, "ok": c.ok, "detail": c.detail} for c in checks
        ]
        state.notes.pop("stack_vm_error", None)
    except CloudError as exc:
        state.notes = dict(state.notes or {})
        state.notes["stack_vm_error"] = str(exc)[:500]
        save_state(state)
        raise
    clock.mark("verify", log)

    try:
        gpu_checks = verify_gpu_env(cfg, ip, log, timeout=600.0)
        state.notes = dict(state.notes or {})
        state.notes["gpu_env"] = [
            {"name": c.name, "ok": c.ok, "detail": c.detail} for c in gpu_checks
        ]
        state.notes.pop("gpu_env_error", None)
    except CloudError as exc:
        state.notes = dict(state.notes or {})
        state.notes["gpu_env_error"] = str(exc)[:500]
        save_state(state)
        raise
    clock.mark("gpu-env", log)

    try:
        ensure_boot_snapshot(
            conn,
            boot_volume_id=state.boot_volume_id,
            cfg=cfg,
            log=log,
        )
    except CloudError as exc:
        log(f"snapshot: {exc}")
    notify_ready(cfg, log)
    try:
        from gpu_rent.balance import (
            BalanceWatchState,
            fetch_balance_rub,
            save_balance_state,
        )

        if cfg.selectel_api_token:
            rub = fetch_balance_rub(cfg.selectel_api_token)
            save_balance_state(
                BalanceWatchState(baseline_rub=rub, last_balance_rub=rub)
            )
            log(
                f"balance: baseline {rub:.0f} ₽ "
                f"(уведомление каждые {cfg.balance_notify_step_rub:.0f} ₽, "
                "нужен local-watchdog)"
            )
        else:
            log("balance: SELECTEL_API_TOKEN не задан — шаги по ₽ skip")
    except Exception as exc:
        log(f"balance baseline: {exc}")
    state.bootstrapped = True
    state.phase = "ready_cloud"
    state.notes = dict(state.notes or {})
    state.notes["enable_swarmui"] = swarm
    state.notes["up_timing"] = clock.summary_line()
    save_state(state)
    for line in clock.summary_lines():
        log(line)
    return state


def adopt_server(cfg: Config, log: Log = _log_default, *, update: bool = True) -> SessionState:
    conn = connect(cfg)
    server = pick_existing_server(conn)
    if not server:
        raise CloudError("нечего adopt: нет сервера с тегом/именем gpu-rent")
    state = load_state()
    state.server_id = server.id
    state.server_name = getattr(server, "name", None)
    state.flavor_id = getattr(server, "flavor_id", None) or (
        (server.flavor or {}).get("id") if isinstance(getattr(server, "flavor", None), dict) else None
    )
    state.phase = "ready_cloud"
    state.availability_zone = cfg.gpu_rent_az
    save_state(state)
    log(f"подхватили {server.id} статус {server_status(server)}")
    if server_status(server) == "ACTIVE":
        state.phase = "bootstrapping"
        save_state(state)
        _bind_access(conn, server, state, cfg, log, update=update)
    return state


def cmd_up(
    cfg: Config,
    *,
    no_spot: bool = False,
    flavor: str | None = None,
    yes: bool = False,
    adopt: bool = False,
    update: bool | None = None,
    confirm: Callable[[str], bool] | None = None,
    ask: Callable[[str, str], str] | None = None,
    log: Log = _log_default,
) -> SessionState:
    do_update = cfg.update_git if update is None else update
    with SessionLock():
        if adopt:
            return adopt_server(cfg, log=log, update=do_update)
        conn = connect(cfg)
        _require_gpu_quota(conn)
        state = load_state()
        existing = pick_existing_server(conn)
        if existing:
            status = server_status(existing)
            state.server_id = existing.id
            state.server_name = getattr(existing, "name", None)
            if status == "ACTIVE":
                # Prefer FIP from state; may still need ensure below.
                fip = state.floating_ip
                # Warm if local state says so, OR VM still has bootstrap marker
                # (local bootstrapped often cleared on stop while disk/VM remain).
                vm_warm = False
                if fip and not (state.bootstrapped and state.floating_ip):
                    try:
                        if probe_ssh(cfg, fip, attempts=1) == "ok":
                            probe = run_ssh(
                                cfg,
                                fip,
                                "test -f /opt/swarmui/.gpu-rent-bootstrapped && echo yes || echo no",
                                check=False,
                                timeout=20,
                            ).strip()
                            vm_warm = probe.endswith("yes") or probe == "yes"
                    except Exception:
                        vm_warm = False

                if (state.bootstrapped and state.floating_ip) or vm_warm:
                    if update is None:
                        do_update = False
                        log(
                            "warm ACTIVE — без git pull SwarmUI/extensions "
                            "(нужен свежий: --update)"
                        )
                    else:
                        do_update = update
                    if vm_warm and not state.bootstrapped:
                        log("warm: маркер на VM — восстанавливаю local bootstrapped")
                        state.bootstrapped = True
                    log("сервер уже ACTIVE — второй GPU не создаём")
                    state.phase = "bootstrapping"
                    save_state(state)
                    _bind_access(conn, existing, state, cfg, log, update=do_update)
                    return state

                # Bootstrap не завершён: почти всегда VM без authorized_keys.
                if not fip:
                    try:
                        fip, fip_id = ensure_floating_ip(
                            conn, existing, state.floating_ip_id, state.floating_ip, log
                        )
                        state.floating_ip = fip
                        if fip_id:
                            state.floating_ip_id = fip_id
                        save_state(state)
                    except Exception as exc:
                        log(f"FIP: {exc}")

                outcome = probe_ssh(cfg, fip, attempts=2) if fip else "down"
                if outcome == "ok":
                    log("сервер ACTIVE, SSH ок — продолжаем bootstrap")
                    state.phase = "bootstrapping"
                    save_state(state)
                    _bind_access(conn, existing, state, cfg, log, update=do_update)
                    return state

                log(
                    f"ACTIVE без bootstrap, SSH={outcome} — "
                    "удаляем compute (диски оставляем), create с user_data Base64"
                )
                delete_server(conn, existing, log)
                state.server_id = None
                state.server_name = None
                state.bootstrapped = False
                state.phase = "idle"
                save_state(state)
                for vid in (state.boot_volume_id, state.data_volume_id):
                    if not vid:
                        continue
                    try:
                        wait_volume(conn, conn.block_storage.get_volume(vid), "available")
                    except Exception as vol_exc:
                        log(f"wait volume {vid}: {vol_exc}")
                existing = None
            if existing is not None and status in {"EXPIRED", "SHELVED", "SHELVED_OFFLOADED"}:
                existing = unshelve(conn, existing, log)
                state.server_id = existing.id
                state.phase = "bootstrapping"
                state.unshelved_at = utc_now()
                save_state(state)
                _bind_access(conn, existing, state, cfg, log, update=do_update)
                return state
            if existing is not None:
                raise CloudError(f"сервер gpu-rent в статусе {status} — разбери в панели")

        if state.server_id:
            log(f"в state был server {state.server_id}, в облаке нет — создаём заново")
            state.server_id = None
            state.server_name = None
            state.floating_ip = None
            state.floating_ip_id = None
            state.bootstrapped = False
            state.phase = "idle"
            save_state(state)

        flavors = list(iter_flavors(conn))
        try:
            picked = resolve_flavor(
                flavors,
                cfg.flavor_preference,
                explicit=flavor,
                default_id=cfg.default_flavor_id or None,
                fallback=cfg.flavor_fallback,
            )
        except ValueError as exc:
            raise CloudError(str(exc)) from exc

        image = pick_boot_image(list(iter_images(conn)))
        if not image:
            raise CloudError("нет GPU-образа Ubuntu Driver 580/535 в Glance")
        vtype = pick_volume_type(list(iter_volume_types(conn)), cfg.gpu_rent_az)
        spot = cfg.default_spot and not no_spot

        try:
            offers = scan_pools(cfg)
            for line in format_pool_scan(offers, cfg.flavor_preference):
                log(line)
            best = best_offer(offers)
            if best and best.region != cfg.os_region_name:
                log(
                    f"! сейчас OS_REGION_NAME={cfg.os_region_name}, "
                    f"а предпочтение лучше закрывается в {best.region} — "
                    f"поставь OS_REGION_NAME={best.region} и GPU_RENT_AZ={best.region}a "
                    f"в .env (диски ещё не созданы) и повтори up"
                )
        except Exception as exc:
            log(f"скан пулов: {exc}")

        print_up_preview(cfg, flavors, picked=picked, spot=spot, log=log)

        if not yes and ask is not None and flavor is None:
            try:
                plan = prompt_server_plan(
                    cfg,
                    flavors,
                    picked=picked,
                    spot=spot,
                    ask=ask,
                    confirm=confirm,
                )
            except ValueError as exc:
                raise GpuRentError(str(exc)) from exc
            picked = plan.flavor
            spot = plan.spot
            if plan.data_gb != cfg.data_volume_size_gb:
                from dataclasses import replace

                cfg = replace(cfg, data_volume_size_gb=plan.data_gb)
            log(
                f"выбрано: {'preemptible ' if spot else ''}{picked.name}, "
                f"data {cfg.data_volume_size_gb} GB"
            )

        prompt = (
            f"Создать {'preemptible ' if spot else ''}GPU {picked.name} "
            f"в {cfg.gpu_rent_az}, образ {getattr(image, 'name', image.id)}, "
            f"data {cfg.data_volume_size_gb} GB.\n"
            f"Диск тарифицируется всегда; idle-killer через {cfg.idle_minutes} мин "
            f"простоя (льгота {cfg.idle_grace_minutes} мин). Продолжить?"
        )
        if not yes:
            ok = confirm(prompt) if confirm else False
            if not ok:
                raise GpuRentError("отменено")

        private_path, pub = ensure_ed25519(cfg.ssh_private_key_path)
        del private_path
        public_key = pub.read_text(encoding="utf-8")
        ensure_keypair(conn, public_key, log)

        net, _subnet = ensure_network(conn, log)
        cidr = guess_operator_cidr()
        if cidr == "0.0.0.0/0":
            import os

            if (os.environ.get("GPU_RENT_SSH_CIDR") or "").strip():
                log("SG SSH: GPU_RENT_SSH_CIDR=0.0.0.0/0")
            else:
                log("не удалось узнать твой IP — SG откроет SSH с 0.0.0.0/0")
        sg = ensure_security_group(conn, cidr, log)

        boot = ensure_boot_volume(
            conn,
            az=cfg.gpu_rent_az,
            volume_type=vtype,
            image_id=image.id,
            snapshot_name=cfg.boot_snapshot_name,
            existing_id=state.boot_volume_id or cfg.boot_volume_id,
            log=log,
        )
        data = ensure_data_volume(
            conn,
            az=cfg.gpu_rent_az,
            volume_type=vtype,
            size_gb=cfg.data_volume_size_gb,
            existing_id=state.data_volume_id or cfg.data_volume_id,
            log=log,
        )

        state.phase = "provisioning"
        state.flavor_id = picked.id
        state.flavor_name = picked.name
        state.boot_volume_id = boot.id
        state.data_volume_id = data.id
        state.image_id = image.id
        state.network_id = net.id
        state.security_group_id = sg.id
        state.availability_zone = cfg.gpu_rent_az
        state.spot = spot
        state.keypair_name = KEYPAIR_NAME
        save_state(state)

        server = create_gpu_server(
            conn,
            flavor_id=picked.id,
            net_id=net.id,
            sg_name=SG_NAME,
            boot_volume_id=boot.id,
            data_volume_id=data.id,
            az=cfg.gpu_rent_az,
            spot=spot,
            public_key=public_key,
            log=log,
        )
        state.server_id = server.id
        state.server_name = getattr(server, "name", None)
        state.created_at = utc_now()
        state.unshelved_at = None
        state.phase = "bootstrapping"
        save_state(state)
        _bind_access(conn, server, state, cfg, log, update=do_update)
        return state


def cmd_stop(
    cfg: Config,
    *,
    destroy_disks: bool = False,
    no_pull: bool = False,
    log: Log = _log_default,
) -> SessionState:
    with SessionLock():
        conn = connect(cfg)
        state = load_state()
        if cfg.pull_output and not no_pull and state.floating_ip:
            try:
                from gpu_rent.sync_files import pull_tree

                pull_tree(cfg, state.floating_ip, "/mnt/swarm_data/Output", cfg.local_output_dir, log)
            except Exception as exc:
                log(f"pull Output не удался: {exc}")
        server = None
        if state.server_id:
            try:
                server = conn.compute.get_server(state.server_id)
            except Exception:
                server = None
        if server is None:
            server = pick_existing_server(conn)
        if server:
            delete_server(conn, server, log)
        else:
            log("compute уже нет")

        try:
            from gpu_rent.idle_killer import revoke_old_credentials

            revoke_old_credentials(conn, log)
        except Exception as exc:
            log(f"revoke idle-killer app cred: {exc}")

        if not cfg.keep_floating_ip:
            delete_floating_ip(conn, state.floating_ip_id, state.floating_ip, log)
            state.floating_ip = None
            state.floating_ip_id = None

        if destroy_disks:
            for vid in (state.data_volume_id, state.boot_volume_id):
                if not vid:
                    continue
                try:
                    conn.block_storage.delete_volume(vid, ignore_missing=True)
                    log(f"удалён том {vid}")
                except Exception as exc:
                    log(f"том {vid} не удалился: {exc}")
            state.boot_volume_id = None
            state.data_volume_id = None

        state.server_id = None
        state.server_name = None
        state.bootstrapped = False
        state.phase = "idle"
        save_state(state)
        try:
            from gpu_rent.local_watchdog import clear_lease

            clear_lease()
        except Exception:
            pass
        try:
            from gpu_rent.balance import clear_balance_state

            clear_balance_state()
        except Exception:
            pass
        log("фаза idle" + ("" if destroy_disks else " (диски на месте)"))
        return state