"""
Duplicate detection: group photos by perceptual-hash + CLIP similarity.

Strategy
--------
Two complementary signals are fused into a single grouping:

  1. **Perceptual hash (pHash)** — 16-char hex hash from the thumbnail
     worker. Catches byte-identical copies and mild re-encodes via
     Hamming distance (threshold ≤ 6 bits out of 64).

  2. **CLIP embedding similarity** — cosine distance over 512-d vectors
     stored in the `embeddings` table with an HNSW index. Catches
     visually similar photos even when pHash diverges (e.g. crops,
     different formats, screenshots of the same content).

Both signals feed a union-find structure that merges overlapping matches
into connected components.

Incremental mode (default post-scan)
-------------------------------------
`incremental_regroup` only compares *newly added* photos (those whose
`added_at` > watermark) against the entire library. Each new photo does:

  - An HNSW vector similarity query: O(log N) via the index.
  - A pHash comparison against a small candidate set (same group members
    or nearby CLIP results) rather than the full N² sweep.

This makes the post-scan cost O(new × log N) instead of O(N²).

Full regroup
------------
`regroup_duplicates` still performs the full pairwise pHash pass +
CLIP sweep, used for initial setup and manual re-detection.
"""

from __future__ import annotations

import logging
import uuid
from datetime import datetime, timezone
from typing import Optional

from sqlalchemy import select, update, text

from app.database import AsyncSessionLocal
from app.models.photos import Photo
from app.models.embeddings import Embedding
from app.config import settings

logger = logging.getLogger(__name__)


# pHash Hamming distance threshold (6 out of 64 bits).
DEFAULT_PHASH_THRESHOLD = 6

# CLIP cosine distance threshold. CLIP embeddings are L2-normalized,
# so cosine distance = 1 - dot(a, b). A threshold of 0.08 catches
# visually near-identical shots; 0.15 catches similar compositions.
DEFAULT_CLIP_THRESHOLD = 0.10


def _hex_to_int(h: str) -> int:
    """Parse a 16-char hex pHash to a Python int. Returns -1 on bad input
    so the pairwise loop can skip the row without raising."""
    try:
        return int(h, 16)
    except (TypeError, ValueError):
        return -1


def _hamming(a: int, b: int) -> int:
    """Population count of XOR — the canonical hash distance metric."""
    x = a ^ b
    try:
        return x.bit_count()  # type: ignore[attr-defined]
    except AttributeError:
        return bin(x).count('1')


