import os
import requests
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import io
import time

# Configuration
HF_API_TOKEN = os.environ.get("HF_API_TOKEN")
MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0"  # High-quality model
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"
headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
WATERMARK_TEXT = "SelamGPT"
MAX_RETRIES = 3

def add_watermark(image_bytes):
    """Add watermark to generated image"""
    try:
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        draw = ImageDraw.Draw(image)
        
        # Use default font (no need for external files)
        try:
            font = ImageFont.truetype("arialbd.ttf", 40)
        except:
            font = ImageFont.load_default(size=40)  # Fallback
        
        # Calculate text position (bottom-right corner)
        bbox = draw.textbbox((0, 0), WATERMARK_TEXT, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
        margin = 20
        position = (image.width - text_width - margin, image.height - text_height - margin)
        
        # Draw watermark with semi-transparent white text and black outline
        draw.text(
            position,
            WATERMARK_TEXT,
            font=font,
            fill=(255, 255, 255, 128),
            stroke_width=2,
            stroke_fill=(0, 0, 0, 128)
        )
        
        return image
    except Exception as e:
        print(f"Watermark error: {str(e)}")
        return Image.open(io.BytesIO(image_bytes))  # Return original if watermark fails

def generate_image(prompt):
    """Generate image with retry logic"""
    if not prompt.strip():
        return "Error: Please enter a valid prompt"
    
    for attempt in range(MAX_RETRIES):
        try:
            response = requests.post(
                API_URL,
                headers=headers,
                json={"inputs": prompt, "options": {"wait_for_model": True}},
                timeout=30
            )
            
            if response.status_code == 200:
                return add_watermark(response.content)
            elif response.status_code == 503:  # Model loading
                time.sleep(10 * (attempt + 1))  # Exponential backoff
                continue
            else:
                return f"API Error: {response.text}"
        except requests.Timeout:
            return "Error: Request timed out (30s)"
        except Exception as e:
            return f"Unexpected error: {str(e)}"
    
    return "Failed after multiple attempts. Try again later."

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎨 SelamGPT Image Generator
    *Powered by Stable Diffusion XL with built-in watermark*
    """)
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Describe your image",
                placeholder="A futuristic city at sunset...",
                lines=3
            )
            generate_btn = gr.Button("Generate", variant="primary")
            examples = gr.Examples(
                examples=["A cute robot reading a book", "Ethiopian landscape in oil painting style"],
                inputs=prompt_input
            )
        
        with gr.Column():
            output_image = gr.Image(label="Generated Image", height=512)
            error_output = gr.Textbox(label="Status", visible=False)
    
    generate_btn.click(
        fn=generate_image,
        inputs=prompt_input,
        outputs=[output_image, error_output],
        show_progress="full"
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)