# app.py
import streamlit as st
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel
from huggingface_hub import login
import warnings

warnings.filterwarnings("ignore")

# Set page configuration
st.set_page_config(
    page_title="DuckDuckGo Search Tool",
    page_icon="🔍",
    layout="wide",
)

# Function to initialize the API with user-provided token
def initialize_api(token=None):
    if token:
        login(token)
        try:
            # Initialize the model
            return HfApiModel(model_id="meta-llama/Llama-3.3-70B-Instruct")
        except Exception as e:
            st.error(f"Error initializing model: {str(e)}")
            return None
    else:
        st.warning("Please enter your Hugging Face API token to use this app.")
        return None

# Initialize tools
search_tool = DuckDuckGoSearchTool()

# Streamlit App
st.title("DuckDuckGo Search Tool")
st.markdown("Search the web using DuckDuckGo and get AI-powered responses")

# Initialize session state for token
if "hf_token" not in st.session_state:
    st.session_state.hf_token = ""
    st.session_state.model = None

# Sidebar for token input
with st.sidebar:
    st.title("Configuration")
    
    # Token input 
    token_input = st.text_input(
        "Enter your Hugging Face API Token:", 
        value=st.session_state.hf_token,
        type="password",
        help="Get your token from huggingface.co/settings/tokens"
    )
    
    if token_input != st.session_state.hf_token:
        st.session_state.hf_token = token_input
        # Reset model when token changes
        st.session_state.model = None
    
    # Button to initialize/test the token
    if st.button("Initialize API"):
        with st.spinner("Testing your token..."):
            model = initialize_api(st.session_state.hf_token)
            if model:
                st.session_state.model = model
                st.success("✅ API initialized successfully!")
            else:
                st.error("❌ Failed to initialize the API with the provided token.")

# Check if the model is initialized
if not st.session_state.model and st.session_state.hf_token:
    with st.spinner("Initializing model..."):
        st.session_state.model = initialize_api(st.session_state.hf_token)
        if st.session_state.model:
            st.success("Model initialized successfully!")

# Main content area - only show if token is provided
if st.session_state.hf_token and st.session_state.model:
    st.header("AI Web Search")
    st.markdown("Ask any question and the AI will search the web for answers")
    
    query = st.text_input("Enter your search query:", placeholder="What are the latest advancements in renewable energy?")
    
    if st.button("Search"):
        if query:
            with st.spinner("Searching and processing..."):
                agent = CodeAgent(tools=[search_tool], model=st.session_state.model)
                response = agent.run(query)
                st.success("Search complete!")
                st.markdown("### Results")
                st.markdown(response)
        else:
            st.warning("Please enter a search query")
else:
    # Token not provided
    st.info("Please enter your Hugging Face API token in the sidebar to get started.")
    
    # Show more information to help users understand what they need
    st.markdown("""
    ### How to Get a Hugging Face Token
    
    1. Go to [huggingface.co](https://huggingface.co/) and create an account if you don't have one
    2. Visit your [settings page](https://huggingface.co/settings/tokens)
    3. Create a new token with read access
    4. Copy the token and paste it in the sidebar
    """)