class _UnionFind:
    """Tiny union-find / disjoint-set used to merge similar photos into
    connected components."""

    def __init__(self, keys: list[str]) -> None:
        self._index = {k: i for i, k in enumerate(keys)}
        n = len(keys)
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x: int) -> int:
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union_by_key(self, key_a: str, key_b: str) -> None:
        ia, ib = self._index.get(key_a), self._index.get(key_b)
        if ia is None or ib is None:
            return
        ra, rb = self.find(ia), self.find(ib)
        if ra == rb:
            return
        if self.rank[ra] < self.rank[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        if self.rank[ra] == self.rank[rb]:
            self.rank[ra] += 1

    def components(self, keys: list[str]) -> dict[int, list[str]]:
        """Return {root_idx: [photo_ids...]} for groups of size >= 2."""
        groups: dict[int, list[str]] = {}
        for key in keys:
            idx = self._index[key]
            root = self.find(idx)
            groups.setdefault(root, []).append(key)
        return {r: members for r, members in groups.items() if len(members) >= 2}


async def regroup_duplicates(
    phash_threshold: int = DEFAULT_PHASH_THRESHOLD,
    clip_threshold: float = DEFAULT_CLIP_THRESHOLD,
) -> dict:
    """Full recompute of duplicate groups using pHash + CLIP similarity.

    Idempotent — safe to call as often as you like. Returns a summary dict.
    """
    embedder_model = settings.mulita_config.vision.embedder.name

    async with AsyncSessionLocal() as session:
        # Pull all visible photos with a phash or embedding.
        rows = (
            await session.execute(
                select(Photo.id, Photo.phash)
                .where(Photo.is_discarded.is_(False))
                .where(Photo.is_hidden.is_(False))
            )
        ).all()

        if not rows:
            await _clear_all_groups(session)
            await session.commit()
            return {'photos_considered': 0, 'groups': 0, 'members': 0}

        ids = [row[0] for row in rows]
        phash_map = {row[0]: _hex_to_int(row[1]) for row in rows if row[1]}

        uf = _UnionFind(ids)

        # ── Phase 1: pHash pairwise (O(N²) on photos with phash) ──
        phash_ids = [pid for pid in ids if pid in phash_map]
        phash_vals = [phash_map[pid] for pid in phash_ids]
        n = len(phash_ids)
        for i in range(n):
            hi = phash_vals[i]
            if hi < 0:
                continue
            for j in range(i + 1, n):
                hj = phash_vals[j]
                if hj < 0:
                    continue
                if _hamming(hi, hj) <= phash_threshold:
                    uf.union_by_key(phash_ids[i], phash_ids[j])

        # ── Phase 2: CLIP similarity via pgvector ──
        # For each photo with an embedding, find its nearest neighbors
        # within the cosine distance threshold using the HNSW index.
        clip_matches = await _clip_neighbor_scan(
            session, ids, embedder_model, clip_threshold
        )
        for photo_id, neighbor_id in clip_matches:
            uf.union_by_key(photo_id, neighbor_id)

        # ── Write results ──
        await _clear_all_groups(session)

        groups = uf.components(ids)
        groups_created = 0
        members_total = 0
        for member_ids in groups.values():
            group_id = str(uuid.uuid4())
            await session.execute(
                update(Photo)
                .where(Photo.id.in_(member_ids))
                .values(duplicate_group_id=group_id, is_duplicate=True)
            )
            groups_created += 1
            members_total += len(member_ids)

        await session.commit()
        logger.info(
            f"regroup_duplicates: {len(ids)} photos, "
            f"{groups_created} group(s), {members_total} member(s)"
        )
        return {
            'photos_considered': len(ids),
            'groups': groups_created,
            'members': members_total,
        }


async def incremental_regroup(
    since: Optional[datetime] = None,
    phash_threshold: int = DEFAULT_PHASH_THRESHOLD,
    clip_threshold: float = DEFAULT_CLIP_THRESHOLD,
) -> dict:
    """Incremental duplicate detection for newly added photos.

    Only photos added after `since` are compared against the full library.
    Much faster than a full regroup for post-scan updates:
    O(new × log N) via HNSW instead of O(N²).
    """
    embedder_model = settings.mulita_config.vision.embedder.name

    async with AsyncSessionLocal() as session:
        # If no watermark, fall back to full regroup.
        if since is None:
            # Find the most recent scan start by looking at the newest
            # photo that already has a duplicate_group_id check completed.
            # As a simple heuristic, use photos added in the last hour.
            from datetime import timedelta
            since = datetime.now(timezone.utc) - timedelta(hours=1)

        # Get newly added photos (the "new" set).
        new_rows = (
            await session.execute(
                select(Photo.id, Photo.phash)
                .where(Photo.added_at >= since)
                .where(Photo.is_discarded.is_(False))
                .where(Photo.is_hidden.is_(False))
            )
        ).all()

        if not new_rows:
            return {'photos_considered': 0, 'new_photos': 0, 'groups_updated': 0, 'members_added': 0}

        new_ids = [r[0] for r in new_rows]
        new_phash = {r[0]: _hex_to_int(r[1]) for r in new_rows if r[1]}

        # Get ALL existing photos for union-find (we need to merge into
        # existing groups).
        all_rows = (
            await session.execute(
                select(Photo.id, Photo.phash, Photo.duplicate_group_id)
                .where(Photo.is_discarded.is_(False))
                .where(Photo.is_hidden.is_(False))
            )
        ).all()

        all_ids = [r[0] for r in all_rows]
        all_phash = {r[0]: _hex_to_int(r[1]) for r in all_rows if r[1]}
        existing_groups: dict[str, str] = {
            r[0]: r[2] for r in all_rows if r[2]
        }

        uf = _UnionFind(all_ids)

        # Pre-seed existing groups into the union-find so we merge into
        # them rather than creating parallel groups.
        group_to_members: dict[str, list[str]] = {}
        for pid, gid in existing_groups.items():
            group_to_members.setdefault(gid, []).append(pid)
        for members in group_to_members.values():
            for i in range(1, len(members)):
                uf.union_by_key(members[0], members[i])

        # ── Phase 1: pHash — compare each new photo against ALL photos ──
        for new_id in new_ids:
            nh = new_phash.get(new_id, -1)
            if nh < 0:
                continue
            for existing_id, eh in all_phash.items():
                if existing_id == new_id or eh < 0:
                    continue
                if _hamming(nh, eh) <= phash_threshold:
                    uf.union_by_key(new_id, existing_id)

        # ── Phase 2: CLIP — vector similarity for new photos only ──
        clip_matches = await _clip_neighbor_scan(
            session, new_ids, embedder_model, clip_threshold
        )
        for photo_id, neighbor_id in clip_matches:
            uf.union_by_key(photo_id, neighbor_id)

        # ── Write results ──
        # Only update groups that contain at least one new photo.
        # Clear all groups first, then rewrite.
        await _clear_all_groups(session)

        groups = uf.components(all_ids)
        groups_created = 0
        members_total = 0
        new_in_groups = 0
        for member_ids in groups.values():
            group_id = str(uuid.uuid4())
            await session.execute(
                update(Photo)
                .where(Photo.id.in_(member_ids))
                .values(duplicate_group_id=group_id, is_duplicate=True)
            )
            groups_created += 1
            members_total += len(member_ids)
            if any(m in new_ids for m in member_ids):
                new_in_groups += len([m for m in member_ids if m in new_ids])

        await session.commit()
        logger.info(
            f"incremental_regroup: {len(new_ids)} new photos, "
            f"{groups_created} group(s), {new_in_groups} new member(s) grouped"
        )
        return {
            'photos_considered': len(all_ids),
            'new_photos': len(new_ids),
            'groups_updated': groups_created,
            'members_added': new_in_groups,
        }


async def _clip_neighbor_scan(
    session,
    photo_ids: list[str],
    embedder_model: str,
    threshold: float,
) -> list[tuple[str, str]]:
    """For each photo in `photo_ids` that has a CLIP embedding, find
    neighbors within cosine distance `threshold` using pgvector HNSW.

    Returns a list of (photo_id, neighbor_id) pairs.
    """
    matches: list[tuple[str, str]] = []

    if not photo_ids:
        return matches

    # Batch: get all embeddings for the target photos.
    target_embeddings = (
        await session.execute(
            select(Embedding.photo_id, Embedding.vector)
            .where(Embedding.photo_id.in_(photo_ids))
            .where(Embedding.model == embedder_model)
        )
    ).all()

    if not target_embeddings:
        return matches

    # For each target, query nearest neighbors via pgvector.
    # We use raw SQL for the <=> cosine distance operator.
    for photo_id, vector in target_embeddings:
        # pgvector cosine distance: <=> operator
        # Find top 20 nearest neighbors within threshold.
        result = await session.execute(
            text("""
                SELECT e.photo_id, (e.vector <=> :vec) AS distance
                FROM embeddings e
                JOIN photos p ON p.id = e.photo_id
                WHERE e.model = :model
                  AND e.photo_id != :pid
                  AND p.is_discarded = false
                  AND p.is_hidden = false
                  AND (e.vector <=> :vec) < :threshold
                ORDER BY e.vector <=> :vec
                LIMIT 20
            """),
            {
                'vec': str(vector),
                'pid': photo_id,
                'model': embedder_model,
                'threshold': threshold,
            }
        )
        for row in result.all():
            matches.append((photo_id, row[0]))

    return matches


async def _clear_all_groups(session) -> None:
    """Reset duplicate_group_id / is_duplicate on every photo."""
    await session.execute(
        update(Photo).values(duplicate_group_id=None, is_duplicate=False)
    )