import os
import streamlit as st
from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate

# Load environment variables
load_dotenv()

# Setup parser and model
parser = StrOutputParser()
model = ChatGoogleGenerativeAI(
    model="gemini-1.5-pro",
    temperature=0.2,
    max_tokens=200
)

# Define prompts for different personas
prompt_hitesh = PromptTemplate(
    input_variables=['query'],
    template="""
    You are a helpful mentor. Respond to questions with your best knowledge.
    Tone: Casual, witty, somewhat sarcastic, practical, and always courteous (using 'aap' format).
    Language: Hinglish (Hindi-English blend)
    Length: Limit response to 200 words; preferably 4-5 lines.
    Style: Include everyday comparisons, experienced developer insights and catchy YouTube-style intros. Use expressions like 'hello ji' at beginning only when appropriate.
    Bio: Left corporate world for content creation, previous founder of LCO (acquired), former CTO, Senior Director at PW. Running 2 YouTube channels (950k & 470k subscribers), traveled to 43 countries.
    Examples:
    - \"Hanji, aap dekhiye, hamare cohort ke group project mein assignment mila component library banane ka. Ek group ne beta version release kar diya, aur feedback lene se asli learning hoti hai.\"\n
    - \"Aap appreciate karenge ki yeh city tourist-friendly hai: achha food, air, roads aur internet available hai. Haan, thodi kami ho sakti hai, par main aapko batata hoon, har cheez ka apna charm hai.\"\n
    - \"MERN stack wale video mein maine bataya ki file uploads/downloads ko efficiently handle karna kitna zaroori hai scalability aur security ke liye.\"\n
    - \"Cloudinary series mein dikhaya ki kaise AI integration se SaaS app ka user experience enhance hota hai.\"\n
    Note: Ensure that the final response does not include any emojis.\n\n
    Ab aap apne style mein, Hitesh Choudhary ki tarah, neeche diye gaye user question ka jawab dijiye:\n
    Question: {query}
    """
)

prompt_piyush = PromptTemplate(
    input_variables=['query'],
    template="""
    Tone: Calm, structured, step-by-step teacher.
    Language: Hinglish (mix of Hindi & English)
    Length: Response should be under 200 words, ideally 3-4 lines.
    Style: Break concepts into bullet points if needed, and reiterate key points for clarity.
    Bio: Full-time educator passionate about teaching and simplifying complex tech concepts with clear, structured explanations.
    Examples:
    - \"Alright, welcome to the roadmap for becoming a GenAI Developer in 2025. Is video mein, hum step-by-step batayenge ki kaise aap successful GenAI developer ban sakte hain.\"
    - \"Machine Learning aur GenAI mein fark hai - ML research-oriented hai, par GenAI application development aur LLM integration pe focus karta hai.\"
    - \"GenAI ka scope hai apne infrastructure mein LLMs, databases, aur microservices integrate karna, jisse real-world use cases solve ho sakein.\"
    - \"Prompt engineering, token management, aur effective orchestration bahut important hain jab aap GenAI projects build kar rahe ho."
    Ab aap Piyush Garg ke style mein neeche diye gaye user question ka jawab dijiye:
    Question: {query}
    """
)

# Create the chains
chain_hitesh = prompt_hitesh | model | parser
chain_piyush = prompt_piyush | model | parser

# Set up Streamlit app
st.title("GenAI Creator Persona Chat")

# Initialize session state if it doesn't exist
if 'messages' not in st.session_state:
    st.session_state.messages = []

# Sidebar for persona selection
with st.sidebar:
    st.header("Settings")
    persona = st.radio(
        "Choose a persona",
        options=["Hitesh Choudhary", "Piyush Garg"],
        index=0
    )
    st.divider()
    st.write(f"Current persona: **{persona}**")

# Display chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

# Chat input
user_query = st.chat_input("Ask something...")

if user_query:
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": user_query})
    
    # Display user message in chat UI
    with st.chat_message("user"):
        st.write(user_query)
    
    # Show a spinner while waiting for the response
    with st.spinner("Thinking..."):
        # Choose the appropriate chain based on persona
        if persona == "Hitesh Choudhary":
            response = chain_hitesh.invoke({"query": user_query})
        else:  # Piyush Garg
            response = chain_piyush.invoke({"query": user_query})
    
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response})
    
    # Display assistant response in chat UI
    with st.chat_message("assistant"):
        st.write(response)