"""Unit tests for assistent personas overlay seed helpers."""

from __future__ import annotations

import json
from pathlib import Path
from unittest.mock import MagicMock

from gpu_rent.provision import _safe_persona_id, count_wanted_models_on_vm, seed_assistent_personas


def test_safe_persona_id():
    assert _safe_persona_id("neutral") == "neutral"
    assert _safe_persona_id("lewd_v2") == "lewd_v2"
    assert _safe_persona_id("../x") is None
    assert _safe_persona_id("a/b") is None
    assert _safe_persona_id("") is None


def test_seed_assistent_personas_writes_overlay(tmp_path: Path, monkeypatch):
    root = tmp_path / "assistent-personas"
    (root / "_base").mkdir(parents=True)
    (root / "personas" / "cinema").mkdir(parents=True)
    (root / "_base" / "assistant.json").write_text(
        json.dumps({"default_persona": "cinema"}), encoding="utf-8"
    )
    (root / "personas" / "cinema" / "persona.json").write_text(
        json.dumps({"title": "Кино", "accent": "#7aa2f7"}), encoding="utf-8"
    )
    (root / "personas" / "cinema" / "bio.json").write_text(
        json.dumps({"role": "DP"}), encoding="utf-8"
    )
    (root / "personas" / "cinema" / "humor.json").write_text(
        json.dumps({"frequency": "rare"}), encoding="utf-8"
    )

    puts: dict[str, str] = {}
    ssh_cmds: list[str] = []

    def fake_put(cfg, host, remote, text, mode=0o644):
        puts[remote] = text

    def fake_ssh(cfg, host, cmd, check=False, timeout=60):
        ssh_cmds.append(cmd)
        if ".gpu-rent-gpu.json" in cmd:
            return json.dumps({"vram_mib": 24576, "name": "4090", "compute_cap": "8.9"})
        return ""

    monkeypatch.setattr("gpu_rent.provision.put_text", fake_put)
    monkeypatch.setattr("gpu_rent.provision.run_ssh", fake_ssh)
    monkeypatch.setattr("gpu_rent.paths.assistent_personas_dir", lambda: root)
    monkeypatch.setattr(
        "gpu_rent.paths.assistent_personas_example_dir", lambda: tmp_path / "missing"
    )
    monkeypatch.setattr(
        "gpu_rent.paths.assistent_personas_manifest_path", lambda: tmp_path / "no.yaml"
    )

    cfg = MagicMock()
    cfg.assistent_personas_dir = str(root)
    logs: list[str] = []
    seed_assistent_personas(cfg, "1.2.3.4", logs.append)

    assert any(p.endswith("personas/cinema/persona.json") for p in puts)
    assert any(p.endswith("personas/cinema/bio.json") for p in puts)
    assert any(p.endswith("personas/cinema/humor.json") for p in puts)
    assert not any(p.endswith("extra.md") for p in puts)
    assert "/mnt/swarm_data/Assistent/_base/assistant.json" in puts
    asst = json.loads(puts["/mnt/swarm_data/Assistent/_base/assistant.json"])
    assert asst["default_persona"] == "cinema"
    assert asst["num_ctx"] == 16384
    assert any("rm -f" in c and "personas.json" in c for c in ssh_cmds)
    assert any("overlay personas/1" in m for m in logs)


def test_seed_skips_exact_if_remote_exists(tmp_path: Path, monkeypatch):
    root = tmp_path / "assistent-personas"
    (root / "personas" / "leonid").mkdir(parents=True)
    (root / "personas" / "leonid" / "persona.json").write_text(
        '{"title":"L"}', encoding="utf-8"
    )
    (root / "personas" / "leonid" / "exact.json").write_text(
        '{"controls":{"preference_bias":0.9}}', encoding="utf-8"
    )
    puts: dict[str, str] = {}

    def fake_put(cfg, host, remote, text, mode=0o644):
        puts[remote] = text

    def fake_ssh(cfg, host, cmd, check=False, timeout=60):
        if "exact.json" in cmd and "test -f" in cmd:
            return "yes"
        if ".gpu-rent-gpu.json" in cmd:
            return "{}"
        return ""

    monkeypatch.setattr("gpu_rent.provision.put_text", fake_put)
    monkeypatch.setattr("gpu_rent.provision.run_ssh", fake_ssh)
    monkeypatch.setattr("gpu_rent.paths.assistent_personas_dir", lambda: root)
    monkeypatch.setattr(
        "gpu_rent.paths.assistent_personas_example_dir", lambda: tmp_path / "missing"
    )
    monkeypatch.setattr(
        "gpu_rent.paths.assistent_personas_manifest_path", lambda: tmp_path / "no.yaml"
    )

    cfg = MagicMock()
    cfg.assistent_personas_dir = str(root)
    seed_assistent_personas(cfg, "h", lambda *_: None)
    assert not any(p.endswith("exact.json") for p in puts)
    assert any(p.endswith("persona.json") for p in puts)


def test_count_wanted_models(monkeypatch):
    raw = "lora:\n  - url: https://civitai.red/x?modelVersionId=1\n    title: A\n"

    monkeypatch.setattr(
        "gpu_rent.provision.run_ssh",
        lambda *a, **k: raw,
    )
    assert count_wanted_models_on_vm(MagicMock(), "h") == 1