import json
import logging
import os
import time
from pathlib import Path
from typing import Callable, TypeVar, Union
from uuid import uuid4
from zipfile import BadZipFile, ZipFile

import requests
from directory_tree import display_tree
from spb_label.exceptions import (
    APIException,
    BadRequestException,
    ConflictException,
    # InvaildObjectException,
)
from phy_credit.exceptions import InvalidObjectException

logger = logging.getLogger("superb_label")


def separate_batches(image_batch_size: int, to_batch_list: list):
    """
    Separates a list of items into batches of a specified size.

    Args:
        image_batch_size (int): The size of each batch.
        to_batch_list (list): The list of items to be batched.

    Returns:
        list: A list of batches, where each batch is a sublist of the original list.
    """
    image_paths = []
    number_of_iteration = len(to_batch_list) // image_batch_size + 1
    for i in range(number_of_iteration):
        start = i * image_batch_size
        end = (
            (i + 1) * image_batch_size
            if (i + 1) * image_batch_size < len(to_batch_list)
            else False
        )
        if end is False:
            subset_image_paths = to_batch_list[start:]
        else:
            subset_image_paths = to_batch_list[start:end]
        image_paths.append(subset_image_paths)

    return image_paths


def read_info_from_zip_yolo(
    dataset_file: str, require_images_folder: bool = True
):
    """
    Reads and extracts information from a YOLO formatted zip file.

    Args:
        dataset_file (str): The path to the dataset zip file.
        require_images_folder (bool, optional): Flag to require an 'images' folder within the zip file. Defaults to True.

    Returns:
        tuple: A tuple containing lists of image paths, annotation paths, and the path to the classes.txt file.
    """
    try:
        with ZipFile(dataset_file, "r") as zip_file:
            file_list = [
                f
                for f in zip_file.namelist()
                if "__MACOSX" not in f and "DS_Store" not in f
            ]
            top_level_folder = os.path.commonprefix(file_list)
            second_level_folders, classes_txt = extract_folders_and_classes(
                file_list, top_level_folder
            )
            image_path_list = extract_paths(
                file_list,
                ["jpg", "jpeg", "png", "bmp"],
                "images",
                require_images_folder,
            )
            annotation_path_list = extract_paths(
                file_list, ["txt"], "labels", True
            )

            validate_contents(
                second_level_folders,
                classes_txt,
                image_path_list,
                annotation_path_list,
                require_images_folder,
            )
            zip_file.extractall(path=Path(dataset_file).parent)

        return image_path_list, annotation_path_list, classes_txt

    except BadZipFile as e:
        raise Exception(
            f"[Invalid zip file] '{dataset_file.split('/')[-1]}' {e}."
        )
    except Exception as e:
        display_error_and_extract(dataset_file, e)


def extract_folders_and_classes(file_list, top_level_folder):
    """
    Extracts second level folders and the classes.txt file from a list of file paths.

    Args:
        file_list (list): The list of file paths within the zip file.
        top_level_folder (str): The common prefix of all file paths, representing the top level folder.

    Returns:
        tuple: A tuple containing a set of second level folder names and the path to the classes.txt file.
    """
    second_level_folders = set()
    classes_txt = None
    for item in file_list:
        split_path = item.split("/")
        if top_level_folder and split_path[-1] == "classes.txt":
            classes_txt = item
        elif not top_level_folder and item.endswith("classes.txt"):
            classes_txt = item
        elif (top_level_folder and len(split_path) == 3) or (
            not top_level_folder and len(split_path) == 2
        ):
            second_level_folders.add(split_path[-2])
    return second_level_folders, classes_txt


def extract_paths(file_list, extensions, folder_name, required):
    """
    Extracts file paths from a list based on file extensions and folder name requirements.

    Args:
        file_list (list): The list of file paths within the zip file.
        extensions (list): A list of file extensions to filter by.
        folder_name (str): The name of the folder to filter by.
        required (bool): Flag to require the specified folder name in the file path.

    Returns:
        list: A list of file paths that match the specified extensions and folder name requirements.
    """
    return [
        f
        for f in file_list
        if f.split("/")[-1].split(".")[-1].lower() in extensions
        and (not required or f.split("/")[-2] == folder_name)
    ]


