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

from __future__ import annotations

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

import pytest

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):
    yaml_path = tmp_path / "assistent-personas.yaml"
    yaml_path.write_text(
        "default: cinema\n"
        "personas:\n"
        "  - id: cinema\n"
        "    title: Кино\n"
        "    prompt: |\n"
        "      You are a DP.\n"
        "  - id: neutral\n"
        "    title: Нейтральный\n"
        "    prompt: Calm.\n",
        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_manifest_path", lambda: yaml_path
    )
    monkeypatch.setattr(
        "gpu_rent.paths.assistent_personas_example_path", lambda: tmp_path / "missing"
    )

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

    assert any("personas/cinema/persona.json" in p for p in puts)
    assert any("personas/cinema/extra.md" in p 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 not any(p.endswith("personas.yaml") for p in puts)
    assert not any(p.endswith("personas.json") for p in puts)
    assert any("overlay personas/2" in m for m in logs)


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