"""After OS bootstrap: extensions, autocomplete, Civitai, then start SwarmUI."""

from __future__ import annotations

import json
import shlex
from collections.abc import Callable
from datetime import datetime, timezone
from importlib.resources import files
from pathlib import Path

import httpx

from gpu_rent.civitai import fetch_model_version, pick_primary_file
from gpu_rent.config import Config
from gpu_rent.errors import CloudError, GpuRentError
from gpu_rent.manifests import (
    MODEL_DIRS,
    extract_version_id,
    parse_extensions,
    parse_models,
    remote_root_for,
    repo_dirname,
    repo_matches_runtime,
)
from gpu_rent.ssh_ops import put_text, remote_exists, run_python, run_ssh
from gpu_rent.sync_files import pull_tree, push_tree

Log = Callable[[str], None]
DATA = "/mnt/swarm_data"

# Forwarded to remote install_*.sh (from .env / gpu-rent.vars → os.environ).
_OLLAMA_INSTALL_ENV = ("OLLAMA_VERSION", "OLLAMA_SHA256")


def _remote_llm_env(cfg: Config, *keys: str) -> dict[str, str]:
    import os

    env: dict[str, str] = {"SWARM_USER": cfg.ssh_user}
    for key in keys:
        value = (os.environ.get(key) or "").strip()
        if value:
            env[key] = value
    return env


def _pkg_text(name: str) -> str:
    return files("gpu_rent.remote").joinpath(name).read_text(encoding="utf-8")


def probe_gpu(cfg: Config, host: str, log: Log) -> dict:
    """Write /mnt/swarm_data/.gpu-rent-gpu.json; return parsed dict."""
    out = run_python(
        cfg,
        host,
        _pkg_text("gpu_probe.py"),
        remote_path="/tmp/gpu-rent-gpu_probe.py",
        timeout=60,
        log=log,
    )
    # Last JSON line from script stdout
    data: dict = {}
    for line in reversed(out.splitlines()):
        line = line.strip()
        if line.startswith("{"):
            try:
                data = json.loads(line)
                break
            except json.JSONDecodeError:
                continue
    if data.get("ok"):
        log(
            f"GPU: {data.get('name')} "
            f"{data.get('vram_mib')} MiB cap={data.get('compute_cap')} "
            f"tier={data.get('tier')}"
        )
    else:
        log(f"GPU probe: {data.get('error') or 'нет данных'}")
    return data


def tune_swarm_perf(cfg: Config, host: str, log: Log) -> bool:
    """Install sage/triton into Comfy venv + ExtraArgs. Returns True if SwarmUI restart needed."""
    out = run_python(
        cfg,
        host,
        _pkg_text("tune_swarm_perf.py"),
        remote_path="/tmp/gpu-rent-tune_swarm_perf.py",
        timeout=900,
        log=log,
    )
    return "RESTART_SWARMUI=1" in out


def ensure_swarm_comfy_installed(cfg: Config, host: str, log: Log) -> None:
    """Headless SwarmUI InstallConfirmWS (ComfyUI) when backends are still empty."""
    log("SwarmUI first-install: ComfyUI backend (если ещё empty)…")
    run_python(
        cfg,
        host,
        _pkg_text("install_swarm_comfy.py"),
        remote_path="/tmp/gpu-rent-install_swarm_comfy.py",
        # Cold: git clone Comfy + pip torch — often 15–40 min.
        timeout=3900,
        log=log,
    )


def seed_extensions(cfg: Config, host: str, log: Log, *, update: bool = True) -> bool:
    from gpu_rent.llm_runtime import normalize_runtime

    runtime = normalize_runtime(cfg.llm_runtime)
    all_repos = parse_extensions(cfg.extensions_manifest)
    repos = [r for r in all_repos if repo_matches_runtime(r, runtime)]
    skipped = [r for r in all_repos if r not in repos]
    for repo in skipped:
        log(
            f"extensions: пропуск {repo_dirname(repo)} "
            f"(requires={repo.requires}, LLM_RUNTIME={runtime})"
        )
    if not repos and not update:
        if all_repos:
            log("extensions: все строки отфильтрованы по requires — стоковый SwarmUI")
        else:
            log("extensions.yaml пуст — стоковый SwarmUI")
        return False
    jobs = []
    for repo in repos:
        jobs.append(
            {
                "kind": repo.kind,
                "url": repo.url,
                "ref": repo.ref,
                "dest": remote_root_for(repo),
            }
        )
    put_text(cfg, host, "/tmp/gpu-rent-ext.json", json.dumps(jobs, indent=2))
    put_text(cfg, host, "/tmp/gpu-rent-update-git", "1\n" if update else "0\n")
    if cfg.git_token:
        put_text(cfg, host, "/tmp/gpu-rent-git.token", cfg.git_token + "\n", mode=0o600)
    if update:
        log(f"extensions: clone/update {len(jobs)} из yaml + установленные на data")
    else:
        log(f"extensions: только недостающие из yaml ({len(jobs)}), без git pull (--no-update)")
    out = run_python(
        cfg,
        host,
        _pkg_text("clone_ext.py"),
        remote_path="/tmp/gpu-rent-clone_ext.py",
        timeout=1800,
        log=log,
    )
    return "cloned " in out or "updated " in out


