from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

from .hash_codec import DEFAULT_HASH_WIDTH, hash_bytes_to_cam_row

if TYPE_CHECKING:
    from .scenegraph import SimpleSceneGraph


def xnor_popcount_score(
    query_row: int,
    stored_row: int,
    *,
    width: int = DEFAULT_HASH_WIDTH,
) -> int:
    """Compute CAM-style same-bit score for two integer hash rows."""
    if width <= 0:
        raise ValueError("width must be greater than 0")
    mask = (1 << width) - 1
    return int((~(int(query_row) ^ int(stored_row)) & mask).bit_count())


@dataclass(frozen=True)
class CamMatch:
    row_index: int
    obj_id: str
    score: int
    similarity: float
    hash_bytes: bytes


@dataclass(frozen=True)
class SoftwareCamIndex:
    obj_ids: tuple[str, ...]
    rows: tuple[int, ...]
    hashes: tuple[bytes, ...]
    width: int = DEFAULT_HASH_WIDTH

    @classmethod
    def from_scene_graph(
        cls,
        scene_graph: "SimpleSceneGraph",
        *,
        width: int = DEFAULT_HASH_WIDTH,
    ) -> "SoftwareCamIndex":
        obj_ids: list[str] = []
        rows: list[int] = []
        hashes: list[bytes] = []

        for obj_id, node in scene_graph.objects.items():
            hash_bytes = node.visual_hash
            obj_ids.append(obj_id)
            hashes.append(hash_bytes)
            rows.append(hash_bytes_to_cam_row(hash_bytes, width=width))

        return cls(
            obj_ids=tuple(obj_ids),
            rows=tuple(rows),
            hashes=tuple(hashes),
            width=width,
        )

    def query(self, query_hash_bytes: bytes, *, top_k: int = 1) -> list[CamMatch]:
        if top_k <= 0:
            raise ValueError("top_k must be greater than 0")
        if not self.rows:
            raise ValueError("cannot query an empty SoftwareCamIndex")

        query_row = hash_bytes_to_cam_row(query_hash_bytes, width=self.width)
        matches = [
            CamMatch(
                row_index=row_index,
                obj_id=self.obj_ids[row_index],
                score=score,
                similarity=score / float(self.width),
                hash_bytes=self.hashes[row_index],
            )
            for row_index, row in enumerate(self.rows)
            for score in [xnor_popcount_score(query_row, row, width=self.width)]
        ]
        matches.sort(key=lambda match: (-match.score, match.row_index))
        return matches[: min(top_k, len(matches))]