import os
import torch
from .core import ImageSimilaritySearch
from typing import Literal, List, TypedDict


class Results(TypedDict):
    path: str
    distance: float


def search(
    query_image: str,
    index_path: str,
    indexed_images_path: str,
    index_type: str = "L2",
    top_k: int = 10,
    model_type: Literal[
        "resnet18",
        "resnet34",
        "resnet50",
        "resnet101",
        "resnet152",
        "dinov2_vits14",
        "dinov2_vits14_reg",
        "dinov2_vitb14",
        "dinov2_vitb14_reg",
        "dinov2_vitl14",
        "dinov2_vitl14_reg",
        "dinov2_vitg14",
        "dinov2_vitg14_reg",
    ] = "dinov2_vits14_reg",
    device: Literal["cuda", "mps", "cpu"] = "cuda",
) -> List[Results]:
    """

    Args:
        query_image (str): Path to query for kNN
        index_path (str): Path to saved index
        indexed_images_path (str): Path to file containing indexed image paths
        index_type (str): Type of index to use (currently only using L2)
        top_k (int): top k-Nearest Neighbors to return
        model_type (str): Type of model to use
        device (str): Device to use for feature extraction ("cuda", "mps", "cpu")

    Return:
        results (List[Results]): A list of dictionaries containing:
            - path (str): File path of the indexed image
            - distance (float): Distance between the query image and the result
    """
    if not (os.path.exists(index_path) and os.path.exists(indexed_images_path)):
        print(
            "No index found, check your index path and image paths file or first index a dataset."
        )
        return None

    if "_reg" in model_type and torch.mps.is_available():
        print(
            "Some PyTorch functionality using DINOv2 models with Registers "
            "not supported by MPS yet, reverting to CPU"
        )
        device = "cpu"
    else:
        device = (
            device
            if device
            else (
                "cuda"
                if torch.cuda.is_available()
                else "mps" if torch.mps.is_available() else "cpu"
            )
        )

    searcher = ImageSimilaritySearch(
        model_type=model_type, index_type=index_type, device=device
    )  # Initialize the similarity search w/same model during indexing
    searcher.load(
        index_path, indexed_images_path
    )  # Load the previously saved index and image paths file
    results = searcher.find_similar(query_image, top_k=top_k)  # Run sample search
    print(f"Querying top {top_k} similar images to: {query_image}")
    print("Matches:")
    for idx, result in enumerate(results):
        print(f"{idx + 1}: {result["path"]} - Distance: {result["distance"]}")

    return results


if __name__ == "__main__":
    search()
