import base64
import requests
from PIL import Image
import io

def enhance_image_via_endpoint(
    image_path: str,
    endpoint_url: str,
    hf_token: str,
    prompt: str = "masterpiece, best quality, highres",
    upscale_factor: float = 2.0
):
    """
    Call the Hugging Face Inference Endpoint to enhance an image
    """
    
    # Load and encode image
    with open(image_path, "rb") as f:
        image_bytes = f.read()
    
    image_base64 = base64.b64encode(image_bytes).decode('utf-8')
    
    # Prepare request
    headers = {
        "Authorization": f"Bearer {hf_token}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "image": image_base64,
        "prompt": prompt,
        "upscale_factor": upscale_factor,
        "seed": 42
    }
    
    # Make request
    response = requests.post(endpoint_url, headers=headers, json=payload)
    
    if response.status_code == 200:
        result = response.json()
        
        if "error" in result:
            raise Exception(f"Enhancement failed: {result['error']}")
        
        # Decode enhanced image
        enhanced_base64 = result["enhanced_image"]
        enhanced_bytes = base64.b64decode(enhanced_base64)
        enhanced_image = Image.open(io.BytesIO(enhanced_bytes))
        
        print(f"Original size: {result.get('original_size')}")
        print(f"Enhanced size: {result.get('enhanced_size')}")
        
        return enhanced_image
    else:
        raise Exception(f"API call failed: {response.status_code} - {response.text}")

# Example usage
if __name__ == "__main__":
    # Replace with your values
    ENDPOINT_URL = "https://YOUR_ENDPOINT_ID.us-east-1.aws.endpoints.huggingface.cloud"
    HF_TOKEN = "hf_YOUR_TOKEN_HERE"
    
    try:
        enhanced = enhance_image_via_endpoint(
            image_path="input.jpg",
            endpoint_url=ENDPOINT_URL,
            hf_token=HF_TOKEN,
            prompt="masterpiece, best quality, highres, sharp details",
            upscale_factor=2.0
        )
        
        enhanced.save("enhanced_output.png")
        print("Image enhanced and saved!")
        
    except Exception as e:
        print(f"Error: {e}")