"""Module providing server functionality."""

import logging
import os
import traceback
from queue import Queue
from typing import Any, Dict, List, Optional
import requests
from PIL import Image
from matrice.data_processing.data_formats.mscoco_detection import (
    get_msococo_images_details,
    add_mscoco_dataset_items_details,
)
from matrice.data_processing.data_formats.imagenet_classification import (
    add_imagenet_dataset_items_details,
)
from matrice.data_processing.data_formats.pascalvoc_detection import (
    get_pascalvoc_image_details,
    add_pascalvoc_dataset_items_details,
)
from matrice.data_processing.data_formats.labelbox_detection import (
    get_labelbox_image_details,
    add_labelbox_dataset_items_details,
    add_labelbox_dataset_item_local_file_path,
    download_labelbox_dataset_items,
)
from matrice.data_processing.data_formats.labelbox_classification import (
    get_labelbox_classification_image_details,
    add_labelbox_classification_dataset_items_details,
    add_labelbox_classification_dataset_item_local_file_path,
)
from matrice.data_processing.data_formats.yolo_detection import (
    get_yolo_image_details,
    add_yolo_dataset_items_details,
    convert_payload_to_coco_format,
)
from matrice.data_processing.data_formats.unlabelled import (
    add_unlabelled_dataset_items_details,
)
from matrice.data_processing.server_utils import (
    download_file,
    rpc_get_call,
    get_batch_pre_signed_download_urls,
    get_filename_from_url,
    update_partition_status,
    update_video_frame_partition_status,
    get_unprocessed_partitions,
    extract_dataset,
    get_partition_items,
    chunk_items,
    handle_source_url_dataset_download,
    get_video_frame_partition_items,
)
from matrice.data_processing.client_utils import (
    scan_folder,
)
from matrice.data_processing.pipeline import (
    Pipeline,
)
from matrice.data_processing.data_formats.video_youtube_bb_tracking import (
    get_youtube_bb_video_frame_details,
    add_youtube_bb_dataset_items_details,
)
from matrice.data_processing.data_formats.video_mot_tracking import (
    get_mot_annotations,
    add_mot_dataset_items_details,
)
from matrice.data_processing.data_formats.video_davis_segmentation import (
    get_davis_annotations,
    add_davis_dataset_items_details,
)
from matrice.data_processing.data_formats.video_imagenet_classification import (
    add_video_imagenet_dataset_items_details,
)
from matrice.data_processing.data_formats.video_kinetics_activity_recognition import (
    get_kinetics_annotations,
    add_kinetics_dataset_items_details,
)
from matrice.data_processing.data_formats.video_detection_mscoco import (
    add_video_mscoco_dataset_items_details,
    get_video_mscoco_annotations,
)

TMP_FOLDER = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "tmp",
)
os.makedirs(TMP_FOLDER, exist_ok=True)
logging.info("Created temporary folder at %s", TMP_FOLDER)


def download_labelbox_dataset(dataset_id, rpc, dataset_version, source_url):
    """Download a dataset from Labelbox.

    Args:
        dataset_id: ID of the dataset
        rpc: RPC client for making API calls
        dataset_version: Version of the dataset
        source_url: Optional source URL to download from

    Returns:
        Path to the downloaded dataset
    """
    if source_url:
        dataset_path = handle_source_url_dataset_download(source_url)
    else:
        logging.info("Downloading annotation file for labelbox dataset")
        dataset_path = get_annotation_files(
            rpc,
            dataset_id,
            dataset_version,
            False,
        )[0]
    logging.info("Downloading dataset from labelbox")
    dataset_path = download_labelbox_dataset_items(dataset_id, dataset_path)
    return dataset_path


def partition_items_producer(
    rpc: Any,
    dataset_id: str,
    partition: int,
    pipeline_queue: Queue,
    download_images_required: bool = False,
    request_batch_size: int = 1000,
    processing_batch_size: int = 10,
) -> None:
    """Get items for a partition and add them to the pipeline queue.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        partition: Partition number
        pipeline_queue: Queue to add items to
        download_images_required: Whether to get presigned URLs for images
        request_batch_size: Number of items to fetch per API request
        processing_batch_size: Size of batches to add to pipeline queue
    """
    try:
        all_dataset_items = get_partition_items(
            rpc,
            dataset_id,
            partition,
            download_images_required,
            request_batch_size,
        )
        processing_batches = chunk_items(
            all_dataset_items,
            processing_batch_size,
        )
        for batch in processing_batches:
            pipeline_queue.put(batch)
        logging.info(
            "Successfully fetched %s items for partition %s",
            len(all_dataset_items),
            partition,
        )
    except Exception as e:
        logging.error(
            "Error processing partition %s: %s",
            partition,
            e,
        )
        traceback.print_exc()


def video_frame_partition_items_producer(
    rpc: Any,
    dataset_id: str,
    partition: int,
    pipeline_queue: Queue,
    download_images_required: bool = False,
    request_batch_size: int = 1000,
    processing_batch_size: int = 10,
    isFileInfoRequired: bool = True,
) -> None:
    """Get items for a partition and add them to the pipeline queue.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        partition: Partition number
        pipeline_queue: Queue to add items to
        download_images_required: Whether to get presigned URLs for images
        request_batch_size: Number of items to fetch per API request
        processing_batch_size: Size of batches to add to pipeline queue
    """
    try:
        all_dataset_items = get_video_frame_partition_items(
            rpc,
            dataset_id,
            partition,
            download_images_required,
            request_batch_size,
            isFileInfoRequired,
        )
        processing_batches = chunk_items(
            all_dataset_items,
            processing_batch_size,
        )
        for batch in processing_batches:
            pipeline_queue.put(batch)
        logging.info(
            "Successfully fetched %s items for partition %s",
            len(all_dataset_items),
            partition,
        )
    except Exception as e:
        logging.error(
            "Error processing partition %s: %s",
            partition,
            e,
        )
        traceback.print_exc()


def download_samples(
    image_details: Dict[str, Any],
    rpc: Any,
    bucket_alias: str = "",
    account_number: str = "",
) -> Dict[str, Any]:
    """Download sample and update sample details.

    Args:
        image_details: Dictionary containing image metadata
        rpc: RPC client for making API calls
        bucket_alias: Bucket alias
        account_number: Account number

    Returns:
        Updated sample details dictionary
    """
    if image_details.get("is_complete"):
        return image_details
    dataset_item = image_details.get("sample_details")
    try:
        if not dataset_item.get("cloudPath"):
            dataset_item["cloudPath"] = get_batch_pre_signed_download_urls(
                dataset_item.get("fileLocation"),
                rpc,
                bucket_alias,
                account_number,
            )[dataset_item.get("fileLocation")]
        if not dataset_item.get("local_file_path"):
            dataset_item["local_file_path"] = os.path.join(
                TMP_FOLDER,
                dataset_item["filename"],
            )
        os.makedirs(
            os.path.dirname(dataset_item["local_file_path"]),
            exist_ok=True,
        )
        download_file(
            dataset_item["cloudPath"],
            dataset_item["local_file_path"],
        )
        return {
            "sample_details": dataset_item,
            "is_complete": False,
        }
    except Exception as e:
        logging.error(
            "Error downloading image %s: %s",
            dataset_item.get("filename"),
            e,
        )
        return {
            "sample_details": dataset_item,
            "is_complete": False,
        }


