"""Script to convert Label Studio output JSON files to an Annotation."""

import json
from uuid import UUID
from dotenv import load_dotenv
from python_repositories.adapters import MinioAdapter
from data_store.dto import Annotation
from data_store.repositories import AnnotationRepository


def list_label_studio_output_files(minio_adapter: MinioAdapter) -> list[str]:
    """List the Label Studio output files in MinIO."""
    # Check input
    if not isinstance(minio_adapter, MinioAdapter):
        raise ValueError("minio_adapter must be an instance of MinioAdapter")
    if not minio_adapter.is_connected:
        raise ConnectionError("minio_adapter must be connected to MinIO")
    # List objects with the Label Studio output prefix
    prefix = "label-studio-output/"
    return minio_adapter._list_objects(prefix)


def get_label_studio_output(minio_adapter: MinioAdapter, object_name: str) -> dict:
    """Get a Label Studio output file from MinIO."""
    # Check input
    if not isinstance(minio_adapter, MinioAdapter):
        raise ValueError("minio_adapter must be an instance of MinioAdapter")
    if not minio_adapter.is_connected:
        raise ConnectionError("minio_adapter must be connected to MinIO")
    if not isinstance(object_name, str) or len(object_name) == 0:
        raise ValueError("object_name must be a non-empty string")
    # Get object from MinIO
    buffer = minio_adapter._get(object_name)
    if buffer is None:
        raise FileNotFoundError(f"File not found in MinIO: {object_name}")
    # Convert buffer to JSON
    data: dict = json.load(buffer)
    return data


def convert_label_studio_output(label_studio_output: dict) -> tuple[Annotation, UUID]:
    """Convert Label Studio output to Annotation and corresponding job ID."""
    # Check input
    if not isinstance(label_studio_output, dict):
        raise ValueError("label_studio_output must be a dictionary")
    # Extract job id
    image_path = label_studio_output["task"]["data"]["image"]
    job_id_str = image_path.rsplit("/")[-1].split(".")[0]
    job_id = UUID(job_id_str)
    # Extract relevant information from the Label Studio output
    annotation_data: dict[str, str | None | bool] = {}
    for elem in label_studio_output["result"]:
        # Extract key and value
        key = elem["from_name"]
        choices = elem["value"]["choices"]
        # Check for presence of uncertainty flag
        uncertainty_key = f"{key}_uncertainty"
        uncertainty_value = bool("uncertain" in choices)
        if uncertainty_value:
            choices.remove("uncertain")
        # If no choices left, set to None
        value: str | None = choices[0] if choices else None
        # Store in annotation data
        annotation_data[key] = value
        annotation_data[uncertainty_key] = uncertainty_value
    # Create the Annotation object
    annotation = Annotation(**annotation_data)
    return annotation, job_id


def sync_label_studio_annotations(
    minio_adapter: MinioAdapter,
    annotation_repository: AnnotationRepository,
) -> None:
    """Sync Label Studio output files to AnnotationRepository."""
    # Check inputs
    if not isinstance(minio_adapter, MinioAdapter):
        raise ValueError("minio_adapter must be an instance of MinioAdapter")
    if not minio_adapter.is_connected:
        raise ConnectionError("minio_adapter must be connected to MinIO")
    if not isinstance(annotation_repository, AnnotationRepository):
        raise ValueError("annotation_repository must be an instance of AnnotationRepository")
    if not annotation_repository.is_connected:
        raise ConnectionError("annotation_repository must be connected to MinIO")
    # List annotations in repository
    annotation_ids = annotation_repository.list_all()
    print(f"Found {len(annotation_ids)} annotations in repository.")
    # Count Label Studio output files
    label_studio_files = list_label_studio_output_files(minio_adapter)
    print(f"Found {len(label_studio_files)} Label Studio output files.")
    # Stop early if equal number of annotations and Label Studio files
    if len(annotation_ids) == len(label_studio_files):
        print("No new Label Studio output files to process. Exiting.")
        return
    # Process each Label Studio output file
    for file_name in label_studio_files:
        print(f"Processing file: {file_name}")
        # Get Label Studio output
        try:
            label_studio_output = get_label_studio_output(
                minio_adapter,
                file_name,
            )
        except Exception as exc:
            print(f"Failed to get Label Studio output for {file_name}: {exc}")
            continue
        # Convert to Annotation
        try:
            annotation, job_id = convert_label_studio_output(label_studio_output)
        except Exception as exc:
            print(f"Failed to convert Label Studio output for {file_name}: {exc}")
            continue
        # Store Annotation in repository
        try:
            annotation_repository.put(annotation, job_id)
        except Exception as exc:
            print(f"Failed to store annotation for {file_name}: {exc}")
            continue
        print(f"Successfully processed and stored annotation for {file_name}.")


if __name__ == "__main__":
    # Load environment variables from .env file
    load_dotenv()
    # Connect to data sources
    with MinioAdapter() as minio_adapter, AnnotationRepository() as annotation_repo:
        sync_label_studio_annotations(minio_adapter, annotation_repo)