"""Unit tests for Civitai Krea2 dataset scrape / split (no live network)."""

from __future__ import annotations

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

import httpx
import pytest

from gpu_rent.civitai import list_images, list_models
from gpu_rent.civitai_dataset import (
    cmd_split,
    looks_minor,
    normalize_image,
    rating_from_nsfw,
    reaction_score,
    search_row,
    tags_from_image,
    train_row,
)
from gpu_rent.errors import CloudError


def test_rating_from_nsfw():
    assert rating_from_nsfw("None") == "pg"
    assert rating_from_nsfw("Soft") == "pg13"
    assert rating_from_nsfw("Mature") == "r"
    assert rating_from_nsfw("X") == "x"
    assert rating_from_nsfw(True) == "x"
    assert rating_from_nsfw(False) == "pg"


def test_reaction_score():
    assert reaction_score({"likeCount": 2, "heartCount": 3, "laughCount": 1}) == 6
    assert reaction_score(None) == 0


def test_looks_minor():
    assert looks_minor("loli school", [])
    assert looks_minor("portrait", ["shota"])
    assert not looks_minor("adult woman redhead", ["stockings"])


def test_tags_from_prompt_fallback():
    tags = tags_from_image({}, "woman, redhead, cinematic lighting, <lora:x:0.7>", cap=5)
    assert "woman" in tags
    assert "redhead" in tags
    assert not any("lora:" in t.lower() for t in tags)


def test_normalize_image_skips_no_prompt():
    item = {
        "id": 1,
        "type": "image",
        "nsfwLevel": "None",
        "stats": {"likeCount": 10},
        "meta": {},
    }
    assert normalize_image(item, kind="checkpoint", model_id=1, model_version_id=2, ours=False) is None


def test_normalize_image_ok():
    item = {
        "id": 42,
        "type": "image",
        "url": "https://example/x.jpg",
        "username": "u",
        "createdAt": "2026-01-01T00:00:00Z",
        "nsfwLevel": "Soft",
        "width": 832,
        "height": 1216,
        "stats": {"likeCount": 10, "heartCount": 2},
        "meta": {
            "prompt": "a woman in soft light",
            "negativePrompt": "blur",
            "steps": 8,
            "cfgScale": 1,
            "sampler": "euler",
            "civitaiResources": [{"type": "lora", "modelVersionId": 9, "weight": 0.7}],
        },
        "tags": ["woman", "portrait"],
    }
    row = normalize_image(
        item, kind="checkpoint", model_id=100, model_version_id=200, ours=True, min_score=5
    )
    assert row is not None
    assert row["id"] == 42
    assert row["rating"] == "pg13"
    assert row["score"] == 12
    assert row["ours"] is True
    assert row["params"]["steps"] == 8
    assert "woman" in row["tags"]


def test_normalize_skips_low_score():
    item = {
        "id": 3,
        "type": "image",
        "nsfwLevel": "None",
        "stats": {"likeCount": 1},
        "meta": {"prompt": "hello world"},
    }
    assert (
        normalize_image(item, kind="lora", model_id=1, model_version_id=2, ours=False, min_score=5)
        is None
    )


def test_train_and_search_rows():
    row = {
        "id": 7,
        "rating": "r",
        "score": 50,
        "kind": "checkpoint",
        "modelVersionId": 3231611,
        "tags": ["woman", "cinematic"],
        "prompt": "A woman in cinematic light",
        "negativePrompt": "blur",
        "params": {"steps": 8, "cfgScale": 1, "sampler": "euler"},
        "resources": [{"type": "lora", "modelVersionId": 1, "weight": 0.5}],
    }
    tr = train_row(row)
    assert "Tags: woman, cinematic" in tr["instruction"]
    assert "Rating: r" in tr["instruction"]
    assert "A woman" in tr["output"]
    assert "cfg: 1" in tr["output"]
    sr = search_row(row)
    assert sr["id"] == 7
    assert sr["loras"] == [{"versionId": 1, "weight": 0.5}]
    assert sr["params"]["cfg"] == 1