def get_pre_signed_upload_urls(
    cloud_file_paths,
    rpc,
    file_type,
    bucket_alias="",
    account_number="",
):
    """Get pre-signed upload URLs for files.

    Args:
        cloud_file_paths: Paths of files in cloud storage
        rpc: RPC client for making API calls
        file_type: Type of files
        bucket_alias: Bucket alias
        account_number: Account number

    Returns:
        Response from API containing pre-signed URLs
    """
    logging.debug(
        "Getting presigned URLs for %d files",
        len(cloud_file_paths),
    )
    logging.debug(
        "cloud_file_paths for upload: %s",
        cloud_file_paths,
    )
    if not isinstance(cloud_file_paths, list):
        cloud_file_paths = [cloud_file_paths]
    logging.debug(
        "final cloud_file_paths for upload: %s",
        cloud_file_paths,
    )
    payload_get_presigned_url = {
        "fileNames": cloud_file_paths,
        "type": "samples",
        "isPrivateBucket": (True if bucket_alias else False),
        "bucketAlias": bucket_alias,
        "accountNumber": account_number,
    }
    resp = rpc.post(
        "/v2/dataset/get_batch_pre_signed_upload_urls",
        payload=payload_get_presigned_url,
    )
    logging.debug(
        "payload for getting the presigned urls for first frames: %s",
        payload_get_presigned_url,
    )
    if resp["success"]:
        logging.debug(
            "presiged urls for upload: %s",
            resp["data"],
        )
        return resp["data"]
    else:
        logging.error(
            "Failed to get presigned URLs: %s",
            resp["message"],
        )
        return resp["message"]

def download_video_samples(
    image_details: Dict[str, Any],
    rpc: Any,
    bucket_alias: str = "",
    account_number: str = "",
) -> Dict[str, Any]:
    """Download sample and update sample details.

    Args:
        image_details: Dictionary containing image metadata
        rpc: RPC client for making API calls
        bucket_alias: Bucket alias
        account_number: Account number

    Returns:
        Updated sample details dictionary
    """
    logging.debug(
        "keys of image details are %s",
        image_details.keys(),
    )
    logging.debug("%s", image_details.get("is_complete"))
    logging.debug(
        "image details for download: %s",
        image_details,
    )
    if image_details.get("is_complete"):
        return image_details
    dataset_item = image_details.get("sample_details")
    logging.debug("dataset item is %s", dataset_item)
    file_info_frames = dataset_item["fileInfoResponse"]["frames"]
    _, file_info = next(iter(file_info_frames.items()))
    logging.debug("file info is %s", file_info)
    try:
        if not file_info.get("cloudPath"):
            file_info["cloudPath"] = get_batch_pre_signed_download_urls(
                file_info.get("fileLocation"),
                rpc,
                bucket_alias,
                account_number,
            )[file_info.get("fileLocation")]
        if not dataset_item.get("local_file_path"):
            dataset_item["local_file_path"] = os.path.join(
                TMP_FOLDER,
                file_info["filename"],
            )
        os.makedirs(
            os.path.dirname(dataset_item["local_file_path"]),
            exist_ok=True,
        )
        download_file(
            file_info["cloudPath"],
            dataset_item["local_file_path"],
        )
        return {
            "sample_details": dataset_item,
            "is_complete": False,
        }
    except Exception as e:
        logging.error(
            "Error downloading image %s: %s",
            dataset_item.get("filename"),
            e,
        )
        return {
            "sample_details": dataset_item,
            "is_complete": False,
        }


def upload_file(local_path, presigned_url, max_attempts=5):
    """Upload a file to a presigned URL.

    Args:
        local_path: Local path of the file to upload
        presigned_url: Pre-signed URL to upload to
        max_attempts: Maximum number of upload attempts

    Returns:
        Boolean indicating success of upload
    """
    logging.debug(
        "Uploading %s to %s",
        local_path,
        presigned_url,
    )
    for attempt in range(max_attempts):
        try:
            with open(local_path, "rb") as f:
                response = requests.put(
                    presigned_url,
                    data=f,
                    allow_redirects=True,
                    timeout=30,
                )
                logging.info(f"response from uploading is {response}")
                if response.status_code == 200:
                    logging.info(
                        "Successfully uploaded %s to %s",
                        local_path,
                        presigned_url,
                    )
                    return True
                else:
                    logging.warning(
                        "Failed to upload %s (status: %s), attempt %s/%s",
                        local_path,
                        response.status_code,
                        attempt + 1,
                        max_attempts,
                    )
                    response.raise_for_status()
        except Exception as e:
            if attempt == max_attempts - 1:
                logging.error(
                    "Failed to upload %s after %s attempts. Error: %s",
                    local_path,
                    max_attempts,
                    e,
                )
                return False
            else:
                logging.warning(
                    "Attempt %s/%s failed for %s. Retrying... Error: %s",
                    attempt + 1,
                    max_attempts,
                    local_path,
                    e,
                )


def upload_video_samples(
    image_details: Dict[str, Any],
    rpc: Any,
    bucket_alias: str = "",
    account_number: str = "",
) -> Dict[str, Any]:
    """Upload video samples to storage.

    Args:
        image_details: Dictionary with sample details
        rpc: RPC client for API calls
        bucket_alias: Bucket alias for private storage
        account_number: Account number for private storage

    Returns:
        Updated image details
    """
    dataset_item = image_details.get("sample_details")
    is_complete = image_details.get("is_complete")
    logging.debug(
        "dataset item is %s for upload",
        dataset_item,
    )
    try:
        logging.debug(
            "dataset item key for before fetching url and uploading- %s",
            dataset_item.get("first_frame_upload_cloud_path"),
        )
        if not dataset_item.get("first_frame_upload_cloud_path"):
            logging.debug(
                "getting presigned url by updating the key first_frame_upload_cloud_path %s",
                dataset_item.get("first_frame_upload_cloud_path"),
            )
            cloud_paths_presigned_url_dict = get_pre_signed_upload_urls(
                dataset_item.get("bucket_upload_first_frame_path"),
                rpc,
                bucket_alias,
                account_number,
            )
            logging.info(f"presigned url for upload dictionary is-:{cloud_paths_presigned_url_dict}"),
            dataset_item.update(
                {
                    "first_frame_upload_cloud_path": cloud_paths_presigned_url_dict.get(
                        dataset_item.get("bucket_upload_first_frame_path")
                    )
                }
            )
            logging.info(f"The first frame upload url is {dataset_item.get('first_frame_upload_cloud_path')}")
        upload_file(
            dataset_item["first_frame_path"],
            dataset_item["first_frame_upload_cloud_path"],
        )
        return {
            "sample_details": dataset_item,
            "is_complete": is_complete,
        }
    except Exception as e:
        logging.error(
            "Error uploading sample %s: %s",
            dataset_item.get("filename", "unknown"),
            e,
        )
        return {
            "sample_details": dataset_item,
            "is_complete": False,
        }


def batch_download_samples(
    batch_image_details: List[Dict[str, Any]],
    rpc: Any,
    bucket_alias: str = "",
    account_number: str = "",
) -> List[Dict[str, Any]]:
    """Download a batch of samples.

    Args:
        batch_image_details: List of image details dictionaries
        rpc: RPC client for making API calls

    Returns:
        List of updated sample details
    """
    logging.debug(
        "Processing batch of %s samples for download",
        len(batch_image_details),
    )
    return [
        download_samples(
            image_details,
            rpc,
            bucket_alias,
            account_number,
        )
        for image_details in batch_image_details
    ]