def _github_blob(cfg: Config) -> dict | None:
    url = (
        f"https://api.github.com/repos/{cfg.autocomplete_github_repo}/contents/"
        f"{cfg.autocomplete_github_path}?ref={cfg.autocomplete_github_ref}"
    )
    headers = {"Accept": "application/vnd.github+json", "User-Agent": "gpu-rent"}
    if cfg.git_token:
        headers["Authorization"] = f"Bearer {cfg.git_token}"
    try:
        with httpx.Client(timeout=20.0, follow_redirects=True) as client:
            response = client.get(url, headers=headers)
    except httpx.HTTPError as exc:
        log_skip = str(exc)
        return {"error": log_skip}
    if response.status_code == 403:
        return {"error": "GitHub rate limit — autocomplete не обновляю"}
    if response.status_code != 200:
        return {"error": f"GitHub HTTP {response.status_code}"}
    data = response.json()
    if not isinstance(data, dict):
        return {"error": "неожиданный JSON GitHub"}
    return data


def seed_autocomplete(cfg: Config, host: str, log: Log) -> bool:
    if not cfg.autocomplete_enabled:
        log("autocomplete выключен")
        return False
    dest_dir = f"{DATA}/Data/Autocompletions"
    dest = f"{dest_dir}/{cfg.autocomplete_filename}"
    meta_path = f"{dest}.gpu-rent-meta.json"
    blob = _github_blob(cfg)
    if blob is None:
        return False
    if blob.get("error"):
        log(str(blob["error"]))
        return False
    sha = str(blob.get("sha") or "")
    download_url = str(blob.get("download_url") or "")
    old_sha = ""
    if remote_exists(cfg, host, meta_path):
        raw = run_ssh(cfg, host, f"cat {meta_path}", check=False)
        try:
            old_sha = str(json.loads(raw).get("github_blob_sha") or "")
        except json.JSONDecodeError:
            old_sha = ""
    changed = sha != old_sha or not remote_exists(cfg, host, dest)
    if changed:
        if not download_url:
            log("GitHub не дал download_url")
            return False
        log(f"качаю {cfg.autocomplete_filename}")
        run_ssh(
            cfg,
            host,
            "mkdir -p {dir} && curl -fsSL -o {part} {url} && mv {part} {dest}".format(
                dir=shlex.quote(dest_dir),
                part=shlex.quote(dest + ".partial"),
                url=shlex.quote(download_url),
                dest=shlex.quote(dest),
            ),
            timeout=180,
        )
        meta = {
            "repo": cfg.autocomplete_github_repo,
            "path": cfg.autocomplete_github_path,
            "ref": cfg.autocomplete_github_ref,
            "github_blob_sha": sha,
            "filename": cfg.autocomplete_filename,
            "fetched_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
            "settings_applied": True,
        }
        put_text(cfg, host, meta_path, json.dumps(meta, indent=2) + "\n")
    settings = f"{DATA}/Data/Settings.fds"
    applied = False
    if remote_exists(cfg, host, meta_path):
        raw = run_ssh(cfg, host, f"cat {meta_path}", check=False)
        try:
            applied = bool(json.loads(raw).get("settings_applied"))
        except json.JSONDecodeError:
            applied = False
    if not remote_exists(cfg, host, settings) or not applied:
        fds = (
            "DefaultUser:\n"
            "    AutoComplete:\n"
            f"        Source: {cfg.autocomplete_filename}\n"
            "        EscapeParens: true\n"
        )
        put_text(cfg, host, settings, fds)
        log(f"Settings.fds AutoComplete.Source = {cfg.autocomplete_filename}")
    return changed


def _download_url(host: str, version_id: int, file_info: dict) -> str:
    raw = str(file_info.get("downloadUrl") or "")
    if "civitai." in raw and "/api/download/" in raw:
        # NSFW downloadUrl often points at .com — use the host that had files.
        return f"https://{host}/api/download/models/{version_id}"
    if raw.startswith("https://"):
        return raw
    return f"https://{host}/api/download/models/{version_id}"