def validate_contents(
    second_level_folders,
    classes_txt,
    image_path_list,
    annotation_path_list,
    require_images_folder,
):
    """
    Validates the contents of the zip file against expected structure and files.

    Args:
        second_level_folders (set): A set of second level folder names extracted from the zip file.
        classes_txt (str): The path to the classes.txt file.
        image_path_list (list): A list of image file paths.
        annotation_path_list (list): A list of annotation file paths.
        require_images_folder (bool): Flag to require an 'images' folder within the zip file.

    Raises:
        Exception: If the validation fails due to missing required folders or files.
    """
    expected_folders = (
        {"labels", "images"} if require_images_folder else {"labels"}
    )
    if second_level_folders != expected_folders or not classes_txt:
        missing = expected_folders - second_level_folders
        missing_info = (
            f"Required folder(s), {', '.join(missing)}, are missing."
            if missing
            else "Required 'classes.txt' is missing."
        )
        raise Exception(missing_info)
    if require_images_folder and not image_path_list:
        raise Exception(
            "The 'images' folder contains no image files with valid formats."
        )
    if not annotation_path_list:
        raise Exception(
            "The 'labels' folder contains no label files with valid formats."
        )


def display_error_and_extract(dataset_file, e):
    """
    Displays an error message and extracts the contents of the zip file for inspection.

    Args:
        dataset_file (str): The path to the dataset zip file.
        e (Exception): The exception that was raised.

    Raises:
        Exception: Re-raises the exception after displaying the error message and extracting the zip file.
    """
    print("[Zip File]")
    with ZipFile(dataset_file, "r") as zip_file:
        zip_file.extractall(path=Path(dataset_file).parent)
        display_tree(str(Path(dataset_file).parent), max_depth=2)
    raise Exception(f"[Invalid zip file] {e}")


T = TypeVar("T")


def call_with_retry(
    fn: Callable[..., T],
    *args,
    _max_retries: int = 5,
    _sleep_time: int = 6,
    **kwargs,
) -> Union[T, None]:
    retry_count = 0
    while retry_count < _max_retries:
        if retry_count == 0:
            log_id = str(uuid4()).split("-")[0]
        try:
            res = fn(*args, **kwargs)
            break
        except requests.HTTPError as e:
            error_code = e.response.status_code
            if error_code == 503:  # 503: Service Unavailable
                logger.warning(
                    f"({log_id}) 503 Error occured, will be retried || fn: {fn.__name__} || Retry count: {retry_count}"
                )
            retry_count += 1
            time.sleep(_sleep_time)
        except APIException as e:
            logger.warning(
                f"({log_id}) 429 Error occured, will be retried || fn: {fn.__name__} || Retry count: {retry_count}"
            )
            retry_count += 1
            time.sleep(_sleep_time)
        except InvalidObjectException as e:
            raise Exception(
                f"({log_id}) Class name is not vaild, check the class list || fn: {fn.__name__} || Retry count: {retry_count}"
            )
        except Exception as e:

            if "429" in e.message:  # 429: too many request
                logger.error(
                    f"({log_id}) 429 Error occured, will be retried || fn: {fn.__name__} || Retry count: {retry_count}"
                )
                time.sleep(_sleep_time)
                retry_count += 1
            elif "503" in e.message:  # 503: too many request
                logger.error(
                    f"({log_id}) 503 Error occured, will be retried || fn: {fn.__name__} || Retry count: {retry_count}"
                )
                time.sleep(_sleep_time)
                retry_count += 1
            else:
                raise Exception(
                    f"({log_id}) Unknown error occured, will be retried || fn: {fn.__name__} || Retry count: {retry_count}"
                )
        if retry_count > 0:
            logger.debug(
                f"({log_id}) Retrying || fn: {fn.__name__} || Retry count: {retry_count}"
            )
    if retry_count == _max_retries:
        raise Exception(
            f"({log_id}) Max retries reached || fn: {fn.__name__} || Retry count: {retry_count}"
        )
    if retry_count > 0:
        logger.debug(
            f"({log_id}) Successfully completed || fn: {fn.__name__} || Retry count: {retry_count}"
        )

    return res
