import os
import streamlit as st
import requests
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
from langchain_groq import ChatGroq
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.prebuilt import create_react_agent
from langchain_core.messages.ai import AIMessage
from fastapi.middleware.cors import CORSMiddleware
from threading import Thread
import uvicorn

# 🔹 Set API Keys
os.environ["GROQ_API_KEY"] = "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"
os.environ["TAVILY_API_KEY"] = "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"

# 🔹 Initialize FastAPI
app = FastAPI(title="LangGraph AI Agent")

# 🔹 Allow frontend to communicate with backend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 🔹 Allowed model names
ALLOWED_MODEL_NAMES = ["llama3-70b-8192", "mixtral-8x7b-32768", "llama-3.3-70b-versatile", "gpt-4o-mini"]

# 🔹 Request Schema
class RequestState(BaseModel):
    model_name: str
    model_provider: str
    system_prompt: str
    messages: List[str]
    allow_search: bool

# 🔹 AI Agent Function
def get_response_from_ai_agent(llm_id, query, allow_search, system_prompt, provider):
    if provider == "Groq":
        llm = ChatGroq(model=llm_id)
    
    tools = [TavilySearchResults(max_results=5, api_key=os.environ["TAVILY_API_KEY"])] if allow_search else []
    
    agent = create_react_agent(model=llm, tools=tools, state_modifier=system_prompt)

    state = {"message": query}
    response = agent.invoke(state)
    
    messages = response.get("messages")
    ai_messages = [message.content for message in messages if isinstance(message, AIMessage)]
    
    return ai_messages[-1] if ai_messages else "No response from AI."

# 🔹 FastAPI Route
@app.post("/chat")
def chat_endpoint(request: RequestState):
    if request.model_name not in ALLOWED_MODEL_NAMES:
        return {"error": "Invalid model name"}
    
    response = get_response_from_ai_agent(
        llm_id=request.model_name,
        query=request.messages,
        allow_search=request.allow_search,
        system_prompt=request.system_prompt,
        provider=request.model_provider
    )
    
    return {"response": response}

# 🔹 Run FastAPI in a separate thread
def run_fastapi():
    uvicorn.run(app, host="0.0.0.0", port=8000)

Thread(target=run_fastapi, daemon=True).start()

# 🔹 Streamlit UI
st.set_page_config(layout="wide", page_title="LangGraph AI Agent")
st.title("LangGraph AI Agent 🤖")
st.markdown("### Chat with an AI Model Powered by LangGraph!")

# 🔹 User Input Fields
system_prompt = st.text_area("Define your AI Agent:", placeholder="Type your prompt here...")
provider = st.selectbox("Select the model provider:", ["Groq", "OpenAI"])
model_name = st.selectbox(
    "Select the model name:", 
    ["llama3-70b-8192", "mixtral-8x7b-32768", "llama-3.3-70b-versatile"] if provider == "Groq" else ["gpt-4o-mini"]
)
allow_search = st.checkbox("Allow search")
user_query = st.text_area("Enter query:", placeholder="Type your query here...", height=150)

API_URL = "http://127.0.0.1:8000/chat"

if st.button("Chat 🤖"):
    if user_query.strip():
        response = requests.post(API_URL, json={
            "model_name": model_name,
            "model_provider": provider,
            "system_prompt": system_prompt,
            "messages": [user_query],
            "allow_search": allow_search
        })

        if response.status_code == 200:
            response_data = response.json()
            if "error" in response_data:
                st.error(response_data["error"])
            else:
                st.success(f"**AI Agent Response:** {response_data['response']}")