def seed_swarmui_api_keys(cfg: Config, host: str, log: Log) -> None:
    """Write CIVITAI_API_TOKEN (and HF if set) into SwarmUI user keys via SetAPIKey.

    Swarm stores them in Users.ldb GenericData — needed for Model Downloader in the UI.
    Call after SwarmUI HTTP is up (after wait_backend / verify).
    """
    keys: dict[str, str] = {}
    if cfg.civitai_api_token:
        keys["civitai_api"] = cfg.civitai_api_token
    hf = (cfg.hf_token or "").strip()
    if hf:
        keys["huggingface_api"] = hf
    if not keys:
        log("SwarmUI API keys: нет CIVITAI_API_TOKEN / HF_TOKEN — skip")
        return
    put_text(
        cfg,
        host,
        "/tmp/gpu-rent-swarm-api-keys.json",
        json.dumps(keys) + "\n",
        mode=0o600,
    )
    names = ", ".join(keys)
    log(f"SwarmUI: прокидываю API keys ({names})")
    try:
        run_python(
            cfg,
            host,
            _pkg_text("swarmui_set_api_keys.py"),
            remote_path="/tmp/gpu-rent-swarmui_set_api_keys.py",
            timeout=180,
            log=log,
        )
    except CloudError as exc:
        log(f"⚠ SwarmUI API keys: {exc}")
        run_ssh(cfg, host, "rm -f /tmp/gpu-rent-swarm-api-keys.json", check=False)


def seed_civitai(cfg: Config, host: str, log: Log) -> None:
    from gpu_rent.huggingface import is_huggingface_url

    entries = parse_models(cfg.models_manifest)
    if not entries:
        log("model-seed пропущен: манифест пуст — дефолт SwarmUI")
        return

    jobs: list[dict] = []
    for entry in entries:
        url = (entry.url or "").strip()
        if url and is_huggingface_url(url):
            name = url.rstrip("/").rsplit("/", 1)[-1].split("?", 1)[0] or "model.safetensors"
            folder = MODEL_DIRS.get(entry.kind, entry.kind)
            dest = f"{DATA}/Models/{folder}/{name}"
            stem = Path(name).stem
            swarm = {
                "name": stem,
                "title": stem,
                "description": f"Hugging Face: {url}",
                "trigger_phrase": "",
                "author": "",
                "tags": ["huggingface"],
            }
            jobs.append(
                {
                    "dest": dest,
                    "url": url,
                    "sha256": "",
                    "auth": "hf",
                    "sidecars": {
                        f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2),
                    },
                }
            )
            continue

        if not cfg.civitai_api_token:
            continue
        vid = entry.version_id or (extract_version_id(entry.url) if entry.url else None)
        if not vid:
            log(f"пропуск {entry.kind}: нет version_id / HF url")
            continue
        try:
            api_host, version = fetch_model_version(
                cfg.civitai_api_token, cfg.civitai_api_host, vid
            )
        except CloudError as exc:
            log(str(exc))
            continue
        info = pick_primary_file(version)
        if not info:
            log(f"version {vid}: нет files[]")
            continue
        name = str(info.get("name") or f"{vid}.safetensors")
        folder = MODEL_DIRS.get(entry.kind, entry.kind)
        dest = f"{DATA}/Models/{folder}/{name}"
        hashes = info.get("hashes") or {}
        sha = str((hashes.get("SHA256") or hashes.get("sha256") or "")).lower()
        stem = Path(name).stem
        civitai_json = json.dumps(version, ensure_ascii=False, indent=2)
        trained = version.get("trainedWords") or []
        phrase = trained[0] if isinstance(trained, list) and trained else ""
        swarm = {
            "name": stem,
            "title": version.get("name") or stem,
            "description": (version.get("description") or "")[:2000],
            "trigger_phrase": phrase,
            "author": (
                ((version.get("model") or {}) if isinstance(version.get("model"), dict) else {}).get(
                    "name"
                )
            ),
            "tags": version.get("tags") or [],
        }
        jobs.append(
            {
                "dest": dest,
                "url": _download_url(api_host, vid, info),
                "sha256": sha,
                "auth": "civitai",
                "sidecars": {
                    f"{stem}.civitai.json": civitai_json,
                    f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2),
                },
            }
        )

    if not jobs:
        if not cfg.civitai_api_token:
            log("Civitai-seed пропущен: нет CIVITAI_API_TOKEN — дефолт SwarmUI")
        else:
            log("Civitai-seed: ни одной скачиваемой строки")
        return

    hf_n = sum(1 for j in jobs if j.get("auth") == "hf")
    civ_n = len(jobs) - hf_n
    if hf_n and not cfg.hf_token:
        log(
            "⚠ в манифесте есть Hugging Face URL, но нет HF_TOKEN — "
            "gated/abliterated файлы дадут 401. Токен: huggingface.co/settings/tokens"
        )
    if civ_n and not cfg.civitai_api_token:
        log("⚠ Civitai-строки пропущены: нет CIVITAI_API_TOKEN")

    if not any(e.kind == "checkpoint" for e in entries):
        log("в манифесте нет checkpoint — генерация может не стартовать")
    put_text(cfg, host, "/tmp/gpu-rent-civitai-jobs.json", json.dumps(jobs, indent=2))
    if cfg.civitai_api_token:
        put_text(
            cfg, host, "/tmp/gpu-rent-civitai.token", cfg.civitai_api_token + "\n", mode=0o600
        )
    if cfg.hf_token:
        put_text(cfg, host, "/tmp/gpu-rent-hf.token", cfg.hf_token + "\n", mode=0o600)
    log(
        f"model-seed: {len(jobs)} файл(ов) "
        f"(civitai={civ_n}, huggingface={hf_n}) — прогресс [N/{len(jobs)}]"
    )
    run_python(
        cfg,
        host,
        _pkg_text("civitai_fetch.py"),
        remote_path="/tmp/gpu-rent-civitai_fetch.py",
        timeout=7200,
        log=log,
    )