def batch_upload_video_samples(
    batch_image_details: List[Dict[str, Any]],
    rpc: Any,
    bucket_alias: str = "",
    account_number: str = "",
) -> List[Dict[str, Any]]:
    """Download a batch of samples.

    Args:
        batch_image_details: List of image details dictionaries
        rpc: RPC client for making API calls

    Returns:
        List of updated sample details
    """
    logging.debug(
        "Processing batch of %s samples for upload",
        len(batch_image_details),
    )
    return [
        upload_video_samples(
            image_details,
            rpc,
            bucket_alias,
            account_number,
        )
        for image_details in batch_image_details
    ]

def batch_download_video_samples(
    batch_image_details: List[Dict[str, Any]],
    rpc: Any,
    bucket_alias: str = "",
    account_number: str = "",
) -> List[Dict[str, Any]]:
    """Download a batch of samples.

    Args:
        batch_image_details: List of image details dictionaries
        rpc: RPC client for making API calls

    Returns:
        List of updated sample details
    """
    logging.debug(
        "Processing batch of %s samples for download",
        len(batch_image_details),
    )
    return [
        download_video_samples(
            image_details,
            rpc,
            bucket_alias,
            account_number,
        )
        for image_details in batch_image_details
    ]


def calculate_image_properties(
    image_details: Dict[str, Any],
) -> Dict[str, Any]:
    """Calculate properties of an image.

    Args:
        image_details: Dictionary containing image metadata

    Returns:
        Updated image details with calculated properties
    """
    if image_details.get("is_complete"):
        return image_details
    dataset_item = image_details.get("sample_details")
    try:     
        with Image.open(dataset_item["local_file_path"]) as image:
            width, height = image.size
            dataset_item.update(
                {
                    "image_height": height,
                    "image_width": width,
                    "image_area": height * width,
                }
            )
            os.remove(dataset_item["local_file_path"])
            return {
                "sample_details": dataset_item,
                "is_complete": True,
            }
    except Exception as e:
        logging.error(
            "Error processing image %s: %s",
            dataset_item.get("filename"),
            e,
        )
        return {
            "sample_details": dataset_item,
            "is_complete": False,
        }


def batch_calculate_sample_properties(
    batch_sample_details: List[Dict[str, Any]],
    properties_calculation_fn: callable,
) -> List[Dict[str, Any]]:
    """Calculate properties for a batch of samples.

    Args:
        batch_image_details: List of image details dictionaries

    Returns:
        List of processed image details
    """
    logging.debug(
        "Processing batch of %s samples for property calculation",
        len(batch_sample_details),
    )
    processed_batch = []
    for dataset_item in batch_sample_details:
        dataset_item = properties_calculation_fn(dataset_item)
        if dataset_item.get("is_complete"):
            processed_batch.append(dataset_item["sample_details"])
    return processed_batch


def batch_update_video_dataset_items(
    batch_image_details: List[Dict[str, Any]],
    rpc: Any,
    dataset_id: str,
    version: str,
    attempts: int = 3,
    is_yolo: bool = False,
) -> List[Dict[str, Any]]:
    """Update video dataset items in batch.

    Args:
        batch_image_details: List of dictionaries containing image details
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        version: Version of the dataset
        attempts: Number of retry attempts
        is_yolo: Whether using YOLO format

    Returns:
        Updated batch image details
    """
    retry_count = 0
    logging.debug(
        "Batch image details are %s",
        batch_image_details,
    )
    while retry_count < attempts:
        try:
            logging.debug(
                "Attempting to update batch of %s items (attempt %s/%s)",
                len(batch_image_details),
                retry_count + 1,
                attempts,
            )
            items = []
            count_dataset_item=0
            for dataset_item in batch_image_details:
                count_dataset_item+=1
                logging.debug(
                    "Processing dataset item for updating-: %s",
                    dataset_item,
                )
                frame_annotations = {}
                file_info = dataset_item.get("fileInfoResponse", {})
                frames = file_info.get("frames", {})
                for (
                    frame_id,
                    ann,
                ) in frames.items():
                    frame_annotations[frame_id] = []
                    frame_annotations[frame_id].append(
                        {
                            "id": ann.get("id"),
                            "segmentation": ann.get("segmentation", []),
                            "isCrowd": ann.get("isCrowd", []),
                            "confidence": 0.0,
                            "bbox": [
                                ann["bbox"][0],
                                ann["bbox"][1],
                                ann["bbox"][2] + ann["bbox"][0],
                                ann["bbox"][3] + ann["bbox"][1],
                            ],
                            "height": ann.get("height"),
                            "width": ann.get("width"),
                            "center": ann.get("center", []),
                            "area": ann.get("area", 0),
                            "category": ann.get("category"),
                            "masks": ann.get("masks", []),
                        }
                    )
                item = {
                    "datasetItemId": str(file_info.get("_idVideoDatasetItem")),
                    "version": str(version),
                    "splitType": str(dataset_item.get("splitType")),
                    "frameWiseAnnotations": frame_annotations,
                    "height": dataset_item.get("video_height"),
                    "width": dataset_item.get("video_width"),
                    "area": int(dataset_item.get("video_height"))
                    * int(dataset_item.get("video_width")),
                    "fps": dataset_item.get("frame_rate"),
                }
                items.append(item)
            payload = {
                "datasetId": str(dataset_id),
                "items": items,
            }
            resp = rpc.put(
                path="/v2/dataset/update-video-dataset-items/",
                payload=payload,
            )
            logging.debug(
                "Update dataset items payload: %s",
                payload,
            )
            if resp.get("success"):
                logging.debug(
                    "Successfully updated batch of %s items",
                    len(batch_image_details),
                )
                logging.debug("successfully updated %s number of dataset items",
                              count_dataset_item)
                for item in batch_image_details:
                    item["status"] = "processed"
                return batch_image_details
            logging.error(
                "Failed to update batch: %s",
                resp.get("data"),
            )
            retry_count += 1
        except Exception as e:
            logging.error("Error updating batch: %s", e)
            retry_count += 1
    for item in batch_image_details:
        item["status"] = "errored"
    return batch_image_details


