import base64
import os
import re
import shutil
import zipfile
from io import BytesIO
from pathlib import Path

import requests
from PIL import Image


def encode_img_to_base64(image_path: str) -> str:
    """
    Encode an image file to a base64 string.

    Args:
        image_path (str): The path to the image file.

    Returns:
        str: A base64 encoded string representation of the image.
    """
    try:
        image = Image.open(image_path)
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        base64_string = base64.b64encode(buffered.getvalue()).decode("utf-8")
        return f"data:image/png;base64,{base64_string}"
    except FileNotFoundError:
        raise FileNotFoundError(f"Image file not found: {image_path}")
    except IOError as e:
        raise IOError(f"Error reading BMP image file: {e}")
    except Exception as e:
        raise ValueError(f"Error processing the image file into Base64 format: {e}")


def separate_options(options_string: str) -> list[str]:
    """
    Separates a string of options into a list of individual options.

    Args:
        options_string (str): A string containing multiple options, where each option starts with a letter followed by a colon.

    Returns:
        list[str]: A list of individual options with the letter prefix removed and whitespace stripped.
    """
    # Ensure we're working with normalized line endings
    options_string = options_string.strip()

    # If no content, return empty list
    if not options_string:
        return []

    # Pattern to match option boundaries (lines starting with a letter and colon)
    pattern = re.compile(r"(?:^|\n)([A-Za-z]:)", re.MULTILINE)

    # Split the string at option boundaries
    parts = pattern.split(options_string)

    # First part might be text before any option marker
    if not re.match(r"^[A-Za-z]:$", parts[0]):
        parts = parts[1:]

    # Process the parts in pairs (marker + content)
    options = []
    for i in range(0, len(parts), 2):
        if i + 1 < len(parts):
            # Clean up the content
            content = parts[i + 1].strip()
            content = "\n".join(
                line.strip() for line in content.split("\n") if line.strip()
            )
            if content:
                options.append(content)

    # If no options found, fall back to simple newline splitting
    if not options:
        return [line.strip() for line in options_string.split("\n") if line.strip()]

    return options


DATA_FILES = [
    "Low-level-semantic",
    "High-level-sub-semantic",
    "High-level-obj-semantic",
    "Discrete-temporal",
    "Continuous-temporal",
    "2D-spatial",
    "3D-spatial",
]


def ensure_data(
    data_dir: Path, force_download: bool, data_files: list[str] = DATA_FILES
) -> None:
    """
    Ensure that the required data files are downloaded and extracted.

    Args:
        data_dir (Path): Directory to store the downloaded and extracted data.
        force_download (bool): If True, force re-download even if files exist.
        data_files (list[str]): List of data file names to download and extract.

    Returns:
        None
    """
    for file in data_files:
        data_url = f"https://huggingface.co/datasets/FanqingM/MMIU-Benchmark/resolve/main/{file}.zip?download=true"
        extracted_dir = data_dir / file
        zip_path = extracted_dir / f"{file}.zip"

        if extracted_dir.is_dir() and force_download:
            shutil.rmtree(extracted_dir)

        # make sure data_dir exists
        extracted_dir.parent.mkdir(parents=True, exist_ok=True)

        if not extracted_dir.is_dir():
            if not zip_path.exists():
                download_large_file(data_url, zip_path)
            unzip_file(zip_path, extracted_dir)
            zip_path.unlink()  # Remove the zip file after extraction


def download_large_file(
    url: str, output_path: Path, chunk_size: int = 1024 * 1024
) -> None:
    """Downloads a large file in chunks and saves it to disk."""
    print(f"Downloading {url} to {output_path}")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with requests.get(url, stream=True) as response:
        response.raise_for_status()  # Raise an error for failed requests
        with open(output_path, "wb") as file:
            for chunk in response.iter_content(chunk_size=chunk_size):
                if chunk:  # Ensure chunk is not empty
                    file.write(chunk)


def unzip_file(zip_path: Path, extract_to: Path) -> None:
    """
    Unzip a file to a specified directory.

    Args:
        zip_path (Path): Path to the zip file to be extracted.
        extract_to (Path): Directory where the contents will be extracted.

    Returns:
        None
    """
    print(f"Unzipping {zip_path}")

    # Create the target directory if it doesn't exist
    os.makedirs(extract_to, exist_ok=True)

    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        zip_ref.extractall(extract_to)

    print(f"Extracted to {extract_to}")
