import streamlit as st
from langchain.agents import AgentType, initialize_agent
from langchain import HuggingFaceHub
from langchain.tools import Tool
from datetime import datetime
from langchain.tools import DuckDuckGoSearchRun
import os

# Set up Hugging Face Hub
token = os.environ['HF_TOKEN']
hub_llm = HuggingFaceHub(
    repo_id='meta-llama/Llama-3.1-8B-Instruct',
    huggingfacehub_api_token=token
)

# Set the page title and icon
st.set_page_config(
    page_title="AI Driven Search",
    page_icon="🔍",
    layout="wide",
)

# Custom CSS style for the title block
st.markdown(
    """
    <style>
        .title-block {
            background-color: #3498db;
            color: #ffffff;
            padding: 20px;
            border-radius: 10px;
            margin-bottom: 20px;
        }
        .subtitle {
            color: #2c3e50;
        }
    </style>
    """,
    unsafe_allow_html=True,
)

# Title block with custom styling
st.markdown('<div class="title-block">', unsafe_allow_html=True)
st.title("🌐 AI Powered Search Engine")
st.markdown("### Find what you're looking for with the power of AI!")
st.markdown("</div>", unsafe_allow_html=True)

# Subtitle and description with custom styling
st.markdown('<div class="subtitle">', unsafe_allow_html=True)
st.subheader("How it works:")
st.write(
    "Our search engine is powered by DuckDuckGo search and uses language models "
    "that understand your queries and provide accurate results."
)
st.markdown("</div>", unsafe_allow_html=True)

# Form for user input
with st.form(key="form"):
    user_input = st.text_input("Ask your question")
    submit_clicked = st.form_submit_button("Search")

# Run search if form is submitted
if submit_clicked:
    datetime_tool = Tool(
        name="Datetime",
        func=lambda x: datetime.now().isoformat(),
        description="Returns the current datetime"
    )
    search = DuckDuckGoSearchRun()
    search_tool = Tool(
        name="Search",
        func=search,
        description="Performs an internet search using DuckDuckGo"
    )

    # Initialize the agent
    agent_chain = initialize_agent(
        [search_tool, datetime_tool],
        hub_llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
        handle_parsing_errors=True,
    )
    result = agent_chain.run(user_input)
    st.success(result)

# Footer with custom styling
st.markdown(
    '<p style="text-align:center; color:#7f8c8d;">Built with ❤️ by Abhishek | '
    '<a href="https://github.com/your_username/your_repo" style="color:#3498db;">GitHub Repo</a></p>',
    unsafe_allow_html=True,
)