import gradio as gr
from google import genai
from google.genai import types
from PIL import Image
import json
import os

def load_prompts(file_path="prompts.json"):
    """Load system prompts and style text from JSON file."""
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception as e:
        print(f"⚠️ Could not load prompts.json: {e}")
        return {
            "base_system": "You are a bilingual translator fluent in Ebonics and English.",
            "styles": {},
            "direction_prompts": {},
            "platform_notes": {}
        }

PROMPTS = load_prompts()

def merge_platform_notes():
    """Combine all platform_notes into a formatted string for contextual grounding."""
    notes = PROMPTS.get("platform_notes", {})
    if not notes:
        return ""
    return "\n\n".join(
        [f"[{k.replace('_', ' ').title()}]\n{v}" for k, v in notes.items()]
    )

PLATFORM_CONTEXT = merge_platform_notes()

def create_client(api_key: str):
    """Initialize Gemini client using provided API key."""
    os.environ["GOOGLE_API_KEY"] = api_key
    return genai.Client(api_key=api_key)

def get_style_instruction(direction: str, style: str) -> str:
    """Retrieve appropriate style text from prompts.json."""
    try:
        return PROMPTS["styles"][direction].get(style, PROMPTS["styles"][direction].get("Default", ""))
    except KeyError:
        return ""

def build_system_instruction(direction: str, style: str) -> str:
    """Combine base system prompt, style instructions, and platform notes."""
    base_system = PROMPTS.get("base_system", "")
    style_instruction = get_style_instruction(direction, style)
    system_instruction = (
        f"{base_system}\n\n"
        f"Adapt the translation style according to this guidance:\n{style_instruction}\n\n"
        f"Use the following platform and authenticity notes as contextual grounding:\n{PLATFORM_CONTEXT}"
    )
    return system_instruction

def gemini_translate(client, text: str, direction: str, style: str):
    """Translate between Ebonics and Standard English using prompts.json + platform notes."""
    system_instruction = build_system_instruction(direction, style)
    direction_prompt = PROMPTS.get("direction_prompts", {}).get(direction, "")

    prompt = f"{direction_prompt}\n\nStyle: {style}\n\n{text}"

    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=prompt,
        config=types.GenerateContentConfig(
            system_instruction=system_instruction,
            temperature=0.7 if "English ➜ Ebonics" in direction else 0.5,
            thinking_config=types.ThinkingConfig(thinking_budget=0)
        ),
    )
    return response.text.strip()

def gemini_translate_image(client, image, direction: str, style: str):
    """Extract text from image and translate using prompts.json + platform notes."""
    system_instruction = build_system_instruction(direction, style)

    prompt = (
        f"Extract all readable text from this image, then translate it according to the direction and style.\n\n"
        f"Direction: {direction}\nStyle: {style}\n\n"
        f"Return only the translated text."
    )

    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=[image, prompt],
        config=types.GenerateContentConfig(
            system_instruction=system_instruction,
            temperature=0.7 if "English ➜ Ebonics" in direction else 0.5,
            thinking_config=types.ThinkingConfig(thinking_budget=0)
        ),
    )
    return response.text.strip()

def translate(api_key, text, image, direction, style):
    """Main handler for text/image translation."""
    if not api_key:
        return "⚠️ Please enter your Gemini API key in the sidebar."

    client = create_client(api_key)
    try:
        if image is not None:
            pil_image = Image.open(image)
            result = gemini_translate_image(client, pil_image, direction, style)
        elif text.strip():
            result = gemini_translate(client, text, direction, style)
        else:
            result = "⚠️ Please enter text or upload an image."
    except Exception as e:
        result = f"❌ Error: {e}"
    return result


# ==== GRADIO UI ====

with gr.Blocks(theme=gr.themes.Soft(), title="Ebonics ⇄ English Translator") as demo:
    gr.Markdown("# Ebonics ⇄ Standard English Translator")

    with gr.Row():
        with gr.Column(scale=1):
            api_key = gr.Textbox(
                label="🔑 Gemini API Key",
                type="password",
                placeholder="Enter your Gemini API key here"
            )

            direction = gr.Radio(
                ["Ebonics ➜ English", "English ➜ Ebonics"],
                label="Translation Direction",
                value="Ebonics ➜ English"
            )

            style = gr.Dropdown(
                label="Translation Style",
                choices=["Formal", "Casual", "Academic", "Default"],
                value="Casual",
                info="Choose a translation tone or style (changes with direction)."
            )

            text_input = gr.Textbox(
                label="Text Input",
                lines=6,
                placeholder="Enter text here or upload an image below..."
            )
            image_input = gr.Image(label="Image Input (optional)", type="filepath")

            translate_button = gr.Button("🔁 Translate")

        with gr.Column(scale=1):
            output_box = gr.Textbox(label="Translated Output", lines=8)

    def update_styles(direction):
        """Update style options dynamically based on direction."""
        if direction == "Ebonics ➜ English":
            return gr.update(
                choices=["Formal", "Casual", "Academic", "Default"],
                value="Casual"
            )
        else:
            return gr.update(
                choices=["Modern", "Classic", "Chill", "Energetic", "Default"],
                value="Modern"
            )

    direction.change(update_styles, inputs=[direction], outputs=[style])

    translate_button.click(
        translate,
        inputs=[api_key, text_input, image_input, direction, style],
        outputs=output_box
    )

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