import numpy as np
import pytest

from backend.app.models import ImageRecord, TagKind
from backend.app.models.tag import image_tag
from backend.app.services.ml.centroids import CentroidService
from backend.app.services.tag_service import TagService

pytestmark = pytest.mark.integration


def _img(sha: str, embedding: list[float] | None) -> ImageRecord:
    return ImageRecord(
        path=f"/images/{sha}.jpg",
        sha256=sha,
        size_bytes=1,
        mime="image/jpeg",
        width=1,
        height=1,
        origin="imported_filesystem",
        integrity_status="unknown",
        siglip_embedding=embedding,
    )


async def _attach(db, image_id: int, tag_id: int):
    await db.execute(
        image_tag.insert().values(
            image_record_id=image_id, tag_id=tag_id, source="manual"
        )
    )


@pytest.mark.asyncio
async def test_recompute_skips_too_few_members(db):
    tags = TagService(db)
    tag = await tags.find_or_create("Lonely", TagKind.character)
    img = _img("a" * 64, [0.1] * 1152)
    db.add(img)
    await db.flush()
    await _attach(db, img.id, tag.id)

    svc = CentroidService(db)
    assert await svc.recompute_for_tag(tag.id) is False


@pytest.mark.asyncio
async def test_recompute_writes_centroid(db):
    tags = TagService(db)
    tag = await tags.find_or_create("Popular", TagKind.character)
    for i in range(5):
        img = _img(f"{i:064d}", [float(i)] * 1152)
        db.add(img)
        await db.flush()
        await _attach(db, img.id, tag.id)

    svc = CentroidService(db)
    assert await svc.recompute_for_tag(tag.id) is True

    from backend.app.models import TagReferenceEmbedding
    cen = await db.get(TagReferenceEmbedding, tag.id)
    assert cen is not None
    assert cen.reference_count == 5
    assert abs(np.array(cen.embedding)[0] - 2.0) < 1e-4


@pytest.mark.asyncio
async def test_recompute_skips_ineligible_kind(db):
    tags = TagService(db)
    tag = await tags.find_or_create("somearchive", TagKind.archive)
    for i in range(5):
        img = _img(f"arch{i:060d}", [1.0] * 1152)
        db.add(img)
        await db.flush()
        await _attach(db, img.id, tag.id)
    svc = CentroidService(db)
    assert await svc.recompute_for_tag(tag.id) is False


@pytest.mark.asyncio
async def test_list_drifted_includes_uncomputed(db):
    tags = TagService(db)
    tag = await tags.find_or_create("Drifty", TagKind.character)
    for i in range(5):
        img = _img(f"d{i:063d}", [0.5] * 1152)
        db.add(img)
        await db.flush()
        await _attach(db, img.id, tag.id)
    svc = CentroidService(db)
    drifted = await svc.list_drifted()
    assert tag.id in drifted


@pytest.mark.asyncio
async def test_find_similar_tags(db):
    tags = TagService(db)
    tag = await tags.find_or_create("SimTag", TagKind.character)
    for i in range(5):
        img = _img(f"s{i:063d}", [1.0] * 1152)
        db.add(img)
        await db.flush()
        await _attach(db, img.id, tag.id)
    svc = CentroidService(db)
    await svc.recompute_for_tag(tag.id)

    query_img = _img("q" * 64, [1.0] * 1152)
    db.add(query_img)
    await db.flush()
    hits = await svc.find_similar_tags(query_img.id, limit=10)
    assert any(h.tag_id == tag.id for h in hits)
    sim = next(h.similarity for h in hits if h.tag_id == tag.id)
    assert sim > 0.99