import gradio as gr
import requests
import json
import os
import time
from collections import defaultdict

BASE_URL = "https://api.jigsawstack.com/v1"
headers = {
    "x-api-key":  os.getenv("JIGSAWSTACK_API_KEY")
}


# Rate limiting configuration
request_times = defaultdict(list)
MAX_REQUESTS = 20  # Maximum requests per time window
TIME_WINDOW = 3600   # Time window in seconds (1 hour)

def get_real_ip(request: gr.Request):
    """Extract real IP address using x-forwarded-for header or fallback"""
    if not request:
        return "unknown"
    
    forwarded = request.headers.get("x-forwarded-for")
    if forwarded:
        ip = forwarded.split(",")[0].strip()  # First IP in the list is the client's
    else:
        ip = request.client.host  # fallback
    return ip

def check_rate_limit(request: gr.Request):
    """Check if the current request exceeds rate limits"""
    if not request:
        return True, "Rate limit check failed - no request info"
    
    ip = get_real_ip(request)
    now = time.time()

    # Clean up old timestamps outside the time window
    request_times[ip] = [t for t in request_times[ip] if now - t < TIME_WINDOW]
    

    # Check if rate limit exceeded
    if len(request_times[ip]) >= MAX_REQUESTS:
        time_remaining = int(TIME_WINDOW - (now - request_times[ip][0]))
        time_remaining_minutes = round(time_remaining / 60, 1)      
        time_window_minutes = round(TIME_WINDOW / 60, 1)
        
        return False, f"Rate limit exceeded. You can make {MAX_REQUESTS} requests per {time_window_minutes} minutes. Try again in {time_remaining_minutes} minutes."
    
    # Add current request timestamp
    request_times[ip].append(now)
    return True, ""

    
def image_translate(source_type, image_url, file_store_key, target_language, request: gr.Request):
    rate_limit_ok, rate_limit_msg = check_rate_limit(request)
    if not rate_limit_ok:
        return rate_limit_msg, None
    
    if not target_language or not target_language.strip():
        return "Error: Target language is required.", None

    payload = {"target_language": target_language.strip()}
    
    if source_type == "URL":
        if not image_url or not image_url.strip():
            return "Error: Image URL is required.", None
        payload["url"] = image_url.strip()
    elif source_type == "File Store Key":
        if not file_store_key or not file_store_key.strip():
            return "Error: File Store Key is required.", None
        payload["file_store_key"] = file_store_key.strip()
    else:
        return "Error: Invalid image source selected.", None

    try:
        response = requests.post(
            f"{BASE_URL}/ai/translate/image",
            headers=headers,
            json=payload
        )
        response.raise_for_status()
        
        if response.headers.get("content-type", "").startswith("image/"):
            # The API returns the translated image directly
            from PIL import Image
            import io
            image = Image.open(io.BytesIO(response.content))
            return "✅ Image translated successfully.", image
        else:
            try:
                error_data = response.json()
                return f"Error: API returned an error - {error_data.get('message', 'Unknown error')}", None
            except:
                return "Error: Received an unexpected response from the API.", None

    except requests.exceptions.RequestException as e:
        return f"Request failed: {str(e)}", None
    except Exception as e:
        return f"An unexpected error occurred: {str(e)}", None

# ----------------- Gradio UI ------------------

with gr.Blocks() as demo:
    gr.Markdown("""
    <div style='text-align: center; margin-bottom: 24px;'>
        <h1 style='font-size:2.2em; margin-bottom: 0.2em;'>🧩 Image Translation</h1>
        <p style='font-size:1.2em; margin-top: 0;'>Extract and translate text from images into multiple languages.</p>
        <p style='font-size:1em; margin-top: 0.5em;'>For more details and API usage, see the <a href='https://jigsawstack.com/docs/api-reference/ai/translate/image-translate' target='_blank'>documentation</a>.</p>
    </div>
    """)

    with gr.Row():
        with gr.Column():
            gr.Markdown("#### Image Source")
            it_source_type = gr.Radio(
                choices=["URL", "File Store Key"],
                label="Choose Image Source",
                value="URL"
            )
            it_image_url = gr.Textbox(
                label="Image URL",
                placeholder="Enter the URL of the image..."
            )
            it_file_key = gr.Textbox(
                label="File Store Key",
                placeholder="your-file-store-key",
                visible=False
            )

            gr.Markdown("#### Translation Options")
            it_target_language = gr.Textbox(
                label="Target Language Code",
                placeholder="es, fr, de, ja, etc."
            )
            it_btn = gr.Button("Translate Image", variant="primary")

        with gr.Column():
            gr.Markdown("#### Results")
            it_status = gr.Textbox(label="Status", interactive=False)
            it_translated_image = gr.Image(label="Translated Image")

    def update_it_source(source_type):
        is_url = source_type == "URL"
        return gr.update(visible=is_url), gr.update(visible=not is_url)

    it_source_type.change(
        update_it_source,
        inputs=it_source_type,
        outputs=[it_image_url, it_file_key]
    )

    it_btn.click(
        image_translate,
        inputs=[it_source_type, it_image_url, it_file_key, it_target_language],
        outputs=[it_status, it_translated_image]
    )

demo.launch()