import requests
import base64
import json
from io import BytesIO
from typing import Optional, Dict

# Pillow is required for image format conversion and normalization.
# Please install it using: pip install Pillow
try:
    from PIL import Image
except ImportError:
    print("Pillow library not found. Please install it using: pip install Pillow")
    exit()

def _download_image_from_url(image_url: str, save_path: str) -> bool:
    """
    Downloads an image from a URL and saves it to a local path.

    Args:
        image_url: The URL of the image to download.
        save_path: The local path to save the downloaded image.

    Returns:
        True if the download was successful, False otherwise.
    """
    print(f"Downloading generated image from: {image_url}")
    try:
        image_response = requests.get(image_url, stream=True)
        # Check if the download request was successful.
        if image_response.status_code == 200:
            content_type = image_response.headers.get('Content-Type', '')
            if 'image' in content_type:
                with open(save_path, 'wb') as f:
                    for chunk in image_response.iter_content(1024):
                        f.write(chunk)
                print(f"Image successfully saved to {save_path}")
                return True
            else:
                print(f"Error: Content at URL is not an image. Content-Type: {content_type}")
                return False
        else:
            print(f"Error: Failed to download image. Status code: {image_response.status_code}")
            return False
    except requests.exceptions.RequestException as e:
        print(f"An error occurred during image download: {e}")
        return False

def generate_image(
    prompt_text: str,
    image_path: str,
    download_path: Optional[str] = None
) -> Optional[Dict]:
    """
    Sends a request to the image generation API and optionally downloads the result.

    Args:
        prompt_text: The instructional text for image modification.
        image_path: The file path to the input image (any common format).
        download_path: If provided, the path to save the generated image.

    Returns:
        A dictionary of the JSON response from the API, or None on error.
    """
    url = "https://kontext-chat.replicate.dev/generate-image"

    try:
        # --- Image Normalization Step ---
        with Image.open(image_path) as img:
            if img.mode != 'RGB':
                img = img.convert('RGB')
            with BytesIO() as output_buffer:
                img.save(
                    output_buffer,
                    format="JPEG",
                    quality=95,
                    subsampling=0,
                    progressive=False
                )
                image_bytes = output_buffer.getvalue()
    except FileNotFoundError:
        print(f"Error: Image file not found at {image_path}")
        return None
    except Exception as e:
        print(f"Error processing image file. Ensure it's a valid image. Details: {e}")
        return None

    encoded_string = base64.b64encode(image_bytes).decode('utf-8')
    input_image_data_uri = f"data:image/jpeg;base64,{encoded_string}"

    payload = {
        "prompt": prompt_text,
        "input_image": input_image_data_uri
    }

    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:139.0) Gecko/20100101 Firefox/139.0",
        "Accept": "*/*",
        "Content-Type": "application/json",
        "Referer": "https://kontext-chat.replicate.dev/",
        "Origin": "https://kontext-chat.replicate.dev",
    }

    try:
        response = requests.post(url, headers=headers, data=json.dumps(payload))
        response.raise_for_status()
        api_response_data = response.json()

        # --- Optional Download Logic ---
        if download_path and isinstance(api_response_data, dict):
            image_url = api_response_data.get("imageUrl")
            if image_url:
                _download_image_from_url(image_url, download_path)
            else:
                print("Warning: 'imageUrl' not found in response, could not download image.")

        return api_response_data
    except requests.exceptions.RequestException as e:
        print(f"API request failed: {e}")
        return None