def batch_update_video_mot_dataset_items(
    batch_image_details: List[Dict[str, Any]],
    rpc: Any,
    dataset_id: str,
    version: str,
    attempts: int = 3,
    is_yolo: bool = False,
    batch_size: int = 30,
) -> List[Dict[str, Any]]:
    """Update MOT video dataset items in batch.

    Args:
        batch_image_details: List of dictionaries containing image details
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        version: Version of the dataset
        attempts: Number of retry attempts
        is_yolo: Whether using YOLO format
        batch_size: Size of processing batches

    Returns:
        Updated batch image details
    """
    retry_count = 0
    logging.debug(
        "Batch image details are %s",
        batch_image_details,
    )
    while retry_count < attempts:
        try:
            logging.debug(
                "Attempting to update batch of %s items (attempt %s/%s)",
                len(batch_image_details),
                retry_count + 1,
                attempts,
            )
            for dataset_item in batch_image_details:
                logging.debug(
                    "Processing dataset item for updating: %s",
                    dataset_item,
                )
                file_info = dataset_item.get("fileInfoResponse", {})
                frames = file_info.get("frames", {})
                sorted_frames = sorted(frames.items())
                for i in range(
                    0,
                    len(sorted_frames),
                    batch_size,
                ):
                    chunk_frames = sorted_frames[i : i + batch_size]
                    frame_annotations = {}
                    for (
                        frame_id,
                        frame_info,
                    ) in chunk_frames:
                        annotations = frame_info.get("annotations", [])
                        frame_annotations[frame_id] = [
                            {
                                "id": ann.get("id"),
                                "segmentation": ann.get(
                                    "segmentation",
                                    [],
                                ),
                                "isCrowd": ann.get("isCrowd", []),
                                "confidence": ann.get(
                                    "confidence",
                                    0.5,
                                ),
                                "bbox": [
                                    ann["bbox"][0],
                                    ann["bbox"][1],
                                    ann["bbox"][2] + ann["bbox"][0],
                                    ann["bbox"][3] + ann["bbox"][1],
                                ],
                                "height": ann.get("height"),
                                "width": ann.get("width"),
                                "center": ann.get("center", []),
                                "area": ann.get("area", 0),
                                "category": ann.get("category"),
                                "masks": ann.get("masks", []),
                            }
                            for ann in annotations
                            if len(ann.get("bbox", [])) >= 4
                        ]
                    payload = {
                        "datasetId": str(dataset_id),
                        "items": [
                            {
                                "datasetItemId": str(file_info.get("_idVideoDatasetItem")),
                                "version": str(version),
                                "splitType": str(dataset_item.get("splitType")),
                                "frameWiseAnnotations": frame_annotations,
                                "height": dataset_item.get("video_height"),
                                "width": dataset_item.get("video_width"),
                                "area": int(
                                    dataset_item.get(
                                        "video_height",
                                        0,
                                    )
                                )
                                * int(
                                    dataset_item.get(
                                        "video_width",
                                        0,
                                    )
                                ),
                                "fps": dataset_item.get("frame_rate"),
                            }
                        ],
                    }
                    try:
                        logging.debug(
                            "Sending update for frames %s to %s: %s",
                            i,
                            i + batch_size,
                            payload,
                        )
                        resp = rpc.put(
                            path="/v2/dataset/update-video-dataset-items/",
                            payload=payload,
                        )
                        logging.debug(
                            "response from update-video-dataset-items: %s",
                            resp,
                        )
                        if resp and resp.get("success"):
                            logging.debug(
                                "Successfully updated frames %s to %s",
                                i,
                                i + batch_size,
                            )
                            dataset_item["status"] = "processed"
                        else:
                            error_msg = resp.get("data") if resp else "No response from server"
                            logging.error(
                                "Failed to update frames %s to %s: %s",
                                i,
                                i + batch_size,
                                error_msg,
                            )
                            dataset_item["status"] = "errored"
                    except Exception as rpc_err:
                        logging.error(
                            "RPC call failed for frames %s to %s: %s",
                            i,
                            i + batch_size,
                            rpc_err,
                        )
                        dataset_item["status"] = "errored"
            return batch_image_details
        except Exception as e:
            logging.error("Error updating batch: %s", e)
            retry_count += 1
    for item in batch_image_details:
        item["status"] = "errored"
    return batch_image_details


def batch_update_video_davis_dataset_items(
    batch_image_details: List[Dict[str, Any]],
    rpc: Any,
    dataset_id: str,
    version: str,
    attempts: int = 3,
    is_yolo: bool = False,
    batch_size: int = 30,
) -> List[Dict[str, Any]]:
    """Update dataset items in batch, processing frames in groups of 30.

    Args:
        batch_image_details: List of image details to update
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        version: Version of the dataset
        attempts: Number of retry attempts
        is_yolo: Whether the dataset is in YOLO format
        batch_size: Number of frames to process in one API call

    Returns:
        List of updated dataset items
    """
    retry_count = 0
    logging.debug(
        "Batch image details are %s",
        batch_image_details,
    )
    while retry_count < attempts:
        try:
            logging.debug(
                "Attempting to update batch of %s items (attempt %s/%s)",
                len(batch_image_details),
                retry_count + 1,
                attempts,
            )
            for dataset_item in batch_image_details:
                logging.debug(
                    "Processing dataset item for updating: %s",
                    dataset_item,
                )
                file_info = dataset_item.get("fileInfoResponse", {})
                frames = file_info.get("frames", {})
                sorted_frames = sorted(frames.items())
                for i in range(
                    0,
                    len(sorted_frames),
                    batch_size,
                ):
                    chunk_frames = sorted_frames[i : i + batch_size]
                    frame_annotations = {}
                    for (
                        frame_id,
                        frame_info,
                    ) in chunk_frames:
                        annotations = frame_info.get("annotations", [])
                        frame_annotations[frame_id] = [
                            {
                                "id": ann.get("id"),
                                "segmentation": ann.get(
                                    "segmentation",
                                    [],
                                ),
                                "isCrowd": ann.get("isCrowd", []),
                                "confidence": ann.get(
                                    "confidence",
                                    0.5,
                                ),
                                "bbox": [
                                    ann["bbox"][0],
                                    ann["bbox"][1],
                                    ann["bbox"][2],
                                    ann["bbox"][3],
                                ],
                                "height": ann.get("height"),
                                "width": ann.get("width"),
                                "center": ann.get("center", []),
                                "area": ann.get("area", 0),
                                "category": ann.get("category"),
                                "masks": ann.get("masks", []),
                            }
                            for ann in annotations
                            if len(ann.get("bbox", [])) >= 4
                        ]
                    payload = {
                        "datasetId": str(dataset_id),
                        "items": [
                            {
                                "datasetItemId": str(file_info.get("_idVideoDatasetItem")),
                                "version": str(version),
                                "splitType": str(dataset_item.get("splitType")),
                                "frameWiseAnnotations": frame_annotations,
                                "height": dataset_item.get("video_height"),
                                "width": dataset_item.get("video_width"),
                                "area": int(
                                    dataset_item.get(
                                        "video_height",
                                        0,
                                    )
                                )
                                * int(
                                    dataset_item.get(
                                        "video_width",
                                        0,
                                    )
                                ),
                                "fps": dataset_item.get("frame_rate"),
                            }
                        ],
                    }
                    try:
                        logging.debug(
                            "Sending update for frames %s to %s: %s",
                            i,
                            i + batch_size,
                            payload,
                        )
                        resp = rpc.put(
                            path="/v2/dataset/update-video-dataset-items/",
                            payload=payload,
                        )
                        logging.debug(
                            "response from update-video-dataset-items: %s",
                            resp,
                        )
                        if resp and resp.get("success"):
                            logging.debug(
                                "Successfully updated frames %s to %s",
                                i,
                                i + batch_size,
                            )
                            dataset_item["status"] = "processed"
                        else:
                            error_msg = resp.get("data") if resp else "No response from server"
                            logging.error(
                                "Failed to update frames %s to %s: %s",
                                i,
                                i + batch_size,
                                error_msg,
                            )
                            dataset_item["status"] = "errored"
                    except Exception as rpc_err:
                        logging.error(
                            "RPC call failed for frames %s to %s: %s",
                            i,
                            i + batch_size,
                            rpc_err,
                        )
                        dataset_item["status"] = "errored"
            return batch_image_details
        except Exception as e:
            logging.error("Error updating batch: %s", e)
            retry_count += 1
    for item in batch_image_details:
        item["status"] = "errored"
    return batch_image_details


