import gradio as gr
import requests
from PIL import Image
import io
import os 
BASE_URL = "https://api.jigsawstack.com/v1"

headers = {"x-api-key": os.getenv("JIGSAWSTACK_API_KEY")}


# ----------------- JigsawStack API Wrappers ------------------

def image_translate(source_type, image_url, file_store_key, target_language):
    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()