import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import arabic_reshaper
from bidi.algorithm import get_display 
import requests
from io import BytesIO
import os

# Your COLOR_MAP dictionary remains the same
COLOR_MAP = {
    "Black": (0, 0, 0),
    "White": (255, 255, 255),
    "Gray": (128, 128, 128),
    "Red": (255, 0, 0),
    "Blue": (0, 0, 255),
}

# --- REVISED AND FINAL FUNCTION ---
def overlay_text_on_image(persian_text, url, upload, username, text_color):
    # --- Stage 1: Image Loading (No changes) ---
    if url and url.strip():
        if not url.startswith(('http://', 'https://')):
            url = 'https://' + url
        try:
            response = requests.get(url)
            response.raise_for_status()
            img = Image.open(BytesIO(response.content))
        except Exception as e:
            raise ValueError(f"Failed to load image from URL: {e}")
    elif upload:
        img = Image.open(upload)
    else:
        raise ValueError("Please provide either an image URL or upload an image.")
    
    img = img.resize((1080, 1080), Image.LANCZOS)
    draw = ImageDraw.Draw(img)
    width, height = img.size
    
    # --- Stage 2: Font and Color Setup (No changes) ---
    font_path = "Vazir.ttf" 
    if not os.path.exists(font_path):
        print("Warning: Vazirmatn font not found. Falling back to default.")
        possible_fonts = [
            "Vazir.ttf",
            "/usr/share/fonts/truetype/noto/NotoNaskhArabic-Regular.ttf",
            "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"
        ]
        font_path = next((f for f in possible_fonts if os.path.exists(f)), "DejaVuSans.ttf")

    selected_color = COLOR_MAP.get(text_color, (0, 0, 0))
    lines = [line for line in persian_text.split('\n') if line.strip()]
    reshaped_lines = [arabic_reshaper.reshape(line) for line in lines]

    # --- Stage 3: UNIFORM FONT SIZE CALCULATION (No changes in logic) ---
    initial_font_size = 80
    max_width = width * 0.90
    uniform_font_size = initial_font_size
    
    longest_line = ""
    if reshaped_lines:
        temp_font = ImageFont.truetype(font_path, uniform_font_size)
        longest_line = max(reshaped_lines, key=lambda line: draw.textlength(line, font=temp_font))
    
    while uniform_font_size > 20:
        font = ImageFont.truetype(font_path, uniform_font_size)
        if draw.textlength(longest_line, font=font) <= max_width:
            break
        uniform_font_size -= 2
    
    final_font = ImageFont.truetype(font_path, uniform_font_size)

    # --- Stage 4: Text Block Height Calculation and Centering ---
    line_spacing = 20
    total_text_height = 0
    line_heights = []
    
    for line in reshaped_lines:
        # *** THE ONLY CHANGE IS ON THIS LINE ***
        # Replaced getbbox with textbbox for older Pillow compatibility
        line_bbox = draw.textbbox((0, 0), line, font=final_font)
        
        line_height = line_bbox[3] - line_bbox[1]
        line_heights.append(line_height)
        total_text_height += line_height

    total_text_height += (len(reshaped_lines) - 1) * line_spacing if len(reshaped_lines) > 1 else 0
    y_start = (height - total_text_height) / 2
    
    # --- Stage 5: Drawing the Text ---
    current_y = y_start
    for i, line in enumerate(reshaped_lines):
        x_center = width / 2
        line_y_center = current_y + line_heights[i] / 2
        
        draw.text(
            (x_center, line_y_center), 
            line, 
            font=final_font, 
            fill=selected_color, 
            anchor="mm"
        )
        
        current_y += line_heights[i] + line_spacing

    # --- Stage 6: Username Drawing (No changes) ---
    username_font_size = 40
    username_font = ImageFont.truetype(font_path, username_font_size)
    display_username = get_display(arabic_reshaper.reshape(f"@{username}"))
    
    username_x = width / 2
    username_y = height - 50
    draw.text((username_x, username_y), display_username, font=username_font, fill=selected_color, anchor="mb")
    
    return img

# --- Your Gradio Interface Code (No changes needed here) ---
with gr.Blocks() as demo:
    gr.Markdown("# Persian Quote Overlay for Instagram Posts")
    gr.Markdown("Paste an image URL from another HF Space or upload an image, enter Persian text (quote, supports multi-line), choose text color, and username. The app overlays the text centered with dynamic font sizing and places the username at the bottom center.")
    
    with gr.Row():
        persian_text = gr.Textbox(label="Persian Quote Text (Multi-line supported)", lines=5)
        image_input = gr.Textbox(label="Image URL (e.g., from another HF Space)")
        username = gr.Textbox(label="Channel Username (e.g., mychannel)")
    
    with gr.Row():
        text_color = gr.Dropdown(label="Text Color", choices=list(COLOR_MAP.keys()), value="Black")
        image_upload = gr.Image(label="Optional: Upload Image (if not using URL)", type="filepath")
    
    output_image = gr.Image(label="Output Image")
    
    submit_btn = gr.Button("Generate Overlay")
    submit_btn.click(overlay_text_on_image, 
                     inputs=[persian_text, image_input, image_upload, username, text_color], 
                     outputs=output_image)

demo.launch()