def batch_update_video_imagenet_dataset_items(
    batch_image_details: List[Dict[str, Any]],
    rpc: Any,
    dataset_id: str,
    version: str,
    attempts: int = 3,
    is_yolo: bool = False,
    batch_size: int = 30,
) -> List[Dict[str, Any]]:
    """Update dataset items in batch, processing frames in groups of 30.

    Args:
        batch_image_details: List of image details to update
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        version: Version of the dataset
        attempts: Number of retry attempts
        is_yolo: Whether the dataset is in YOLO format
        batch_size: Number of frames to process in one API call

    Returns:
        List of updated dataset items
    """
    retry_count = 0
    logging.debug(
        "Batch image %s details are %s",
        dataset_id,
        batch_image_details,
    )
    while retry_count < attempts:
        try:
            logging.debug(
                "Attempting to update batch of %s items (attempt %s/%s)",
                len(batch_image_details),
                retry_count + 1,
                attempts,
            )
            batch_payload_items = []
            for dataset_item in batch_image_details:
                logging.debug(
                    "Processing dataset item for updating: %s",
                    dataset_item,
                )
                frame_annotations = []
                file_info = dataset_item.get("fileInfoResponse", {})
                segment_annotations = dataset_item.get("annotations", [])
                frame_annotations = [
                    {
                        "id": ann.get("id"),
                        "segment": dataset_item.get("segment_duration"),
                        "segmentation": ann.get("segmentation", []),
                        "isCrowd": ann.get("isCrowd", []),
                        "confidence": ann.get("confidence", 0.5),
                        "bbox": ann.get("bbox", []),
                        "height": ann.get("height"),
                        "width": ann.get("width"),
                        "center": ann.get("center", []),
                        "area": ann.get("area", 0),
                        "category": ann.get("category"),
                        "masks": ann.get("masks", []),
                    }
                    for ann in segment_annotations
                ]
                logging.debug(
                    "preview_cloud_path is- %s",
                    dataset_item.get("bucket_upload_first_frame_path"),
                )
                batch_payload_items.append(
                    {
                        "datasetItemId": str(file_info.get("_idVideoDatasetItem")),
                        "version": str(version),
                        "splitType": str(dataset_item.get("splitType")),
                        "previewFileCloudPath": dataset_item.get("bucket_upload_first_frame_path"),
                        "segmentWiseAnnotations": frame_annotations,
                        "height": dataset_item.get("video_height"),
                        "width": dataset_item.get("video_width"),
                        "area": int(dataset_item.get("video_height", 0))
                        * int(dataset_item.get("video_width", 0)),
                        "fps": dataset_item.get("frame_rate"),
                    }
                )
            payload = {
                "datasetId": str(dataset_id),
                "items": batch_payload_items,
            }
            try:
                logging.debug(
                    "Sending update for frames with payload: %s",
                    payload,
                )
                resp = rpc.put(
                    path="/v2/dataset/update-video-dataset-items/",
                    payload=payload,
                )
                logging.debug(
                    "Response from update-video-dataset-items: %s",
                    resp,
                )
                if resp and resp.get("success"):
                    logging.debug("Successfully updated dataset items!!")
                    for dataset_item in batch_image_details:
                        dataset_item["status"] = "processed"
                    return batch_image_details
                else:
                    error_msg = resp.get("data") if resp else "No response from server"
                    logging.error(
                        "Failed to update dataset items: %s",
                        error_msg,
                    )
            except Exception as rpc_err:
                logging.error(
                    "RPC call failed for dataset batch: %s",
                    rpc_err,
                )
        except Exception as e:
            logging.error("Error updating batch: %s", e)
        retry_count += 1
    for item in batch_image_details:
        item["status"] = "errored"
    return batch_image_details


def batch_update_kinetics_dataset_items(
    batch_image_details: List[Dict[str, Any]],
    rpc: Any,
    dataset_id: str,
    version: str,
    attempts: int = 3,
    is_yolo: bool = False,
    batch_size: int = 30,
) -> List[Dict[str, Any]]:
    """Update dataset items in batch, processing frames in groups of 30.

    Args:
        batch_image_details: List of image details to update
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        version: Version of the dataset
        attempts: Number of retry attempts
        is_yolo: Whether the dataset is in YOLO format
        batch_size: Number of frames to process in one API call

    Returns:
        List of updated dataset items
    """
    retry_count = 0
    logging.debug(
        "Batch image details are %s",
        batch_image_details,
    )
    while retry_count < attempts:
        try:
            logging.debug(
                "Attempting to update batch of %s items (attempt %s/%s)",
                len(batch_image_details),
                retry_count + 1,
                attempts,
            )
            batch_payload_items = []
            for dataset_item in batch_image_details:
                logging.debug(
                    "Processing dataset item for updating: %s",
                    dataset_item,
                )
                frame_annotations = []
                file_info = dataset_item.get("fileInfoResponse", {})
                segment_annotations = dataset_item.get("annotations", [])
                for ann in segment_annotations:
                    frame_annotations.append(
                        {
                            "id": ann.get("id"),
                            "segment": ann.get("duration"),
                            "segmentation": ann.get("segmentation", []),
                            "isCrowd": ann.get("isCrowd", []),
                            "confidence": ann.get("confidence", 0.5),
                            "bbox": ann.get("bbox", []),
                            "order_id": 0,
                            "height": ann.get("height"),
                            "width": ann.get("width"),
                            "center": ann.get("center", []),
                            "area": ann.get("area", 0),
                            "category": ann.get("category"),
                            "masks": ann.get("masks", []),
                        }
                    )
                logging.debug(
                    "preview_cloud_path is- %s",
                    dataset_item.get("bucket_upload_first_frame_path"),
                )
                batch_payload_items.append(
                    {
                        "datasetItemId": str(file_info.get("_idVideoDatasetItem")),
                        "previewFileCloudPath": dataset_item.get("bucket_upload_first_frame_path"),
                        "version": str(version),
                        "splitType": str(dataset_item.get("splitType")),
                        "segmentWiseAnnotations": frame_annotations,
                        "height": dataset_item.get("video_height"),
                        "width": dataset_item.get("video_width"),
                        "area": int(dataset_item.get("video_height", 0))
                        * int(dataset_item.get("video_width", 0)),
                        "fps": dataset_item.get("frame_rate"),
                    }
                )
            payload = {
                "datasetId": str(dataset_id),
                "items": batch_payload_items,
            }
            try:
                logging.debug(
                    "Sending update for frames with payload: %s",
                    payload,
                )
                resp = rpc.put(
                    path="/v2/dataset/update-video-dataset-items/",
                    payload=payload,
                )
                logging.debug(
                    "Response from update-video-dataset-items: %s",
                    resp,
                )
                if resp and resp.get("success"):
                    logging.debug("Successfully updated dataset items!!")
                    for dataset_item in batch_image_details:
                        dataset_item["status"] = "processed"
                    return batch_image_details
                else:
                    error_msg = resp.get("data") if resp else "No response from server"
                    logging.error(
                        "Failed to update dataset items: %s",
                        error_msg,
                    )
            except Exception as rpc_err:
                logging.error(
                    "RPC call failed for dataset batch: %s",
                    rpc_err,
                )
        except Exception as e:
            logging.error("Error updating batch: %s", e)
        retry_count += 1
    for item in batch_image_details:
        item["status"] = "errored"
    return batch_image_details


