import requests
import base64
import os
from PIL import Image
import io
import time
import datetime

def encode_image_to_base64(image_path):
    """Convert an image file to base64 string."""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def process_image(
    image_path, 
    mask_path=None, 
    prompt="", 
    height=1632, 
    width=1232, 
    guidance_scale=30, 
    num_inference_steps=50, 
    max_sequence_length=512,
    api_token=None,
    output_path="output_image.jpg"
):
    """
    Send a request to the Hugging Face Inference Endpoint.
    
    Args:
        image_path (str): Path to the input image
        mask_path (str, optional): Path to the mask image
        prompt (str): Text prompt to guide the model
        height (int): Output image height
        width (int): Output image width
        guidance_scale (float): Guidance scale for the model
        num_inference_steps (int): Number of inference steps
        max_sequence_length (int): Maximum sequence length
        api_token (str): Hugging Face API token
        output_path (str): Path to save the output image
    
    Returns:
        The response from the API or the path to the saved image
    """
    # Log start time
    start_time = time.time()
    start_timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
    print(f"Request started at: {start_timestamp}")
    
    # Endpoint URL
    url = "https://pewtn8bsankvfriv.us-east-1.aws.endpoints.huggingface.cloud"
    
    # Get API token (from environment variable or parameter)
    if api_token is None:
        api_token = os.environ.get("HF_API_TOKEN")
        if api_token is None:
            raise ValueError("API token not provided. Please set HF_API_TOKEN environment variable or pass it as a parameter.")
    
    # Encode image to base64
    base64_image = encode_image_to_base64(image_path)
    
    # Encode mask to base64 if provided
    base64_mask = None
    if mask_path:
        base64_mask = encode_image_to_base64(mask_path)
    
    # Prepare payload
    payload = {
        "inputs": {
            "image": base64_image,
            "prompt": prompt
        },
        "parameters": {
            "height": height,
            "width": width,
            "guidance_scale": guidance_scale,
            "num_inference_steps": num_inference_steps,
            "max_sequence_length": max_sequence_length,
        }
    }
    
    # Add mask if provided
    if base64_mask:
        payload["inputs"]["mask"] = base64_mask
    
    # Set up headers
    headers = {
        "Authorization": f"Bearer {api_token}",
        "Content-Type": "application/json"
    }
    
    # Make the request
    print(f"Sending request to {url}...")
    response = requests.post(url, json=payload, headers=headers)
    
    # Log end time
    end_time = time.time()
    end_timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
    elapsed_time = end_time - start_time
    print(f"Request completed at: {end_timestamp}")
    print(f"Total request time: {elapsed_time:.2f} seconds ({elapsed_time/60:.2f} minutes)")
    
    # Handle the response
    if response.status_code == 200:
        try:
            response_data = response.json()
            
            # Check if response has the expected format
            if "image" in response_data and "status" in response_data and response_data["status"] == "success":
                # Decode the base64 image
                image_data = base64.b64decode(response_data["image"])
                
                # Convert to PIL Image
                image = Image.open(io.BytesIO(image_data))
                
                # Save the image to the specified output path
                image.save(output_path)
                print(f"Image successfully saved to {output_path}")
                
                return output_path
            else:
                print("Unexpected response format:", response_data)
                return response_data
        except Exception as e:
            print(f"Error processing response: {e}")
            return response.json()
    else:
        print(f"Error: {response.status_code}")
        print(response.text)
        return None

if __name__ == "__main__":
    # Example usage
    output_file = "generated_image1.jpg"
    
    # 从环境变量获取API令牌，或者提示用户输入
    api_token = os.environ.get("HF_API_TOKEN")
    if not api_token:
        api_token = input("请输入您的Hugging Face API令牌: ")
    
    print("Starting image processing...")
    result = process_image(
        image_path="cup.png",
        mask_path="cup_mask.png",  # Optional
        prompt="a blue paper cup",
        api_token=api_token,
        output_path=output_file
    )
    
    if result == output_file:
        print(f"Processing completed successfully. Image saved to {output_file}")
    else:
        print("Processing completed with unexpected result:", result)