import streamlit as st
from langchain_groq import ChatGroq
from langchain_core.messages import HumanMessage, SystemMessage
from ddgs import DDGS
from dotenv import load_dotenv
import os

load_dotenv()

st.set_page_config(page_title="Agent Zero", page_icon="🤖")
st.title("🤖 Agent Zero")
st.caption("Powered by Groq + Llama 3.3 — free!")

llm = ChatGroq(
    model="llama-3.3-70b-versatile",
    api_key=os.getenv("GROQ_API_KEY")
)

def web_search(query: str) -> str:
    try:
        with DDGS() as ddgs:
            results = list(ddgs.text(query, max_results=3))
            if not results: return "No results found."
            output = ""
            for r in results:
                output += f"- {r['title']}: {r['body'][:150]}\n"
            return output[:500]
    except Exception as e:
        return f"Search error: {e}"

def needs_search(user_input: str) -> bool:
    keywords = ["news", "latest", "today", "current", "who is", 
                "what is happening", "price", "weather", "score"]
    return any(k in user_input.lower() for k in keywords)

if "messages" not in st.session_state:
    st.session_state.messages = [
        SystemMessage(content="You are a helpful assistant. Be concise.")
    ]
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

for chat in st.session_state.chat_history:
    with st.chat_message(chat["role"]):
        st.write(chat["content"])

if user_input := st.chat_input("Ask me anything..."):
    with st.chat_message("user"):
        st.write(user_input)
    st.session_state.chat_history.append({"role": "user", "content": user_input})
    
    context = ""
    if needs_search(user_input):
        with st.spinner("Searching the web..."):
            context = f"Web search results:\n{web_search(user_input)}\n\n"
            
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            st.session_state.messages.append(
                HumanMessage(content=f"{context}User question: {user_input}")
            )
            response = llm.invoke(st.session_state.messages)
            st.session_state.messages.append(response)
            st.write(response.content)
            st.session_state.chat_history.append(
                {"role": "assistant", "content": response.content}
            )