def batch_update_video_mscoco_dataset_items(
    batch_image_details: List[Dict[str, Any]],
    rpc: Any,
    dataset_id: str,
    version: str,
    attempts: int = 3,
    is_yolo: bool = False,
    batch_size: int = 30,
) -> List[Dict[str, Any]]:
    """Update dataset items in batch, processing frames in groups of 30.

    Args:
        batch_image_details: List of image details to update
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        version: Version of the dataset
        attempts: Number of retry attempts
        is_yolo: Whether the dataset is in YOLO format
        batch_size: Number of frames to process in one API call

    Returns:
        List of updated dataset items
    """
    retry_count = 0
    logging.debug(
        "Batch image details are %s",
        batch_image_details,
    )
    while retry_count < attempts:
        try:
            logging.debug(
                "Attempting to update batch of %s items (attempt %s/%s)",
                len(batch_image_details),
                retry_count + 1,
                attempts,
            )
            batch_payload_items = []
            for dataset_item in batch_image_details:
                logging.debug(
                    "Processing dataset item for updating: %s",
                    dataset_item,
                )
                frame_annotations = []
                file_info = dataset_item.get("fileInfoResponse", {})
                segment_annotations = dataset_item.get("annotations", [])
                for ann in segment_annotations:
                    frame_annotations.append(
                        {
                            "id": ann.get("id"),
                            "segment": ann.get("duration"),
                            "segmentation": ann.get("segmentation", []),
                            "isCrowd": ann.get("isCrowd", []),
                            "order_id": ann.get("order_id", 0),
                            "confidence": ann.get("confidence", 0.5),
                            "bbox": ann.get("bbox", []),
                            "height": ann.get("height"),
                            "width": ann.get("width"),
                            "center": ann.get("center", []),
                            "area": ann.get("area", 0),
                            "category": ann.get("category"),
                            "masks": ann.get("masks", []),
                        }
                    )
                logging.debug(f'Split Type is {dataset_item.get("splitType")}')

                batch_payload_items.append(
                    {
                        "datasetItemId": str(file_info.get("_idVideoDatasetItem")),
                        "previewFileCloudPath": "",
                        "version": str(version),
                        "splitType": str(dataset_item.get("splitType")),
                        "segmentWiseAnnotations": frame_annotations,
                        "height": dataset_item.get("video_height"),
                        "width": dataset_item.get("video_width"),
                        "area": int(dataset_item.get("video_height", 0))
                        * int(dataset_item.get("video_width", 0)),
                        "fps": dataset_item.get("frame_rate"),
                    }
                )

            # Send update API request for the batch
            payload = {
                "datasetId": str(dataset_id),
                "items": batch_payload_items,
            }
            try:
                logging.debug(
                    "Sending update for frames with payload: %s",
                    payload,
                )
                resp = rpc.put(
                    path="/v2/dataset/update-video-dataset-items/",
                    payload=payload,
                )
                logging.debug(
                    "Response from update-video-dataset-items: %s",
                    resp,
                )
                if resp and resp.get("success"):
                    logging.debug("Successfully updated dataset items!!")
                    for dataset_item in batch_image_details:
                        dataset_item["status"] = "processed"
                    return batch_image_details
                else:
                    error_msg = resp.get("data") if resp else "No response from server"
                    logging.error(
                        "Failed to update dataset items: %s",
                        error_msg,
                    )
            except Exception as rpc_err:
                logging.error(
                    "RPC call failed for dataset batch: %s",
                    rpc_err,
                )
        except Exception as e:
            logging.error("Error updating batch: %s", e)
        retry_count += 1
    for item in batch_image_details:
        item["status"] = "errored"
    return batch_image_details


def batch_update_dataset_items(
    batch_image_details: List[Dict[str, Any]],
    rpc: Any,
    dataset_id: str,
    version: str,
    attempts: int = 3,
    is_yolo: bool = False,
) -> List[Dict[str, Any]]:
    """Update dataset items in batch.

    Args:
        batch_image_details: List of dictionaries containing image details
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        version: Version of the dataset
        attempts: Number of retry attempts
        is_yolo: Whether using YOLO format

    Returns:
        Updated batch image details
    """
    logging.debug(
        "Processing batch of %s samples for update_dataset_items",
        len(batch_image_details),
    )
    retry_count = 0
    max_retries = attempts
    while retry_count < max_retries:
        try:
            payload = {
                "datasetId": str(dataset_id),
                "items": [
                    {
                        "datasetItemId": str(dataset_item.get("_id")),
                        "version": str(version),
                        "splitType": str(dataset_item.get("splitType")),
                        "annotations": dataset_item.get("annotations"),
                        "height": int(
                            dataset_item.get(
                                "image_height",
                                dataset_item.get("height"),
                            )
                        ),
                        "width": int(
                            dataset_item.get(
                                "image_width",
                                dataset_item.get("width"),
                            )
                        ),
                        "area": int(
                            dataset_item.get(
                                "image_area",
                                dataset_item.get("area"),
                            )
                        ),
                    }
                    for dataset_item in batch_image_details
                ],
            }
            if is_yolo:
                payload = convert_payload_to_coco_format(payload)
            response = rpc.put(
                path="/v2/dataset/update-dataset-items/",
                payload=payload,
            )
            logging.debug(
                "Update dataset items payload: %s",
                payload,
            )
            if response.get("success"):
                logging.debug(
                    "Successfully updated batch of %s items",
                    len(batch_image_details),
                )
                for item in batch_image_details:
                    item["status"] = "processed"
                return batch_image_details
            return batch_image_details
        except Exception as e:
            retry_count += 1
            logging.error(
                "Error updating items (attempt %d/%d): %s",
                retry_count,
                max_retries,
                str(e),
            )
            if retry_count >= max_retries:
                logging.error(
                    "Failed to update items after %d attempts",
                    max_retries,
                )
                return batch_image_details
    for item in batch_image_details:
        item["status"] = "errored"
    return batch_image_details


def submit_partition_status(
    dataset_items_batches: List[List[Dict[str, Any]]],
    rpc: Any,
    action_record_id: str,
    dataset_id: str,
    version: str,
    annotation_type: str,
) -> None:
    """Submit status of processed partition.

    Args:
        dataset_items_batches: List of batches of dataset items
        rpc: RPC client for making API calls
        action_record_id: ID of the action record
        dataset_id: ID of the dataset
        version: Version of the dataset
        annotation_type: Type of annotation
    """
    logging.info(
        "Submitting partition status for dataset %s version %s",
        dataset_id,
        version,
    )
    try:
        partitions_status = {}
        partition_items = {}
        for batch in dataset_items_batches:
            for item in batch:
                partition_num = item.get("partition")
                status = item.get("status")
                if partition_num not in partition_items:
                    partition_items[partition_num] = []
                partition_items[partition_num].append(item)
                if status == "errored":
                    partitions_status[partition_num] = "errored"
                    logging.error(
                        "Partition %s errored",
                        partition_num,
                    )
                elif status == "processed" and partition_num not in partitions_status:
                    partitions_status[partition_num] = "processed"
                    logging.info(
                        "Partition %s processed successfully",
                        partition_num,
                    )
        logging.info(
            "Updating status for %s partitions",
            len(partitions_status),
        )
        for (
            partition_num,
            status,
        ) in partitions_status.items():
            logging.debug(
                "Updating partition %s with status %s and %s items",
                partition_num,
                status,
                len(partition_items[partition_num]),
            )
            update_partition_status(
                rpc,
                action_record_id,
                dataset_id,
                version,
                partition_num,
                status,
                partition_items[partition_num],
                annotation_type,
            )
        logging.debug("Successfully updated partition status")
    except Exception as e:
        logging.error(
            "Error updating partition status: %s",
            e,
        )
        logging.debug(
            "Error details: %s",
            traceback.format_exc(),
        )


