import gradio as gr
from google import genai
from google.genai import types
import os

# Get the Gemini API key from environment variable
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
    raise ValueError("GEMINI_API_KEY environment variable not set.")

# Create a Gemini client using the API key
client = genai.Client(api_key=api_key)

def simplify_bengali_text(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    full_prompt = (
        f"এই লেখাটি সহজ বাংলায় লিখো যাতে সাধারণ পাঠক সহজেই বুঝতে পারে। "
        f"এটি কোনো সারাংশ, তালিকা, বা বিন্দুযুক্ত আকারে লেখা হোক না। "
        f"মূল বক্তব্য যেন অপরিবর্তিত থাকে, কিন্তু ভাষাটি সহজ ও প্রাঞ্জল হয়।\n\n"
        f"মূল পাঠঃ\n{message}"
    )

    part = types.Part.from_text(text=full_prompt)

    contents = [
        types.Content(
            role="user",
            parts=[part]
        )
    ]

    config = types.GenerateContentConfig(
        temperature=temperature,
        top_p=top_p,
        max_output_tokens=max_tokens,
        response_modalities=["TEXT"],
        safety_settings=[
            types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
            types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="OFF"),
            types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="OFF"),
            types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
        ]
    )

    response = ""
    for chunk in client.models.generate_content_stream(
        model="gemini-2.0-flash-001",
        contents=contents,
        config=config,
    ):
        if hasattr(chunk, "text"):
            response += chunk.text
            yield response

# Gradio UI
demo = gr.ChatInterface(
    simplify_bengali_text,
    additional_inputs=[
        gr.Textbox(value="Please simplify the given Bengali text in same meaning but in easy undestandable Bengali that can understood by masses in recent times.", label="System message"),
        gr.Slider(minimum=1, maximum=8192, value=2048, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
    ],
)

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