"""Incremental head retraining (#1317 phase 2). The refit-decision + fingerprint
are split out sklearn-free (train_head itself needs scikit-learn), so they're
tested directly via db_sync."""
import pytest
from sqlalchemy import select

from backend.app.models import (
    ImageRecord,
    MLSettings,
    Tag,
    TagHead,
    TagKind,
    TagSuggestionRejection,
)
from backend.app.models.tag import image_tag
from backend.app.services.ml.heads import (
    _head_fingerprints,
    _heads_needing_retrain,
)

pytestmark = pytest.mark.integration


def _img(db, sha: str) -> ImageRecord:
    img = ImageRecord(
        path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg",
        width=1, height=1, origin="imported_filesystem",
        integrity_status="unknown",
    )
    db.add(img)
    db.flush()
    return img


def _tag(db, name: str) -> Tag:
    t = Tag(name=name, kind=TagKind.general)
    db.add(t)
    db.flush()
    return t


def _apply(db, image_id: int, tag_id: int) -> None:
    db.execute(image_tag.insert().values(
        image_record_id=image_id, tag_id=tag_id, source="manual",
    ))


def _reject(db, image_id: int, tag_id: int) -> None:
    db.add(TagSuggestionRejection(image_record_id=image_id, tag_id=tag_id))
    db.flush()


def _version(db) -> str:
    return db.execute(
        select(MLSettings.embedder_model_version).where(MLSettings.id == 1)
    ).scalar_one()


def _head(db, tag_id: int, fp: str | None, version: str) -> None:
    db.add(TagHead(
        tag_id=tag_id, embedding_version=version,
        weights=[0.0] * 1152, bias=0.0, suggest_threshold=0.5,
        auto_apply_threshold=None, n_pos=10, n_neg=30,
        ap=0.8, precision_cv=0.9, recall=0.6, train_fingerprint=fp,
    ))
    db.flush()


def test_fingerprint_changes_on_new_positive(db_sync):
    tag = _tag(db_sync, "glasses")
    i1 = _img(db_sync, "a" * 64)
    _apply(db_sync, i1.id, tag.id)
    db_sync.commit()
    fp1 = _head_fingerprints(db_sync, [tag.id])[tag.id]

    i2 = _img(db_sync, "b" * 64)
    _apply(db_sync, i2.id, tag.id)
    db_sync.commit()
    assert _head_fingerprints(db_sync, [tag.id])[tag.id] != fp1


def test_fingerprint_changes_on_new_rejection(db_sync):
    tag = _tag(db_sync, "glasses")
    i1 = _img(db_sync, "c" * 64)
    _apply(db_sync, i1.id, tag.id)
    db_sync.commit()
    fp1 = _head_fingerprints(db_sync, [tag.id])[tag.id]

    _reject(db_sync, i1.id, tag.id)
    db_sync.commit()
    assert _head_fingerprints(db_sync, [tag.id])[tag.id] != fp1


def test_needing_retrain_selects_only_changed(db_sync):
    ver = _version(db_sync)
    a = _tag(db_sync, "A")
    _apply(db_sync, _img(db_sync, "d" * 64).id, a.id)
    b = _tag(db_sync, "B")
    _apply(db_sync, _img(db_sync, "e" * 64).id, b.id)
    c = _tag(db_sync, "C")
    _apply(db_sync, _img(db_sync, "f" * 64).id, c.id)
    db_sync.commit()

    ids = [a.id, b.id, c.id]
    fps = _head_fingerprints(db_sync, ids)
    _head(db_sync, a.id, fps[a.id], ver)          # A: head with CURRENT fp → skip
    _head(db_sync, b.id, "stale", ver)            # B: head with STALE fp → retrain
    db_sync.commit()                              # C: no head → retrain

    need = set(_heads_needing_retrain(db_sync, ids, ver, fps, full=False))
    assert a.id not in need
    assert {b.id, c.id} <= need


def test_stale_embedding_version_forces_retrain(db_sync):
    ver = _version(db_sync)
    a = _tag(db_sync, "A")
    _apply(db_sync, _img(db_sync, "g" * 64).id, a.id)
    db_sync.commit()
    fps = _head_fingerprints(db_sync, [a.id])
    # Matching fingerprint but a DIFFERENT embedding space → must refit.
    _head(db_sync, a.id, fps[a.id], "old-model-v0")
    db_sync.commit()
    assert a.id in set(_heads_needing_retrain(db_sync, [a.id], ver, fps, full=False))


def test_full_forces_all(db_sync):
    ver = _version(db_sync)
    a = _tag(db_sync, "A")
    _apply(db_sync, _img(db_sync, "h" * 64).id, a.id)
    db_sync.commit()
    fps = _head_fingerprints(db_sync, [a.id])
    _head(db_sync, a.id, fps[a.id], ver)          # current fp → would be skipped
    db_sync.commit()
    # full=True ignores the fingerprint (nightly reconcile).
    assert a.id in set(_heads_needing_retrain(db_sync, [a.id], ver, fps, full=True))
