import gradio as gr
import google.generativeai as genai
from PIL import Image
import tempfile

LANGUAGE_PROMPTS = {
    "Urdu": {
        "hint": "Extract all visible Urdu text from this image and return it in plain Urdu text."
    },
    "Saraiki": {
        "hint": (
            "The image may contain Saraiki text. "
            "Look for special characters like 'ٻ', 'ڄ', 'ݙ', 'ڳ', 'ݨ'. "
            "Extract all visible text and return it in plain Saraiki text."
        )
    },
}

def extract_text_from_uploaded_image(image, api_key, language):
    if not api_key:
        raise gr.Error("⚠️ Please enter your Google AI Studio API key.")
    if image is None:
        raise gr.Error("⚠️ Please upload an image (jpg, jpeg, png, webp, etc.).")

    # Configure Gemini safely
    genai.configure(api_key=api_key)
    model = genai.GenerativeModel('gemini-2.0-flash')
    prompt = LANGUAGE_PROMPTS[language]["hint"]

    try:
        # Save + reload to ensure PIL compatibility
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
            image.save(tmp.name)
            img = Image.open(tmp.name)

        response = model.generate_content([img, prompt])

        # Gemini block check
        if response.prompt_feedback and response.prompt_feedback.block_reason:
            raise gr.Error(f"❌ Blocked by Gemini: {response.prompt_feedback.block_reason}")

        text = response.text.strip() or "❌ No text found"
        # return text + clear the key field
        return text
    except gr.Error:
        # re-raise Gradio errors to show popups
        raise
    except Exception as e:
        raise gr.Error(f"⚠️ Unexpected error: {e}")

with gr.Blocks(title="HarfSaaz – Urdu & Saraiki OCR") as demo:
    # 1) inject CSS via an HTML component
    gr.HTML("""
      <style>
        .fixed-box1 { 
          border: 3px solid #ccc; 
          width: 450px; 
          height: 800px; 
          padding: 8px; 
          box-sizing: border-box;
        }
        .fixed-box2 { 
          border: 3px solid #ccc; 
          width: 400px; 
          height: 400px; 
          padding: 8px; 
          box-sizing: border-box;
        }
        .tall-orange-btn {
          height: 100%;               /* Match container height */
          background-color: orange;
          color: white;
          font-weight: bold;
          font-size: 16px;
          border: none;
          border-radius: 8px;
          width: 100%;
        }
          .tall-orange-btn:hover {
          background-color: darkorange;
        }
      </style>
    """)

    # 2) API key + language row
    with gr.Row():
        with gr.Column(scale=4):  # Adjust scale to control width
            with gr.Row():
                api_key_input = gr.Textbox(
                    label="API Key", type="password", placeholder="AIzaSy..."
                )
                lang_dropdown = gr.Dropdown(
                    choices=list(LANGUAGE_PROMPTS.keys()),
                    value="Saraiki",
                    label="Language"
                )
        with gr.Column(scale=1):
            submit_btn = gr.Button("Run HarfSaazi", elem_classes="tall-orange-btn")


    # 3) Upload & result side-by-side
    gr.Markdown("## 📤 Upload & Results")
    with gr.Row():
        image_input = gr.Image(
            type="pil",
            label="Upload Image",
            elem_classes="fixed-box1"
        )
        result_output = gr.Textbox(
            interactive=False,
            label="HarfSaaz Result:",
            elem_classes="fixed-box1"
        )

    submit_btn.click(
        extract_text_from_uploaded_image,
        inputs=[image_input, api_key_input, lang_dropdown],
        outputs=[result_output]
    )

    # 4) Static examples below
    gr.Markdown("## 🎓 Examples")
    with gr.Row():
        examples = [
            ("ex1.png", """ملھ مہانگ دے ایس زما نے اچ ہکا شے مفت پی بلدی اے
تے او ہے صلاح۔ تساں مکان بنڑا نوراں ہو وے یا کھوہ کھٹا نوڑاں
ہووے۔ وپار اچ سیر کارڑاں ہو وے یا ست بسم اللہ آکھ تے وسدے
ہمسائے نے دعوئی ٹورٹا ہووے۔ بھنڑ بچے دا عقیقہ کرنا ہووے یا ہال دی
شادی راچ ہتھ پانوڑاں ہووے۔ صلاحیں ڈیوٹی آلے ہر موقعے نے کھکھر
آلی کار ویڑھ ویندن۔ سیالے دے کونے آنگوں چنبڑ ویندن۔ اسیڈی
ایها ذری بھوری عقل جیرهی ریمکدی و دی اے اونکوں حریان تے پریشان
کر ڈیندن ۔ لالہ ! چیتا گردان کر ڈیندن۔ صلاحیں دے ہوا جھکوڑے اچ بندہ
کھلے آنگوں موندھا تھی ویندے۔ کڑا ہوں ایڈے وقبل دے، کیڑا ہوں
اوڑے پھنک دے۔ بہہ بہہ . تے ٹھڈے شکارے بھریندے تے ول"""),
            ("ex2.jpeg", """زھر ڈتے کہیں، دواوی بھل گئے
اے دل ستم کوں ولا وی بھل گئے
دھن کوں ایں آئے فریب ارشد
کہیں دی یاد وچ خداوی بھل گئے"""),
            ("ex3.jpeg", """آج رات تاں ہاں نے سمنھڑ ڈے
وت ویسیں پیا متاں مر پودوں
تیں باج تاں جی نال جھڑے ہیں
مستان سردی بازی ہر یووں
تیڈا دل بھر تصور ہے ساڈا
گر چھوڑیں تھی بے گھر یووں
وت میل نصیب دے ذمے ہے ۔
آج رات جانی گھر یووں"""),
        ]
        for img_path, caption in examples:
            with gr.Column():
                gr.Image(value=img_path, interactive=False, elem_classes="fixed-box2")
                gr.Markdown(f"`{caption}`")

if __name__ == "__main__":
    demo.launch()