import streamlit as st
from transformers import HfApiEngine
from transformers.agents import ReactCodeAgent, ReactJsonAgent, ManagedAgent
from transformers.agents.search import DuckDuckGoSearchTool, VisitWebpageTool
from huggingface_hub import login

# Authenticate with Hugging Face
def authenticate_hf():
    hf_key = st.secrets["HF_API_KEY"]
    login(hf_key)

# Initialize agents
def initialize_agents():
    llm_model_path = "meta-llama/Meta-Llama-3-70B-Instruct"
    llm_engine = HfApiEngine(model=llm_model_path)

    # Agent 1: Extract song and artist
    parse_input_agent = ReactJsonAgent(tools=[], llm_engine=llm_engine)
    managed_parse_input_agent = ManagedAgent(
        agent=parse_input_agent,
        name="parse_input",
        description="Extracts the name of the song and the artist from the user input. If no artist is mentioned, it performs its best guess."
    )

    # Agent 2: Perform a web search
    find_lyrics_site_agent = ReactCodeAgent(
        tools=[DuckDuckGoSearchTool()], 
        llm_engine=llm_engine,
        additional_authorized_imports=['requests']
    )
    managed_find_lyrics_site_agent = ManagedAgent(
        agent=find_lyrics_site_agent,
        name="find_lyrics_site",
        description="Runs web searches to find the lyrics of a song. Give it your query as an argument. Returns the best website for the required lyrics."
    )

    # Agent 3: Extract lyrics from webpage
    web_parsing_agent = ReactCodeAgent(
        tools=[DuckDuckGoSearchTool(), VisitWebpageTool()], 
        llm_engine=llm_engine,
        additional_authorized_imports=['requests', 'bs4']
    )
    managed_web_parsing_agent = ManagedAgent(
        agent=web_parsing_agent,
        name="web_parsing_lyrics",
        description="Visits the website provided (DO NOT MODIFY THE URL) and retrieves lyrics from the song. Returns the extracted text. If it is not possible to clean the text, just return all the content of the webpage."
    )

    # Manager Agent
    manager_agent = ReactCodeAgent(
        tools=[], 
        llm_engine=llm_engine, 
        managed_agents=[
            managed_parse_input_agent, 
            managed_find_lyrics_site_agent, 
            managed_web_parsing_agent
        ]
    )

    return manager_agent

# Streamlit App
def main():
    st.title("Lyrics Finder App 🎵")
    st.write("Enter the name of the song and optionally the artist to retrieve the lyrics.")

    authenticate_hf()
    manager_agent = initialize_agents()

    # Input for user queries
    custom_query = st.text_input("Enter your song and artist (e.g., 'Search for Pelo Suelto by Elsa y Elmar')", "")

    if st.button("Find Lyrics"):
        if custom_query.strip() == "":
            st.error("Please enter a valid song and artist name!")
        else:
            st.write("Processing your request, please wait...")
            with st.spinner("Agents are working hard to find the lyrics..."):

                # Retry logic
                max_attempts = 5
                done = False
                attempts = 0
                final_lyrics = None

                while not done and attempts < max_attempts:
                    try:
                        result = manager_agent.run(custom_query)
                        final_lyrics = result["Task outcome (extremely detailed version)"]
                        done = True
                    except Exception as e:
                        st.warning(f"Attempt {attempts + 1}/{max_attempts} failed: {e}")
                        attempts += 1

                if final_lyrics:
                    st.success("Lyrics retrieved successfully!")

                    # Display lyrics in a larger text area
                    st.text_area("Lyrics Output", final_lyrics, height=400)

                    # Button to download lyrics as a text file
                    lyrics_file = f"{custom_query.replace(' ', '_')}_lyrics.txt"
                    st.download_button(
                        label="Download Lyrics as .txt",
                        data=final_lyrics,
                        file_name=lyrics_file,
                        mime="text/plain"
                    )
                else:
                    st.error("Failed to retrieve lyrics after multiple attempts.")

if __name__ == "__main__":
    main()