import asyncio
import os
import httpx
import base64
import json
from dotenv import load_dotenv
from typing import Optional
import uuid

load_dotenv()
API_KEY = os.getenv("GOOGLE_API_KEY")

async def _upload_file_http(media_data: bytes, mime_type: str) -> Optional[str]:
    """Helper for testing upload"""
    url = f"https://generativelanguage.googleapis.com/upload/v1beta/files?key={API_KEY}"
    num_bytes = len(media_data)
    
    headers = {
        "X-Goog-Upload-Protocol": "resumable",
        "X-Goog-Upload-Command": "start",
        "X-Goog-Upload-Header-Content-Length": str(num_bytes),
        "X-Goog-Upload-Header-Content-Type": mime_type,
        "Content-Type": "application/json"
    }
    
    metadata = {"file": {"display_name": f"debug_upload_{uuid.uuid4()}"}}
    
    async with httpx.AsyncClient(timeout=30.0) as client:
        # 1. Start Upload
        print("Starting resumable upload...")
        response = await client.post(url, headers=headers, json=metadata)
        if response.status_code != 200:
            print(f"Failed to start upload: {response.text}")
            return None
            
        upload_url = response.headers.get("x-goog-upload-url")
        if not upload_url:
            print("No upload URL returned")
            return None
            
        # 2. Upload Bytes
        print(f"Uploading {num_bytes} bytes to {upload_url[:50]}...")
        headers_upload = {
            "Content-Length": str(num_bytes),
            "X-Goog-Upload-Offset": "0",
            "X-Goog-Upload-Command": "upload, finalize"
        }
        
        response_upload = await client.post(upload_url, headers=headers_upload, content=media_data)
        if response_upload.status_code != 200:
            print(f"Failed to upload data: {response_upload.text}")
            return None
            
        data = response_upload.json()
        file_uri = data.get("file", {}).get("uri")
        print(f"File uploaded successfully: {file_uri}")
        return file_uri

async def test_edit():
    # Use gemini-3-pro-image-preview
    model_name = "gemini-3-pro-image-preview" 
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent?key={API_KEY}"

    print(f"Testing Model: {model_name} with Raw HTTP File API")

    # Generate valid JPEG using PIL
    from PIL import Image
    import io
    
    img = Image.new('RGB', (100, 100), color = 'red')
    buf = io.BytesIO()
    img.save(buf, format='JPEG')
    pixel_data = buf.getvalue()
    
    print(f"Generated valid JPEG ({len(pixel_data)} bytes)")
    
    # Upload File
    file_uri = await _upload_file_http(pixel_data, "image/jpeg")
    if not file_uri:
        print("Skipping generation due to upload failure")
        return

    prompt = "Make the image blue"
    final_prompt = f"Edit this image: {prompt}"

    parts = []
    
    # 1. File Data Part
    parts.append({
        "fileData": {
            "mimeType": "image/jpeg",
            "fileUri": file_uri
        }
    })

    parts.append({"text": final_prompt})

    payload = {
        "contents": [{
            "parts": parts
        }]
    }
    
    async with httpx.AsyncClient(timeout=30) as client:
        print("Generating content...")
        response = await client.post(
            url,
            headers={"Content-Type": "application/json"},
            json=payload
        )

        print(f"Status: {response.status_code}")
        # print(response.text)
        
        if response.status_code == 200:
            print("✓ Success! API accepted the request.")
            data = response.json()
            if "candidates" in data and data["candidates"]:
                print("Candidates found.")
        else:
             print(f"✗ Failed: {response.text}")

if __name__ == "__main__":
    asyncio.run(test_edit())