"""
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

from app.database import AsyncSessionLocal
from app.models.photos import Photo

logger = logging.getLogger(__name__)


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


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,
    **_ignored,
) -> dict:
    """Full recompute of duplicate groups using pHash similarity.

    Idempotent — safe to call as often as you like. Returns a summary dict.
    """
    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])

        # ── 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,
    **_ignored,
) -> dict:
    """Incremental duplicate detection for newly added photos using pHash."""
    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)

        # Photo.added_at is stored as TIMESTAMP WITHOUT TIME ZONE, so
        # asyncpg rejects aware datetimes with "can't subtract offset-naive
        # and offset-aware". Normalise: if `since` has a tzinfo, convert
        # it to UTC and drop the tzinfo so the bind parameter is naive.
        if since.tzinfo is not None:
            since = since.astimezone(timezone.utc).replace(tzinfo=None)

        # 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)

        # ── 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 _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)
    )