def submit_video_frame_partition_status(
    dataset_items_batches: List[List[Dict[str, Any]]],
    rpc: Any,
    action_record_id: str,
    dataset_id: str,
    version: str,
    annotation_type: str,
    sample_stats: Optional[Dict[str, Any]] = None,
) -> None:
    """Submit status updates for processed partitions.

    Args:
        dataset_items_batches: List of processed dataset item batches
        rpc: RPC client for making API calls
        action_record_id: ID of the action record
        dataset_id: ID of the dataset
        version: Version of the dataset
        annotation_type: Type of annotations
    """
    logging.info(
        "Submitting partition status for dataset %s version %s",
        dataset_id,
        version,
    )
    try:
        partitions_status = {}
        partition_items = {}
        for batch in dataset_items_batches:
            for item in batch:
                partition_num = item.get("partition")
                status = item.get("status")
                if partition_num not in partition_items:
                    partition_items[partition_num] = []
                partition_items[partition_num].append(item)
                if status == "errored":
                    partitions_status[partition_num] = "errored"
                    logging.error(
                        "Partition %s errored",
                        partition_num,
                    )
                elif status == "processed" and partition_num not in partitions_status:
                    partitions_status[partition_num] = "processed"
                    logging.info(
                        "Partition %s processed successfully",
                        partition_num,
                    )
        logging.info(
            "Updating status for %s partitions",
            len(partitions_status),
        )
        for (
            partition_num,
            status,
        ) in partitions_status.items():
            logging.debug(
                "Updating partition %s with status %s and %s items",
                partition_num,
                status,
                len(partition_items[partition_num]),
            )
            update_video_frame_partition_status(
                rpc,
                action_record_id,
                dataset_id,
                version,
                partition_num,
                status,
                partition_items[partition_num],
                annotation_type,
                sample_stats,
            )
            logging.info(
                "Successfully marked partition %s as %s",
                partition_num,
                status,
            )
    except Exception as e:
        logging.error(
            "Failed to submit partition status: %s",
            e,
        )
        logging.debug(
            "Full error traceback: %s",
            traceback.format_exc(),
        )


def get_annotation_files(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    is_annotations_compressed: bool = False,
) -> List[str]:
    """Download and return paths to annotation files.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        is_annotations_compressed: Whether annotations are in compressed format

    Returns:
        List of local paths to downloaded annotation files
    """
    logging.info(
        "Getting annotation files for dataset %s",
        dataset_id,
    )
    response = rpc_get_call(
        rpc,
        f"/v2/dataset/list-annotation-files/{dataset_id}/{dataset_version}",
        {},
    )
    annotation_files = []
    annotation_dir = os.path.join(TMP_FOLDER, "annotations")
    os.makedirs(annotation_dir, exist_ok=True)
    for s3_url in response:
        try:
            file_name = get_filename_from_url(s3_url)
            file_path = os.path.join(annotation_dir, file_name)
            os.makedirs(
                os.path.dirname(file_path),
                exist_ok=True,
            )
            download_file(s3_url, file_path)
            if is_annotations_compressed:
                annotation_files.extend(scan_folder(extract_dataset(file_path)))
            else:
                annotation_files.append(file_path)
            logging.debug(
                "Downloaded annotation file %s",
                annotation_files,
            )
        except Exception as e:
            logging.error(
                "Error downloading annotation file %s: %s",
                s3_url,
                e,
            )
    logging.info(
        "Found %s annotation files: %s",
        len(annotation_files),
        annotation_files,
    )
    return annotation_files


def get_mscoco_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        annotation_files = get_annotation_files(rpc, dataset_id, dataset_version)
        logging.info("Processing MSCOCO image details")
        images_details = get_msococo_images_details(annotation_files)
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": False,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_mscoco_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=download_images_queue,
            process_params={"images_details": images_details},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "detection",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error("Error setting up pipeline: %s", e)
        traceback.print_exc()
        raise


def get_imagenet_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": True,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_imagenet_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=download_images_queue,
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "classification",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error("Error setting up pipeline: %s", e)
        traceback.print_exc()
        raise


def get_pascalvoc_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        annotation_files = get_annotation_files(
            rpc,
            dataset_id,
            dataset_version,
            is_annotations_compressed=True,
        )
        logging.info("Processing Pascal image details")
        (
            images_details,
            missing_annotations,
            classwise_splits,
        ) = get_pascalvoc_image_details(annotation_files)
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": False,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_pascalvoc_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=download_images_queue,
            process_params={"images_details": images_details},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "detection",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error(
            "Error setting up Pascal VOC pipeline: %s",
            e,
        )
        traceback.print_exc()
        raise


def get_labelbox_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        annotation_files = get_annotation_files(rpc, dataset_id, dataset_version)
        logging.info("Processing Labelbox image details")
        (
            images_details,
            missing_annotations,
            classwise_splits,
        ) = get_labelbox_image_details(annotation_files)
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": False,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_labelbox_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=download_images_queue,
            process_params={"images_details": images_details},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Add Dataset Item Local File Path",
            process_fn=add_labelbox_dataset_item_local_file_path,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={"base_dataset_path": dataset_id},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "detection",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error("Error setting up pipeline: %s", e)
        traceback.print_exc()
        raise


def get_labelbox_classification_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        annotation_files = get_annotation_files(rpc, dataset_id, dataset_version)
        logging.info("Processing Labelbox image details")
        (
            images_details,
            missing_annotations,
            classwise_splits,
        ) = get_labelbox_classification_image_details(annotation_files)
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": False,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_labelbox_classification_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=download_images_queue,
            process_params={"images_details": images_details},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Add Dataset Item Local File Path",
            process_fn=add_labelbox_classification_dataset_item_local_file_path,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={"base_dataset_path": dataset_id},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "classification",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error("Error setting up pipeline: %s", e)
        traceback.print_exc()
        raise


