from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Protocol, runtime_checkable

from PIL import Image

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


class ImageHashPipelineOutput(Protocol):
    """Protocol for the output of ImageHashPipeline.process_batch."""

    hash_bits: Any  # torch.Tensor
    cropped_images: list[Image.Image]


@runtime_checkable
class ImageHashPipeline(Protocol):
    """Protocol for an image hash computation pipeline."""

    def process_batch(
        self,
        images: list[Image.Image],
        text_labels: list[str],
        batch_size: int = 32,
    ) -> ImageHashPipelineOutput:
        ...


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


def query_image_against_scene_graph(
    image: Image.Image,
    *,
    pipeline: ImageHashPipeline,
    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:
        return []

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

    results: list[ImageSceneGraphQueryResult] = []
    for query_index, query_bits in enumerate(hash_bits):
        query_hash = bits_tensor_to_hash_bytes(query_bits)
        if scene_graph.objects:
            matches = scene_graph.query_by_visual_hash(query_hash, top_k=top_k)
        else:
            matches = []
        crop = cropped_images[query_index] if query_index < len(cropped_images) else None
        results.append(
            ImageSceneGraphQueryResult(
                query_index=query_index,
                query_hash=query_hash,
                query_crop=crop,
                matches=matches,
            )
        )

    return results