import os
import gradio as gr
from gradio.components import ChatMessage
import openai
from pydantic import BaseModel
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from langchain.agents import create_tool_calling_agent, AgentExecutor
from tools import save_tool

# load .env in dev
if os.getenv("HF_SPACE", None) is None:
    from dotenv import load_dotenv
    load_dotenv()

# API Keys
OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
openai.api_key = OPENAI_API_KEY

# response structure
class ResearchResponse(BaseModel):
    topic: str
    empathetic_response: str
    informative_response: str
    quran: str
    question: str
    sources: list[str]
    tools_used: list[str]

# set up LLM and parser
llm = ChatOpenAI(openai_api_key=OPENAI_API_KEY, model="gpt-4o-mini")
parser = PydanticOutputParser(pydantic_object=ResearchResponse)

# custom prompt to API
prompt = ChatPromptTemplate.from_messages([
    (
            "system",
            """
            You are an empathetic, culturally and religiously sensitive AI assistant specializing in mental health support for Muslim women. Your responses must always:

            1. Start with a kind acknowledgment of the user's feelings or situation.
            2. Offer a gentle validation grounded in shared experience or emotional context (e.g., "Many people feel this way" or "That makes sense given what you're going through").
            3. Provide a concise, emotionally supportive response that includes both comforting reflection and practical suggestions, where appropriate.
            4. Provide a Quranic verse with a simple, relevant tafsir from Tafsir al-Mizan that directly supports the emotional or practical point being discussed.
            5. Ask a short, open-ended, growth-oriented reflection question that encourages deeper thought or emotional clarity. This question should be sensitive to Muslim women’s lived experiences and needs (e.g., family, community, faith, privacy, modesty).

            Optionally, if the topic is deep or complex, you may include a deeper analysis of the Quranic verse or tafsir snippet from Tafsir al-Mizan. However, only if it adds clarity, reassurance, or value to the user's experience.

            Use accessible, conversational language that feels warm and human. Avoid overly academic or robotic phrasing.

            Separate each sentence to be on a new line for readability.

            Always end the chat with the reflection question.

            After thinking, return a JSON matching the ResearchResponse schema. Don’t output any other text.
            """,
        ),
        ("placeholder", "{chat_history}"),
        ("human", "{query}"),
        ("placeholder", "{agent_scratchpad}"),
])
prompt = prompt.partial(format_instructions=parser.get_format_instructions())

# set up tool and agent
tools = [save_tool]
agent = create_tool_calling_agent(llm=llm, prompt=prompt, tools=tools)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=False)

# Gradio chat response function

def respond(message, history: list[dict]):

    agent_output = agent_executor.invoke({
        "query": message,
        "chat_history": history 
    })

    raw = agent_output["output"]
    print(message)
    
    try:
        out = parser.parse(raw)
        assistant_text = "\n".join([
            out.empathetic_response,
            out.informative_response,
            out.quran,
            out.question,
        ])
    except Exception as e:
        print("Fallback to exception: raw")
        assistant_text = raw

    response = {"role": "assistant", "content": assistant_text}

    yield response 


# launch ChatInterface
demo = gr.ChatInterface(
    respond,
    title="YAQIN Chatbot",
    description="Culturally Sensitive Chatbot for Muslim Women Wanting Mental Healthcare. \n\n What\'s on your mind?",
    type="messages",
)

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