def get_yolo_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
):
    """Create and configure the processing pipeline.
    Args:

        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        annotation_files = get_annotation_files(
            rpc,
            dataset_id,
            dataset_version,
            is_annotations_compressed=True,
        )
        logging.info("Processing Pascal image details")
        (
            images_details,
            missing_annotations,
            classwise_splits,
        ) = get_yolo_image_details(annotation_files)
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": False,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_yolo_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=download_images_queue,
            process_params={"images_details": images_details},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "is_yolo": True,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "detection",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error(
            "Error setting up Pascal VOC pipeline: %s",
            e,
        )
        traceback.print_exc()
        raise


def get_unlabelled_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": True,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_unlabelled_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=download_images_queue,
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "classification",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error("Error setting up pipeline: %s", e)
        traceback.print_exc()
        raise


def get_video_youtube_bb_tracking_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        annotation_files = get_annotation_files(
            rpc,
            dataset_id,
            dataset_version,
            is_annotations_compressed=False,
        )
        logging.debug(
            "Annotation files: %s",
            annotation_files,
        )
        logging.info("Processing Youtube BB Frames details")
        (
            images_details,
            missing_annotations,
            classwise_splits,
        ) = get_youtube_bb_video_frame_details(annotation_files)
        logging.debug(
            "Annotation details: %s",
            images_details,
        )
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=video_frame_partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": True,
                    "isFileInfoRequired": True,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_youtube_bb_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=download_images_queue,
            process_params={"frames_details": images_details},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_video_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_video_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_video_frame_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "detection",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error(
            "Error setting up Youtube BB pipeline: %s",
            e,
        )
        traceback.print_exc()
        raise


def get_video_mot_tracking_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        annotation_files = get_annotation_files(
            rpc,
            dataset_id,
            dataset_version,
            is_annotations_compressed=True,
        )
        logging.debug(
            "Annotation files: %s",
            annotation_files,
        )
        logging.info("Processing MOT Frames details")
        images_details, sample_stats = get_mot_annotations(annotation_files)
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=video_frame_partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": False,
                    "isFileInfoRequired": True,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_mot_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=download_images_queue,
            process_params={"frames_details": images_details},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_video_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_video_mot_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_video_frame_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "object_tracking",
                "sample_stats": sample_stats,
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error(
            "Error setting up Pascal VOC pipeline: %s",
            e,
        )
        traceback.print_exc()
        raise


def get_video_davis_segmentation_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        annotation_files = get_annotation_files(
            rpc,
            dataset_id,
            dataset_version,
            is_annotations_compressed=True,
        )
        logging.debug(
            "Annotation files: %s",
            annotation_files,
        )
        logging.info("Processing MOT Frames details")
        images_details, davis_sample_stats = get_davis_annotations(annotation_files)
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=video_frame_partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": False,
                    "isFileInfoRequired": True,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_davis_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=download_images_queue,
            process_params={"frames_details": images_details},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_video_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_video_davis_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_video_frame_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "segmentation",
                "sample_stats": davis_sample_stats,
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error(
            "Error setting up Pascal VOC pipeline: %s",
            e,
        )
        traceback.print_exc()
        raise


def get_video_imagenet_classification_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        upload_first_frame_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=video_frame_partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": True,
                    "isFileInfoRequired": True,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_video_imagenet_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=upload_first_frame_queue,
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Upload First frame",
            process_fn=batch_upload_video_samples,
            pull_queue=upload_first_frame_queue,
            push_queue=download_images_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_video_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_video_imagenet_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_video_frame_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "classification",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error(
            "Error setting up Pascal VOC pipeline: %s",
            e,
        )
        traceback.print_exc()
        raise


def get_kinetics_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        annotation_files = get_annotation_files(
            rpc,
            dataset_id,
            dataset_version,
            is_annotations_compressed=False,
        )
        logging.debug(
            "Annotation files: %s",
            annotation_files,
        )
        logging.info("Processing Kinetics Frames details")
        images_details = get_kinetics_annotations(annotation_files)
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        upload_first_frame_queue = Queue()
        download_images_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=video_frame_partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": True,
                    "isFileInfoRequired": True,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_kinetics_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=upload_first_frame_queue,
            process_params={"frames_details": images_details},
            num_threads=5,
        )
        pipeline.add_stage(
            stage_name="Upload First frame",
            process_fn=batch_upload_video_samples,
            pull_queue=upload_first_frame_queue,
            push_queue=download_images_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_video_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_kinetics_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_video_frame_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "kinetics",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error(
            "Error setting up Pascal VOC pipeline: %s",
            e,
        )
        traceback.print_exc()
        raise


def get_video_mscoco_server_processing_pipeline(
    rpc: Any,
    dataset_id: str,
    dataset_version: str,
    action_record_id: str,
    bucket_alias: str = "",
    account_number: str = "",
) -> Optional[Pipeline]:
    """Create and configure the processing pipeline.

    Args:
        rpc: RPC client for making API calls
        dataset_id: ID of the dataset
        dataset_version: Version number of the dataset
        action_record_id: ID of the action record

    Returns:
        Configured Pipeline instance
    """
    try:
        logging.info(
            "Setting up processing pipeline for dataset %s version %s",
            dataset_id,
            dataset_version,
        )
        annotation_files = get_annotation_files(
            rpc,
            dataset_id,
            dataset_version,
            is_annotations_compressed=False,
        )
        logging.debug(
            "Annotation files: %s",
            annotation_files,
        )
        logging.info("Processing Kinetics Frames details")
        images_details = get_video_mscoco_annotations(annotation_files)
        unprocessed_partitions = get_unprocessed_partitions(rpc, dataset_id, dataset_version)
        logging.info(
            "Found %s unprocessed partitions",
            len(unprocessed_partitions),
        )
        dataset_items_queue = Queue()
        download_images_queue = Queue()
        upload_first_frame_queue = Queue()
        calculate_image_properties_queue = Queue()
        update_dataset_items_queue = Queue()
        pipeline = Pipeline()
        for partition in unprocessed_partitions:
            pipeline.add_producer(
                process_fn=video_frame_partition_items_producer,
                process_params={
                    "rpc": rpc,
                    "dataset_id": dataset_id,
                    "partition": partition,
                    "pipeline_queue": dataset_items_queue,
                    "download_images_required": True,
                    "isFileInfoRequired": True,
                },
                partition_num=partition,
            )
        logging.info("Configuring pipeline stages")
        pipeline.add_stage(
            stage_name="Add Dataset Items Details",
            process_fn=add_video_mscoco_dataset_items_details,
            pull_queue=dataset_items_queue,
            push_queue=upload_first_frame_queue,
            process_params={"frames_details": images_details},
            num_threads=5,
        )

        pipeline.add_stage(
            stage_name="Upload First frame",
            process_fn=batch_upload_video_samples,
            pull_queue=upload_first_frame_queue,
            push_queue=download_images_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )

        pipeline.add_stage(
            stage_name="Download Images",
            process_fn=batch_download_video_samples,
            pull_queue=download_images_queue,
            push_queue=calculate_image_properties_queue,
            process_params={
                "rpc": rpc,
                "bucket_alias": bucket_alias,
                "account_number": account_number,
            },
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Calculate Image Properties",
            process_fn=batch_calculate_sample_properties,
            pull_queue=calculate_image_properties_queue,
            push_queue=update_dataset_items_queue,
            process_params={"properties_calculation_fn": calculate_image_properties},
            num_threads=10,
        )
        pipeline.add_stage(
            stage_name="Update Dataset Items",
            process_fn=batch_update_video_mscoco_dataset_items,
            pull_queue=update_dataset_items_queue,
            process_params={
                "rpc": rpc,
                "dataset_id": dataset_id,
                "version": dataset_version,
            },
            num_threads=10,
            is_last_stage=True,
        )
        pipeline.add_stop_callback(
            callback=submit_video_frame_partition_status,
            process_params={
                "rpc": rpc,
                "action_record_id": action_record_id,
                "dataset_id": dataset_id,
                "version": dataset_version,
                "annotation_type": "detection_mscoco",
            },
        )
        logging.info("Pipeline configuration complete")
        return pipeline
    except Exception as e:
        logging.error(
            "Error setting up Pascal VOC pipeline: %s",
            e,
        )
        traceback.print_exc()
        raise