def ensure_swarmui_running(cfg: Config, host: str, log: Log, restart: bool) -> None:
    active = run_ssh(cfg, host, "systemctl is-active swarmui 2>/dev/null || true", check=False).strip()
    if restart and active == "active":
        log("systemctl restart swarmui (новые extensions/autocomplete)")
        run_ssh(cfg, host, "sudo -n systemctl restart swarmui", timeout=120)
        return
    if active != "active":
        log("systemctl start swarmui")
        run_ssh(cfg, host, "sudo -n systemctl start swarmui", timeout=120)
    run_ssh(cfg, host, "sudo -n systemctl enable swarmui", check=False)


def provision_llm(cfg: Config, host: str, log: Log) -> None:
    from gpu_rent.llm_runtime import normalize_runtime, parse_ollama_models
    from gpu_rent.ssh_ops import run_script_sudo
    from gpu_rent.state import load_state, save_state

    runtime = normalize_runtime(cfg.llm_runtime)

    def _stop_units(*names: str) -> None:
        for name in names:
            run_ssh(
                cfg,
                host,
                f"sudo -n systemctl stop {name} 2>/dev/null; "
                f"sudo -n systemctl disable {name} 2>/dev/null || true",
                check=False,
            )

    # Drop LLM units that should not hold VRAM for this runtime.
    if runtime == "none":
        log("LLM: none — останавливаю gpu-rent-ollama / gpu-rent-llamacpp если были")
        _stop_units("gpu-rent-ollama", "gpu-rent-llamacpp")
        st = load_state()
        st.notes = dict(st.notes or {})
        st.notes["llm_runtime"] = "none"
        st.notes.pop("llm_error", None)
        save_state(st)
        return
    if runtime == "ollama":
        _stop_units("gpu-rent-llamacpp")
        log("LLM: ставим/запускаем Ollama")
        run_script_sudo(
            cfg,
            host,
            _pkg_text("install_ollama.sh"),
            remote_path="/tmp/gpu-rent-install_ollama.sh",
            timeout=900,
            env=_remote_llm_env(cfg, *_OLLAMA_INSTALL_ENV),
            log=log,
        )
        entries = parse_ollama_models(cfg.ollama_models_manifest)
        defaults = [e.name for e in entries if e.default]
        if defaults:
            log(f"Ollama preferred: {defaults[0]}")
        names = [e.name for e in entries]
        if not names:
            log("ollama-models.yaml пуст — pull skip")
        else:
            put_text(cfg, host, "/tmp/gpu-rent-ollama-models.json", json.dumps(names, indent=2))
            log(f"Ollama: pull {len(names)} из манифеста")
            run_python(
                cfg,
                host,
                _pkg_text("ollama_pull.py"),
                remote_path="/tmp/gpu-rent-ollama_pull.py",
                timeout=7200,
                log=log,
            )
    else:
        raise CloudError(f"неизвестный LLM_RUNTIME={runtime!r}")
    st = load_state()
    st.notes = dict(st.notes or {})
    st.notes["llm_runtime"] = runtime
    st.notes.pop("llm_error", None)
    save_state(st)


