import os
import time
import streamlit as st
from smolagents import CodeAgent, HfApiModel, tool
from huggingface_hub import InferenceClient
import requests
from bs4 import BeautifulSoup

# Retrieve Hugging Face token
hf_token = os.getenv("HF_TOKEN")
if not hf_token:
    raise ValueError("Hugging Face token not found. Please set it in the environment or Streamlit secrets.")

# Initialize the Hugging Face Inference client
client = InferenceClient(token=hf_token)

# Custom tools for SmolAgents
@tool
def search_harry_potter_lore(query: str) -> str:
    """Search for Harry Potter-related lore or facts across the entire Harry Potter Fandom site.
    
    Args:
        query: A specific question or topic about Harry Potter lore.
    
    Returns:
        A concise and informative response based on the query.
    """
    headers = {"User-Agent": "Mozilla/5.0"}
    # Construct the search URL for the Harry Potter Fandom site.
    search_url = f"https://harrypotter.fandom.com/wiki/Special:Search?query={query}"
    
    try:
        # Fetch the search results page.
        search_response = requests.get(search_url, headers=headers)
        if search_response.status_code != 200:
            return f"Error: Received status code {search_response.status_code} from search."
        
        search_soup = BeautifulSoup(search_response.text, 'html.parser')
        
        # Look for the first link that appears to be an article.
        article_url = None
        for a in search_soup.find_all("a", href=True):
            href = a["href"]
            # We want links that start with /wiki/ but skip those that contain "Special:"
            if href.startswith("/wiki/") and "Special:" not in href:
                article_url = "https://harrypotter.fandom.com" + href
                break
        
        if not article_url:
            return "No results found for your query."
        
        # Fetch the article page.
        article_response = requests.get(article_url, headers=headers)
        if article_response.status_code != 200:
            return f"Error: Received status code {article_response.status_code} from the article page."
        
        article_soup = BeautifulSoup(article_response.text, 'html.parser')
        
        # Extract the first meaningful paragraph.
        paragraphs = article_soup.find_all("p")
        for p in paragraphs:
            text = p.get_text().strip()
            if len(text) > 50:  # A simple threshold to ensure the paragraph is informative.
                return text
        
        return "Couldn't extract detailed lore from the article."
    
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Initialize the SmolAgent model
model = HfApiModel(model_id="HuggingFaceTB/SmolLM2-360M", token=hf_token)

# Create the agent
agent = CodeAgent(tools=[search_harry_potter_lore], model=model)

def apply_custom_styles():
    st.markdown(
        """
        <style>
            .stApp {
                background-color: #fefbe9;
                background-image: radial-gradient(circle, rgba(250,240,200,1) 0%, rgba(245,235,185,1) 100%);
                font-family: 'Dancing Script', cursive;
                display: flex;
                justify-content: center;
                align-items: center;
                height: 100vh;
            }
            textarea {
                width: 100%;
                height: 100%;
                font-size: 20px;
                font-family: 'Dancing Script', cursive;
                color: #000000;
                background-color: rgba(255, 255, 255, 0.8);
                border: none;
                box-shadow: none;
                outline: none;
                resize: none;
                padding: 10px;
            }
            textarea::placeholder {
                color: rgba(0, 0, 0, 0.5);
            }
            textarea:focus {
                outline: none;
                border: 2px solid rgba(0, 0, 0, 0.3);
            }
            .response {
                font-size: 22px;
                color: #000000;
                text-align: center;
                font-family: 'Dancing Script', cursive;
                padding: 20px;
                animation: fade-in 2s ease-in;
                position: relative;
            }
            @keyframes fade-in {
                from { opacity: 0; }
                to { opacity: 1; }
            }
        </style>
        """,
        unsafe_allow_html=True
    )

def main():
    apply_custom_styles()

    # Initialize session state variables if not set
    if "state_reset" not in st.session_state:
        st.session_state["state_reset"] = True

    if st.session_state["state_reset"]:
        st.session_state["user_input"] = ""
        st.session_state["response"] = None
        st.session_state["waiting_for_input"] = True
        st.session_state["state_reset"] = False

    def process_input():
        """Fetch Tom Riddle's response based on user input."""
        user_input = st.session_state["user_input"]
        if user_input.strip():
            try:
                response = agent.run(
                    f"You are Tom Riddle, a cunning and enigmatic character from Harry Potter (you're not evil yet). Respond concisely and pragmatically but please don't be evil and remain true to your role in the series: {user_input}"
                )
                st.session_state["response"] = response
                st.session_state["waiting_for_input"] = False

            except Exception as e:
                st.error(f"An error occurred: {e}")

    # Show input area or response based on state
    if st.session_state["waiting_for_input"]:
        st.text_area(
            "",
            placeholder="Write your thoughts here...",
            key="user_input",
            on_change=process_input,
        )
    elif st.session_state["response"]:
        st.markdown(
            f"<div class='response'>{st.session_state['response']}</div>",
            unsafe_allow_html=True,
        )
        # Wait for fade-out to complete
        time.sleep(8.5)
        # Fully reset state to start over
        st.session_state["state_reset"] = True
        st.rerun()

if __name__ == "__main__":
    main()