from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from PIL import Image

from .hash_codec import bits_tensor_to_hash_bytes
from .scenegraph import SceneGraphMatch, SimpleSceneGraph


@dataclass(frozen=True)
class ImageSceneGraphQueryResult:
    query_crop_index: int
    query_hash: bytes
    query_crop: Image.Image
    matches: list[SceneGraphMatch]


def query_image_against_scene_graph(
    image: Image.Image,
    pipeline: Any,
    scene_graph: SimpleSceneGraph,
    text_labels: list[str],
    *,
    top_k: int = 1,
    batch_size: int = 1,
) -> list[ImageSceneGraphQueryResult]:
    output = pipeline.process_batch([image], text_labels, batch_size=batch_size)
    hash_bits = output.hash_bits
    cropped_images = list(output.cropped_images)

    if hash_bits.numel() == 0:
        if cropped_images:
            raise ValueError("hash_bits and cropped_images must align")
        return []

    if hash_bits.dim() == 1:
        hash_bits = hash_bits.unsqueeze(0)

    if hash_bits.shape[0] != len(cropped_images):
        raise ValueError("hash_bits and cropped_images must align")

    results: list[ImageSceneGraphQueryResult] = []
    for crop_index, query_bits in enumerate(hash_bits):
        query_hash = bits_tensor_to_hash_bytes(query_bits)
        matches = scene_graph.query_by_visual_hash(query_hash, top_k=top_k)
        results.append(
            ImageSceneGraphQueryResult(
                query_crop_index=crop_index,
                query_hash=query_hash,
                query_crop=cropped_images[crop_index],
                matches=matches,
            )
        )

    return results