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 langchain_openai import OpenAIEmbeddings
from pinecone import Pinecone, ServerlessSpec
import subprocess
import re
import json
from langchain_core.messages import AIMessage
from langchain_core.output_parsers import JsonOutputParser

from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory


# 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"]
embedding_model = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)

PINECONE_API_KEY = os.environ["PINECONE_API_KEY"]
PINECONE_INDEX = os.environ["PINECONE_INDEX"]
pc = Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index(PINECONE_INDEX)

#Run this when upserting new documents to the pinecone index by adding new files to the data or journal folders
#subprocess.run(["python", "pinecone_rag.py"])

# response structure
class ResearchResponse(BaseModel):
    #topic: str
    empathetic_response: str
    informative_response: str
    quran: str
    question: 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, from the Relevant Islamic Background Information.
        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.

        **Important Output Instructions**

        You **must** return a valid **JSON object** using the following format:
        
        {format_instructions}
        
        - Begin your output immediately with the JSON object.
        - Do **not** wrap it in backticks or markdown.
        - Do **not** include any explanation, greeting, or text before or after the JSON.
        - Only the JSON should be returned. Strictly follow the schema.
        
        If you do not follow this format, your response will be rejected and considered invalid.
        """
    ),
    ("placeholder", "{chat_history}"),
    ("human", "{query}"),
    ("placeholder", "{agent_scratchpad}"),
])
prompt = prompt.partial(format_instructions=parser.get_format_instructions())

# set up agent
tools = []
agent = create_tool_calling_agent(llm=llm, prompt=prompt, tools=tools)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=False)
history = ""

store = {}


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]

agent_with_chat_history = RunnableWithMessageHistory(
    agent_executor,
    get_session_history,
    input_messages_key="query",
    history_messages_key="chat_history",
)


# Gradio chat response function

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

    # Step 1: Embed user message
    try:
        vectorized_input = embedding_model.embed_query(message)
    except Exception as e:
        print("Embedding failed:", e)
        vectorized_input = []

    # Step 2: Query Pinecone
    TOP_K_GLOBAL = 3

    context = ""
    
    if vectorized_input:
        try:
            pinecone_response = index.query(
                namespace="ns4",
                vector=vectorized_input,
                top_k=3,
                include_metadata=True
            )
            matches = pinecone_response.get("matches", [])
            context = "\n".join([match['metadata']['text'] for match in matches if 'metadata' in match])
            print("From Pinecone:", context)
            if username or username.strip().lower() == "Lina":
                pinecone_journal_response = index.query(
                    namespace=username,
                    vector=vectorized_input,
                    top_k=3,
                    include_metadata=True
                )
                journal_matches = pinecone_journal_response.get("matches", [])
                journal_context = "\n".join([match['metadata']['text'] for match in journal_matches if 'metadata' in match])
                print("From Pinecone:", journal_context)
                context_prefix = f"Relevant Islamic background information:\n{context}\nRelevant User previous Journal entry information:\n{journal_context}\n\nUser query:"
            else:
                context_prefix = f"Relevant Islamic background information:\n{context}\n\nUser query:"
        except:
            print("Pinecone query failed.")
    else:
        print("Vectorised input is empty.")
    
    # Append context to the query
    
    full_query = f"{context_prefix} {message}"
    print("\n Full query sent to LLM:\n", full_query)

    # Step 4: Agent invocation
    agent_output = agent_with_chat_history.invoke(
        {"query": full_query},
        config={"configurable": {"session_id": "<foo>"}},
    )

    raw = agent_output["output"]
    print("=== RAW OUTPUT FROM LLM ===")
    print(raw)
    out = parser.parse(raw)    
    assistant_text = "\n".join([
        out.empathetic_response,
        out.informative_response,
        out.quran,
        out.question,
    ])

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


# launch ChatInterface


with gr.Blocks(fill_height=True) as demo:
    username = gr.State()

    with gr.Column(visible=True) as login_container:
        username_input = gr.Textbox(label="Enter your username to start.\nIf you're new, enter 'Guest'")
        submit_button = gr.Button("Start Chat")

    # Wrap ChatInterface inside a column so we can control its visibility
    with gr.Column(visible=False) as chatbot_container:
        chatbot = gr.ChatInterface(
            fn=respond,
            title="YAQIN Chatbot",
            description="Culturally Sensitive Chatbot for Muslim Women Wanting Mental Healthcare.\n\nWhat's on your mind?",
            type="messages",
            additional_inputs=[username],
        )

    def start_chatbot(name):
        return gr.update(visible=False), gr.update(visible=True), name

    submit_button.click(
        fn=start_chatbot,
        inputs=username_input,
        outputs=[login_container, chatbot_container, username]
    )

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