import os
from datetime import datetime

import streamlit as st

from helper import ChatBot, current_year, invoke_search_api, save_to_audio

# API KEY
API_KEY = os.environ["API_KEY"]

# Front-end
st.set_page_config(layout="wide")
st.title("SearchBot 🤖")

# Front-end: Sidebar
with st.sidebar:
    with st.expander("Instruction Manual"):
        st.markdown(
            """
            ## SearchBot 🤖
            This Streamlit app allows you to search anything.

            ### How to Use:
            1. **Source**: Select your favorite source, i.e. default is Google.
            1. **Num**: Select number of output you want to display, i.e. default is 20.
            1. **Location**: Select your location so content such as weather can be customized, i.e. default is New York.
            2. **Response**: The app will display a response table.
            3. **Chat History**: Previous conversations will be shown on the app and can be cleared using "Clear Session" button.

            ### Source:
            - **Site**: The source app is deployed on **HuggingFace** with this [link](https://huggingface.co/spaces/eagle0504/searchbot).

            ### Credits:
            - **Developer**: [Yiqiao Yin](https://www.y-yin.io/) | [App URL](https://huggingface.co/spaces/eagle0504/serpapi-demo) | [LinkedIn](https://www.linkedin.com/in/yiqiaoyin/) | [YouTube](https://youtube.com/YiqiaoYin/)

            Enjoy chatting with SearchBot!
        """
        )

    # Example:
    st.success("Example: Yiqiao Yin.")
    st.success("Example: Weather in Westchester, NY today")
    st.success("Example: Latest news about Tesla")
    st.success(
        "Example: Taking medicine during pregnancy (you can select WebMD to filter only results from WebMD)"
    )

    # Input
    src_key_word = st.selectbox(
        "Source of page you prefer?",
        (
            "Google",
            "Google Finance",
            "Google Scholar",
            "Yahoo",
            "WSJ",
            "Bloomberg",
            "NYTimes",
            "Economist",
            "WebMD",
            "BetterHelp",
            "Zocdoc",
            "Clevelandclinic",
            "Drugs.com",
            "EMedicineHealth.com",
            "Reddit",
        ),
    )
    src_key_word = "" if src_key_word == "Google" else src_key_word
    num = st.number_input(
        "Number of src to display", value=7, step=1, placeholder="Type a number..."
    )
    location = st.selectbox(
        "Where would you like to search from?",
        (
            "New York, New York, United States",
            "Los Angeles, California, United States",
            "Chicago, Illinois, United States",
            "Houston, Texas, United States",
            "Phoenix, Arizona, United States",
            "Philadelphia, Pennsylvania, United States",
            "San Antonio, Texas, United States",
            "San Diego, California, United States",
            "Dallas, Texas, United States",
            "San Jose, California, United States",
            "Austin, Texas, United States",
            "Jacksonville, Florida, United States",
            "Fort Worth, Texas, United States",
            "Columbus, Ohio, United States",
            "Charlotte, North Carolina, United States",
            "San Francisco, California, United States",
            "Indianapolis, Indiana, United States",
            "Seattle, Washington, United States",
            "Denver, Colorado, United States",
            "Washington, D.C., United States",
            "Boston, Massachusetts, United States",
            "El Paso, Texas, United States",
            "Nashville, Tennessee, United States",
            "Detroit, Michigan, United States",
            "Oklahoma City, Oklahoma, United States",
            "Portland, Oregon, United States",
            "Las Vegas, Nevada, United States",
            "Memphis, Tennessee, United States",
            "Louisville, Kentucky, United States",
            "Baltimore, Maryland, United States",
            "Milwaukee, Wisconsin, United States",
            "Albuquerque, New Mexico, United States",
            "Tucson, Arizona, United States",
            "Fresno, California, United States",
            "Sacramento, California, United States",
            "Mesa, Arizona, United States",
            "Kansas City, Missouri, United States",
            "Atlanta, Georgia, United States",
            "Omaha, Nebraska, United States",
            "Colorado Springs, Colorado, United States",
        ),
    )
    only_use_chatbot = st.checkbox("Only use chatbot.")

    # Add a button to clear the session state
    if st.button("Clear Session"):
        st.session_state.messages = []
        st.rerun()

    # Credit:
    current_year = current_year()  # This will print the current year
    st.markdown(
        f"""
            <h6 style='text-align: left;'>Copyright © 2010-{current_year} Present Yiqiao Yin</h6>
        """,
        unsafe_allow_html=True,
    )


# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Ensure messages are a list of dictionaries
if not isinstance(st.session_state.messages, list):
    st.session_state.messages = []
if not all(isinstance(msg, dict) for msg in st.session_state.messages):
    st.session_state.messages = []

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Front-end: Chat Panel - React to user input
if prompt := st.chat_input(
    "😉 Ask any question or feel free to use the examples provided in the left sidebar."
):
    # Display user message in chat message container
    st.chat_message("user").markdown(prompt)

    # Add user message to chat history
    st.session_state.messages.append(
        {
            "role": "system",
            "content": f"You are a helpful assistant. Year now is {current_year}",
        }
    )
    st.session_state.messages.append({"role": "user", "content": prompt})

    # API Call
    query = src_key_word + " " + prompt
    try:
        with st.spinner("Wait for it..."):
            if only_use_chatbot:
                md_data = "<empty>"
                response = "<empty>"
                ref_table_string = "<empty>"
            else:
                md_data = invoke_search_api(
                    api_key=API_KEY, query=query, location=location, num=num
                )
                md_data = md_data["data"]
                response = f"""
                Please see search results below: \n{md_data}
                """
                ref_table_string = response

            # API Call
            bot = ChatBot()
            bot.history = (
                st.session_state.messages.copy()
            )  # Update history from messages
            response = bot.generate_response(
                f"""
                Here's user prompt: {prompt}
                Here's relevant content: {response}

                Answer user prompt based on the relevant content.
                If relevant content is provided, it'd be a table with Title, Link, and Snippet. Use information from the Snippet.
                If relevant content is <empty>, use your best judgement or chat history."""
            )

    except Exception as e:
        st.warning(f"Failed to fetch data: {e}")
        response = "We are fixing the API right now. Please check back later."

    # Save audio
    if response:
        save_to_audio(response)

    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        st.markdown(response, unsafe_allow_html=True)
        st.audio("output.mp3", format="audio/mpeg", loop=True)
        with st.expander("See references:", expanded=False):
            st.markdown(ref_table_string)

    # Add assistant response to chat history
    final_response = f"""
    {response}

    {ref_table_string}"""
    st.session_state.messages.append({"role": "assistant", "content": final_response})