src/services/image_uploader.py

import requests
from typing import Optional, Union
import base64
from pathlib import Path
import os
from pydantic import BaseModel


class ImageInfo(BaseModel):
    filename: str
    name: str
    mime: str
    extension: str
    url: str


class ImgBBData(BaseModel):
    id: str
    title: str
    url_viewer: str
    url: str
    display_url: str
    width: int
    height: int
    size: int
    time: int
    expiration: int
    image: ImageInfo
    thumb: ImageInfo
    medium: ImageInfo
    delete_url: str


class ImgBBResponse(BaseModel):
    data: ImgBBData
    success: bool
    status: int


class ImageUploader:
    """A class to handle image uploads to ImgBB service."""

    def __init__(self, api_key: str):
        """
        Initialize the ImageUploader with an API key.

        Args:
            api_key (str): The ImgBB API key
        """
        self.api_key = api_key
        self.base_url = "https://api.imgbb.com/1/upload"

    def upload(
        self,
        image: Union[str, bytes, Path],
        name: Optional[str] = None,
        expiration: Optional[int] = None,
    ) -> ImgBBResponse:
        """
        Upload an image to ImgBB.

        Args:
            image: Can be:
                - A file path (str or Path)
                - Base64 encoded string
                - Base64 data URI (e.g., data:image/jpeg;base64,...)
                - URL to an image
                - Bytes of an image
            name: Optional name for the uploaded file
            expiration: Optional expiration time in seconds (60-15552000)

        Returns:
            ImgBBResponse containing the parsed upload response from ImgBB

        Raises:
            ValueError: If the image format is invalid or upload fails
            requests.RequestException: If the API request fails
        """
        # Prepare the parameters
        params = {"key": self.api_key}
        if expiration:
            if not 60 <= expiration <= 15552000:
                raise ValueError("Expiration must be between 60 and 15552000 seconds")
            params["expiration"] = expiration

        # Handle different image input types
        if isinstance(image, (str, Path)):
            image_str = str(image)
            files = {}
            if os.path.isfile(image_str):
                # It's a file path
                with open(image_str, "rb") as file:
                    files["image"] = file
            elif image_str.startswith(("http://", "https://")):
                # It's a URL
                files["image"] = (None, image_str)
            elif image_str.startswith("data:image/"):
                # It's a data URI
                # Extract the base64 part after the comma
                base64_data = image_str.split(",", 1)[1]
                files["image"] = (None, base64_data)
            else:
                # Assume it's base64 data
                files["image"] = (None, image_str)

            if name:
                files["name"] = (None, name)
            response = requests.post(self.base_url, params=params, files=files)
        elif isinstance(image, bytes):
            # Convert bytes to base64
            base64_image = base64.b64encode(image).decode("utf-8")
            files = {"image": (None, base64_image)}
            if name:
                files["name"] = (None, name)
            response = requests.post(self.base_url, params=params, files=files)
        else:
            raise ValueError(
                "Invalid image format. Must be file path, URL, base64 string, or bytes"
            )

        # Check the response
        if response.status_code != 200:
            raise ValueError(
                f"Upload failed with status {response.status_code}: {response.text}"
            )

        # Parse the response using Pydantic model
        response_json = response.json()
        return ImgBBResponse.parse_obj(response_json)

    def upload_file(
        self,
        file_path: Union[str, Path],
        name: Optional[str] = None,
        expiration: Optional[int] = None,
    ) -> ImgBBResponse:
        """
        Convenience method to upload an image file.

        Args:
            file_path: Path to the image file
            name: Optional name for the uploaded file
            expiration: Optional expiration time in seconds (60-15552000)

        Returns:
            ImgBBResponse containing the parsed upload response from ImgBB
        """
        return self.upload(file_path, name=name, expiration=expiration)

    def upload_url(
        self,
        image_url: str,
        name: Optional[str] = None,
        expiration: Optional[int] = None,
    ) -> ImgBBResponse:
        """
        Convenience method to upload an image from a URL.

        Args:
            image_url: URL of the image to upload
            name: Optional name for the uploaded file
            expiration: Optional expiration time in seconds (60-15552000)

        Returns:
            ImgBBResponse containing the parsed upload response from ImgBB
        """
        return self.upload(image_url, name=name, expiration=expiration)