def test_cmd_split_writes_artifacts(tmp_path: Path):
    cat = tmp_path / "datasets" / "civitai" / "catalog"
    cat.mkdir(parents=True)
    rows = [
        {
            "id": 1,
            "kind": "checkpoint",
            "rating": "pg",
            "score": 10,
            "modelVersionId": 1,
            "tags": ["a"],
            "prompt": "p1",
            "negativePrompt": "",
            "params": {"steps": 4},
            "resources": [],
        },
        {
            "id": 2,
            "kind": "lora",
            "rating": "x",
            "score": 20,
            "modelVersionId": 2,
            "tags": ["b"],
            "prompt": "p2",
            "negativePrompt": "n",
            "params": {"cfgScale": 1},
            "resources": [],
        },
    ]
    with (cat / "images.jsonl").open("w", encoding="utf-8") as fh:
        for r in rows:
            fh.write(json.dumps(r) + "\n")
    counts = cmd_split(out_root=tmp_path, log=lambda m: None)
    assert counts["train"] == 2
    assert counts["search"] == 2
    assert counts["kind:checkpoint"] == 1
    assert counts["kind:lora"] == 1
    assert counts["rating:pg"] == 1
    assert counts["rating:x"] == 1
    search = (tmp_path / "datasets" / "civitai" / "search.jsonl").read_text(encoding="utf-8")
    assert '"id": 1' in search
    train = (tmp_path / "datasets" / "civitai" / "train.jsonl").read_text(encoding="utf-8")
    assert "Write a Krea 2 prompt" in train


def _mock_response(payload: dict, status: int = 200) -> httpx.Response:
    return httpx.Response(status, json=payload, request=httpx.Request("GET", "https://civitai.red/api/v1/x"))


def test_list_models_failover(monkeypatch):
    calls = []

    class FakeClient:
        def __init__(self, *a, **k):
            pass

        def __enter__(self):
            return self

        def __exit__(self, *a):
            return False

        def get(self, url, headers=None, params=None):
            calls.append((url, dict(params or {})))
            if "civitai.com" in url:
                return _mock_response({"items": [{"id": 1, "name": "Krea"}], "metadata": {}})
            return _mock_response({}, status=404)

    monkeypatch.setattr("gpu_rent.civitai.httpx.Client", FakeClient)
    host, data = list_models("tok", "civitai.red", types="Checkpoint", query="krea2", limit=10)
    assert host == "civitai.com"
    assert data["items"][0]["id"] == 1
    assert any("civitai.red" in u for u, _ in calls)
    # query search must not send page=
    assert all("page" not in p for _, p in calls)


def test_list_images_cursor(monkeypatch):
    class FakeClient:
        def __init__(self, *a, **k):
            pass

        def __enter__(self):
            return self

        def __exit__(self, *a):
            return False

        def get(self, url, headers=None, params=None):
            assert params.get("modelVersionId") == 3231611
            assert params.get("withMeta") == "true"
            assert "modelId" not in params
            return _mock_response(
                {
                    "items": [{"id": 9, "meta": {"prompt": "x"}}],
                    "metadata": {"nextCursor": "abc"},
                }
            )

    monkeypatch.setattr("gpu_rent.civitai.httpx.Client", FakeClient)
    host, data = list_images(
        "tok", "civitai.red", model_version_id=3231611, with_meta=True, limit=50
    )
    assert host == "civitai.red"
    assert data["metadata"]["nextCursor"] == "abc"


def test_list_models_raises_on_fail(monkeypatch):
    class FakeClient:
        def __init__(self, *a, **k):
            pass

        def __enter__(self):
            return self

        def __exit__(self, *a):
            return False

        def get(self, url, headers=None, params=None):
            return _mock_response({}, status=500)

    monkeypatch.setattr("gpu_rent.civitai.httpx.Client", FakeClient)
    with pytest.raises(CloudError):
        list_models("tok", "civitai.red", types="Checkpoint", query="krea")