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

def image_to_base64(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode()

def base64_to_image(base64_str):
    image_data = base64.b64decode(base64_str)
    return Image.open(io.BytesIO(image_data))

def save_results(generated_images, output_dir="results"):
    os.makedirs(output_dir, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    saved_paths = []
    for idx, img_base64 in enumerate(generated_images):
        img = base64_to_image(img_base64)
        output_path = os.path.join(output_dir, f"result_{timestamp}_{idx}.png")
        img.save(output_path)
        saved_paths.append(output_path)
        print(f"Saved image to {output_path}")
    
    return saved_paths

data = {
    "model_image": image_to_base64("imgs/person.jpg"),
    "garment_image": image_to_base64("imgs/garment.jpg"),
    "category": "Upper-body",
    "resolution": "768x1024",
    "n_steps": 30,
    "image_scale": 2.0,
    "num_images": 1
}

try:
    response = requests.post("http://localhost:8001/try-on", json=data)
    response.raise_for_status()
    result = response.json()
    
    if result["status"] == "success":
        saved_files = save_results(result["generated_images"])
        print(f"Successfully generated {len(saved_files)} images")
        print(f"Seed used: {result['seed']}")
    else:
        print("Generation failed:", result.get("detail", "Unknown error"))
        
except requests.exceptions.RequestException as e:
    print(f"Error making request: {e}")
except Exception as e:
    print(f"Error processing results: {e}")