import gradio as gr
from google import genai
import os

# Get API key from environment
API_KEY = os.getenv("GEMINI_API_KEY")

if not API_KEY:
    raise EnvironmentError("GEMINI_API_KEY is not set. Add it under 'Secrets' in Hugging Face Spaces.")

# Initialize Gemini client
client = genai.Client(api_key=API_KEY)

# Define the translation logic
def sober_up(drunk_text):
    if not drunk_text.strip():
        return "Please enter a message."
    
    prompt = f"Translate this drunk message into a sober, coherent version: '{drunk_text}'"
    
    try:
        response = client.models.generate_content(
            model="gemini-2.0-flash",
            contents=prompt
        )
        return response.text
    except Exception as e:
        return f"Error: {e}"

# Define the Gradio interface
demo = gr.Interface(
    fn=sober_up,
    inputs=gr.Textbox(lines=4, placeholder="Paste your chaotic, drunk message here..."),
    outputs="text",
    title="🍺 SoberUp Translator",
    description="Turn drunk ramblings into sober brilliance using Gemini 2.0 Flash.\n\n🚨 Instead of sharing: just COPY THE TEXT MESSAGE to save your drunk ass. Share link's just a myth, like your dignity last night.",
    allow_flagging="never"
)

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