def _try_arm_idle_killer(
    cfg: Config,
    host: str,
    log: Log,
    *,
    conn,
    server_id: str | None,
) -> bool:
    """Arm idle-killer; update state notes. Returns True if armed."""
    from gpu_rent.idle_killer import arm_idle_killer
    from gpu_rent.state import load_state, save_state

    if conn is None or not server_id:
        return False
    try:
        arm_idle_killer(cfg, host, conn, server_id, log)
        st = load_state()
        st.notes = dict(st.notes or {})
        st.notes["idle_killer"] = "armed"
        st.notes.pop("idle_killer_error", None)
        save_state(st)
        return True
    except GpuRentError as exc:
        log(f"idle-killer: {exc}")
        st = load_state()
        st.notes = dict(st.notes or {})
        st.notes["idle_killer"] = "failed"
        st.notes["idle_killer_error"] = str(exc)[:500]
        save_state(st)
        log(
            "⚠ idle-killer НЕ вооружён — GPU может тарифицироваться без авто-stop. "
            "Сделай gpu-rent stop или почини identity/application_credential_create."
        )
        return False


def provision_vm(
    cfg: Config,
    host: str,
    log: Log,
    *,
    conn=None,
    server_id: str | None = None,
    update: bool = True,
) -> None:
    from gpu_rent.llm_runtime import normalize_runtime
    from gpu_rent.state import load_state, save_state

    swarm = bool(getattr(cfg, "enable_swarmui", True))
    rt = normalize_runtime(cfg.llm_runtime)
    if not swarm and rt == "none":
        raise CloudError(
            "llm-only: нужен LLM_RUNTIME=ollama (или --ollama / --llm ollama)"
        )

    # Arm ASAP so mid-provision failures still leave auto-stop on the VM.
    armed = _try_arm_idle_killer(cfg, host, log, conn=conn, server_id=server_id)

    restart = bool(update)
    try:
        if swarm:
            try:
                if seed_extensions(cfg, host, log, update=update):
                    restart = True
            except GpuRentError as exc:
                log(f"extensions: {exc}")
                raise
            try:
                if seed_autocomplete(cfg, host, log):
                    restart = True
            except GpuRentError as exc:
                log(f"autocomplete: {exc}")
            seed_civitai(cfg, host, log)
            push_tree(cfg, host, cfg.local_models_dir, f"{DATA}/Models", log, models=True)
            push_tree(
                cfg, host, cfg.local_wildcards_dir, f"{DATA}/Data/Wildcards", log, models=False
            )
            push_tree(
                cfg,
                host,
                cfg.local_workflows_dir,
                f"{DATA}/CustomWorkflows",
                log,
                models=False,
            )
            if cfg.pull_output:
                pull_tree(cfg, host, f"{DATA}/Output", cfg.local_output_dir, log)
        else:
            log("SwarmUI: skip (llm-only)")
            run_ssh(
                cfg,
                host,
                "sudo -n systemctl stop swarmui 2>/dev/null; "
                "sudo -n systemctl disable swarmui 2>/dev/null || true; "
                "echo llm-only | sudo -n tee /mnt/swarm_data/.gpu-rent-llm-only >/dev/null",
                check=False,
            )

        try:
            probe_gpu(cfg, host, log)
        except Exception as exc:
            log(f"GPU probe: {exc}")

        if swarm:
            ensure_swarmui_running(cfg, host, log, restart=restart)
            run_ssh(
                cfg,
                host,
                "sudo -n rm -f /mnt/swarm_data/.gpu-rent-llm-only",
                check=False,
            )

        try:
            provision_llm(cfg, host, log)
        except Exception as exc:
            st = load_state()
            st.notes = dict(st.notes or {})
            st.notes["llm_error"] = str(exc)[:500]
            st.notes["llm_runtime"] = "none"
            save_state(st)
            raise CloudError(f"LLM runtime: {exc}") from exc
    finally:
        # If first arm failed (SSH race), retry once after seeds.
        if not armed:
            _try_arm_idle_killer(cfg, host, log, conn=conn, server_id=server_id)

    if swarm:
        from gpu_rent.access_card import print_access_card
        from gpu_rent.term import console as rich_console

        print_access_card(
            cfg,
            tunneled=False,
            host=host,
            console=rich_console,
            title="gpu-rent · URL после tunnel (ещё ждём Idle)",
        )
    elif rt == "ollama":
        log(f"Ollama API → localhost:{cfg.ollama_local_port} (туннель)")
    log("Hold killer: gpu-rent hold   |  Стоп GPU: gpu-rent stop")