import base64
from pathlib import Path
from typing import Union


class EncodedImage:
    """Handles base64-encoded images from LLM outputs."""

    def __init__(self, encoded_data: str):
        """
        Initialize with base64-encoded image data.

        Args:
            encoded_data: Base64-encoded string of the image
        """
        self.encoded_data = encoded_data.strip()

    def get_bytes(self) -> bytes:
        """
        Decode base64 string and return image as bytes.
        
        Returns:
            bytes: Image data as bytes (ready for Cloudinary upload)
        """
        return base64.b64decode(self.encoded_data)
    
    def get_base64(self) -> str:
        """
        Return base64 string (without data URI prefix).
        
        Returns:
            str: Base64-encoded image string
        """
        return self.encoded_data
    
    def get_data_uri(self, media_type: str = "image/jpeg") -> str:
        """
        Return data URI format (for direct use or display).
        
        Args:
            media_type: MIME type (default: image/jpeg)
            
        Returns:
            str: Data URI string (data:image/jpeg;base64,...)
        """
        return f"data:{media_type};base64,{self.encoded_data}"

    def save_as_jpg(self, path: str) -> None:
        """
        Save the encoded image as a JPEG file (kept for backward compatibility).

        Args:
            path: File path where the JPEG should be saved
        """
        image_bytes = self.get_bytes()
        
        # Ensure the path has .jpg or .jpeg extension
        output_path = Path(path)
        if output_path.suffix.lower() not in ['.jpg', '.jpeg']:
            output_path = output_path.with_suffix('.jpg')

        # Create parent directories if they don't exist
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # Write the decoded bytes to file
        with open(output_path, 'wb') as f:
            f.write(image_bytes)

