import gradio as gr
from openai import OpenAI
import base64
import io

# Initialize the client pointing to your vLLM server
client = OpenAI(
    base_url="http://172.16.34.29:8006/v1",
    api_key="vllm-token", 
)

def encode_image_to_base64(image):
    """Converts PIL image to raw base64 string (no data-uri prefix)."""
    if image is None:
        return None
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

def run_translation(source_code, target_code, text_input, image_input):
    # Construct the base dictionary
    # The schema requires all these keys to be present in the mapping
    payload = {
        "source_lang_code": source_code,
        "target_lang_code": target_code,
        "text": None,
        "image": None
    }

    if image_input is not None:
        payload["type"] = "image"
        payload["image"] = encode_image_to_base64(image_input)
    else:
        if not text_input.strip():
            return "Please provide text or an image."
        payload["type"] = "text"
        payload["text"] = text_input

    try:
        # Crucial: We pass the payload as the single item in the content list
        response = client.chat.completions.create(
            model="translate_gemma",
            messages=[{
                "role": "user", 
                "content": [payload]  # vLLM expects exactly [ { ... } ]
            }],
            max_tokens=500
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

# --- Gradio UI Layout ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🌍 TranslateGemma 27B")
    gr.Markdown("Corrected schema for vLLM inference.")
    
    with gr.Row():
        src_code = gr.Textbox(label="Source Language Code", value="en")
        tgt_code = gr.Textbox(label="Target Language Code", value="bn")

    with gr.Row():
        with gr.Column():
            text_box = gr.Textbox(label="Text Input", placeholder="Type English here...", lines=5)
            image_box = gr.Image(label="Image Input", type="pil")
            submit_btn = gr.Button("Translate", variant="primary")
        
        with gr.Column():
            output_box = gr.Textbox(label="Bangla Translation", interactive=False, lines=10)

    submit_btn.click(
        fn=run_translation,
        inputs=[src_code, tgt_code, text_box, image_box],
        outputs=output_box
    )

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