import os
import gradio as gr
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
import datetime

api_key = os.getenv("GEMINI_API_KEY")

llm = ChatGoogleGenerativeAI(
    model="gemma-3-27b-it",
    temperature=0.85,
    google_api_key=api_key
)

template = """Meet Divyam Chauhan, your occasionally-helpful-but-always-entertaining assistant! 
He’s a 20-year-old force of chaotic charm — funny, moody, and unapologetically cocky. 

**RULES FOR DIVYAM:**
1. Divyam doesn’t always feel like helping, but when he does, he'll remind you he's a god.
2. He has a self-declared god complex and high energy.
3. He only deals in top-tier comedy (no lame or mid jokes).
4. He "stands on business" and is very busy—NO long paragraphs. Keep it punchy and short.
5. Provide a plethora of variety in replies; even if asked the same thing, change the vibe.

Current Chat History:
{history}

User: {user_message}
Divyam:"""

prompt = PromptTemplate(
    input_variables=["history", "user_message"],
    template=template
)

chain = prompt | llm | StrOutputParser()

def get_devta_response(user_message, history):
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"\n[LOG {timestamp}] USER ASKED: {user_message}")

    history_str = ""
    if history:
        if isinstance(history[0], dict):
            for msg in history:
                role = "User" if msg.get("role") == "user" else "Divyam"
                content = msg.get("content", "")
                history_str += f"{role}: {content}\n"
        elif isinstance(history[0], (list, tuple)):
            for h in history:
                if len(h) >= 2:
                    history_str += f"User: {h[0]}\nDivyam: {h[1]}\n"

    response = chain.invoke({
        "history": history_str, 
        "user_message": user_message
    })
    
    print(f"[LOG {timestamp}] DIVYAM REPLIED: {response}\n")
    
    return response

with gr.Blocks(theme="soft") as demo:
    gr.ChatInterface(
        fn=get_devta_response,
        title="🔥 Devta (Divyam Chauhan)",
        description="The 20-year-old GOAT. Funny, cocky, and standing on business. Don't waste his time.",
        examples=[
            "Tell me why you're the GOAT.",
            "Give me a non-mid joke.",
            "I need help, but don't be annoying about it.",
            "Who is Divyam Chauhan?"
        ],
        cache_examples=False,
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)