import streamlit as st
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel
from huggingface_hub import login
import warnings

warnings.filterwarnings("ignore")

# Initialize tools first (outside session state)
search_tool = DuckDuckGoSearchTool()

def main():
    # Streamlit config
    st.set_page_config(
        page_title="DuckDuckGo Search Tool",
        page_icon="🔍",
        layout="wide",
    )
    
    # Session state initialization
    if "model" not in st.session_state:
        st.session_state.model = None
        st.session_state.hf_token = ""

    # Sidebar logic
    with st.sidebar:
        st.title("Configuration")
        new_token = st.text_input(
            "Hugging Face Token",
            type="password",
            value=st.session_state.hf_token
        )
        
        if new_token != st.session_state.hf_token:
            st.session_state.hf_token = new_token
            st.session_state.model = None  # Reset model on token change

        if st.button("Initialize API"):
            with st.spinner("Initializing..."):
                try:
                    login(st.session_state.hf_token)
                    st.session_state.model = HfApiModel(model_id="meta-llama/Llama-3.3-70B-Instruct")
                    st.success("API Initialized!")
                except Exception as e:
                    st.error(f"Error: {str(e)}")

    # Main app logic
    if st.session_state.model:
        st.title("AI Web Search")
        query = st.text_input("Enter your search query")
        
        if query and st.button("Search"):
            with st.spinner("Searching..."):
                agent = CodeAgent(tools=[search_tool], model=st.session_state.model)
                response = agent.run(query)
                st.markdown(f"## Results\n{response}")
    else:
        st.info("Please enter your Hugging Face token in the sidebar")

if __name__ == "__main__":
    main()  # No infinite loop needed when using proper streamlit run