# Import necessary libraries
import streamlit as st
import pandas as pd
import transformers
import torch
from huggingface_hub import login
from dotenv import load_dotenv
import os

# Import data
music_data = pd.read_csv("Spotify_Youtube.csv")

# Login to HuggingFace Hub
load_dotenv()
HUGGINGFACE_API_KEY = os.environ.get("HUGGINGFACE_API_KEY")
login(HUGGINGFACE_API_KEY)

# Load Meta LLaMA model
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto"
)

# Function to parse user input using Meta LLaMA model
def parse_user_input(user_input):
    messages = [
        {
            "role": "system",
            "content": """You will be provided with an input: '{user_input}', and your task is to determine the following:
            - Valence: a number that is equal to the mood. Positive moods are closer to 1 and negative moods are closer to 0.
            - Number of songs: the number of songs the user requests.
            - Tempo: the tempo of the songs.
            - Danceability: the danceability of the songs.

            Provide this information in the following format with each value separated by a space:
            'valence number_of_songs tempo danceability'
            Example: '0.5 20 120 0.8'
            """
        },
        {
            "role": "user",
            "content": user_input
        },
    ]

    prompt = pipeline.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    terminators = [
        pipeline.tokenizer.eos_token_id,
        pipeline.tokenizer.convert_tokens_to_ids("")
    ]

    outputs = pipeline(
        prompt,
        max_new_tokens=256,
        eos_token_id=terminators,
        do_sample=True,
        temperature=0.6,
        top_p=0.9,
    )
    return outputs[0]["generated_text"][len(prompt):]

# Function to create a new dataframe from the music dataframe based on valence, number of tracks, tempo, and danceability
def get_tracks_by_artist_and_danceability(music_data, valence, num_tracks, tempo, danceability):
    filtered_tracks = music_data[
        (music_data['Valence'].between(valence - 0.1, valence + 0.1)) &
        (music_data['Tempo'].between(tempo - 30, tempo + 30)) &
        (music_data['Danceability'].between(danceability - 0.2, danceability + 0.2))
    ]
    return filtered_tracks.head(num_tracks)[['Track', 'Artist']]

# Streamlit Application
logo = "music_logo.png"

# Sidebar
with st.sidebar:
    st.image(logo, width=100)
    st.header("Navigation")
tab_selection = st.sidebar.radio("Go to", ["Music Generator", "Browse Music", "About Us"])

# Music generator page
if tab_selection == "Music Generator":
    st.header("Mood Playlist Generator")
    st.write("Enter your music preferences in a detailed format and receive a personalized playlist based on your mood")
    user_prompt = st.text_input("Example: 'I want 20 happy songs with high tempo that I can dance to!'")

    if st.button("Generate Playlist"):
        try:
            with st.spinner("Processing your request..."):
                parsed_input = parse_user_input(user_prompt)
                # st.write(f"Parsed input: {parsed_input}")

                # Extract parameters from the parsed input
                valence, num_tracks, tempo, danceability = parsed_input.split()
                valence = float(valence)
                num_tracks = int(num_tracks)
                tempo = int(tempo)
                danceability = float(danceability)

                # st.write(f"Number of tracks: {num_tracks}, Valence: {valence}, Tempo: {tempo}, Danceability: {danceability}")

                tracks = get_tracks_by_artist_and_danceability(music_data, valence, num_tracks, tempo, danceability)
                # st.write(f"Found {len(tracks)} tracks.")

                if tracks.empty:
                    st.write("No tracks found. Please try a different query.")
                else:
                    st.write("Here are your recommended playlist:")
                    st.table(tracks)
                    st.button("Add playlist to Spotify")
        except ValueError:
            st.write("Error: Unable to parse the input. Please make sure the format is correct.")

# Browse music page
elif tab_selection == "Browse Music":
    st.header("Browse Music")
    st.write("Explore the music data used for generating your playlists.")
    df = pd.read_csv("Spotify_Youtube.csv")
    st.dataframe(df)

# About us page
elif tab_selection == "About Us":
    st.header("About Us")