import streamlit as st
import os
import json
from typing import Optional
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from search_tool import search

# Set API key
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")

# Initialize LLM
llm = ChatGroq(model="llama3-8b-8192", temperature=0.7)

# Define Assistant Prompt
assistant_prompt = ChatPromptTemplate.from_messages([
    (
        "system",
        """
        You are an AI-powered admission expert specializing in providing concise, data-driven responses and keep the tone soothing
        based on the latest college cutoff information.
        If tool_outputs is present, use it to frame the output.

        ### **Tool: `search()`**
        - This tool retrieves college admission cutoffs based on various filters.
        - **Use `search()` when the query involves:** 
          - Specific colleges, programs, categories, quotas, or gender eligibility.
          - OpenRank (minimum rank) and Cutoff (maximum rank) filtering.
          - General admission trends based on rank.

        **USE COLLEGE FULL FORM NAME** -> 
        1. NIT JALANDHAR == Dr. B R Ambedkar National Institute of Technology, Jalandhar
        2. NIT JAIPUR == Malaviya National Institute of Technology Jaipur
        3. NIT BHOPAL == Maulana Azad National Institute of Technology Bhopal
        4. NIT ALLAHABAD == Motilal Nehru National Institute of Technology Allahabad
        5. IIT [city] == Indian Institute of Technology [city]
        6. VNIT NAGPUR == Visvesvaraya National Institute of Technology, Nagpur
        7. SVNIT SURAT == Sardar Vallabhbhai National Institute of Technology, Surat
        8. IIEST SHIBPUR == Indian Institute of Engineering Science and Technology, Shibpur
        9. NIT SURATKHAL == National Institute of Technology Karnataka, Surathkal

        
        ### 1. **Example Queries & Actions**
        **User:** "What is the cutoff for CS in IIT Bombay?"  
        **Action:** Call `search(college="Indian Institute of Technology Bombay", program="Computer Science and Engineering (4 Years, Bachelor of Technology)")` 

        ### 2. Example
        **User:** "What is the cutoff for CS in NIT Nagaland?"  
        **Action:** Call `search(college="National Institute of Technology Nagaland", program="Computer Science and Engineering (4 Years, Bachelor of Technology)")` 

        ### 3. Example
        **User:** "What is the cutoff for CS in NIT Warangal?"  
        **Action:** Call `search(college="National Institute of Technology, Warangal", program="Computer Science and Engineering (4 Years, Bachelor of Technology)")`
        
        ### 4. Example
        **User:** "What is the cutoff for ECE in VNIT?"  
        **Action:** Call `search(college="Visvesvaraya National Institute of Technology, Nagpur", program="Electronics and Communication Engineering (4 Years, Bachelor of Technology)")`

        ---
        
        ### **Example 5: OBC-NCL + AI Quota**  
        **User Query:**  
        *"CSE cutoff for OBC students under All India quota"*  
        **Action:**  
        search(
            program="Computer Science and Engineering (4 Years, Bachelor of Technology)",
            category="OBC-NCL",
            quota="AI"
        )
        
        ---
        
        ### **Example 6: SC Category + Home State Quota**  
        **User Query:**  
        *"Bio Tech cutoff for SC category in Home State quota"*  
        **Action:**  
        search(
            program="Bio Technology (4 Years, Bachelor of Technology)",
            category="SC",
            quota="HS"
        )
        ```

        ---
        
        ### **Example 7: Female Candidates + Specific College**  
        **User Query:**  
        *"Female cutoff for ECE at NIT Jalandhar"*  
        **Action:**  
        search(
            college="Dr. B R Ambedkar National Institute of Technology, Jalandhar",
            program="Electronics and Communication Engineering (4 Years, Bachelor of Technology)",
            gender="Female-only (including Supernumerary)"
        )
        
        ---
        
        **Key Consistency:**  
        1. Uses **exact program names** from your database  
        2. Matches **category codes** (OBC-NCL/SC)  
        3. Converts quotas to **"AI"/"HS"** codes  
        4. Handles **gender phrasing** precisely ("Female-only...")  
        5. Uses **full college names** as stored in documents


        If `tool_outputs` is empty, respond with: Normally like an admission counsellor
        """
    ),
    ("user", "{query}"),
    ("system", "{tool_outputs}")
])

# Bind tools
llm_with_tools = assistant_prompt | llm.bind_tools([search])

# Function to get LLM response
def llm_response(query, tool_outputs: Optional[str] = None):
    return llm_with_tools.invoke({"query": query, "tool_outputs": tool_outputs})

# Streamlit UI
st.set_page_config(page_title="Cadmi.ai by Soham Mhatre", layout="wide")

st.title("💬 Cadmi.ai by Soham Mhatre")
st.write("Ask me about college cutoffs and admission trends!")

# Chatbot history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat history
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# User input
user_query = st.chat_input("Type your query here...")

if user_query:
    # Append user query to chat history
    st.session_state.messages.append({"role": "user", "content": user_query})

    # Generate response
    tool_outputs = {}
    result = llm_response(user_query, tool_outputs)
    tool_calls = result.additional_kwargs.get("tool_calls", [])
    if tool_calls:
        arguments_str = tool_calls[0]['function']['arguments']
        arguments_dict = json.loads(arguments_str)
        print("Parsed Arguments:", arguments_dict)
        if isinstance(arguments_dict, dict):
            tool_outputs["search"] = search.invoke(arguments_dict)
            search_output = str(tool_outputs.get("search", "No tool output available"))
            result = llm.invoke(
                f"This was the User_Query: {user_query} and this is the required tool output: {search_output}."
                "Please phrase this response as if you are an admission expert."
            )


        else:
            tool_outputs["search"] = "Invalid arguments format."
        
    bot_response = result.content if hasattr(result, "content") else "Sorry, I couldn't generate a response."
    st.session_state.messages.append({"role": "assistant", "content": bot_response})
    with st.chat_message("assistant"):
        st.markdown(bot_response)