import streamlit as st
import librosa
import httpx  # Recommended library for efficient HTTP requests
import numpy as np
import os

# Replace with your actual API endpoint URL
API_URL = "https://jwry1smafa0ql7pc.us-east-1.aws.endpoints.huggingface.cloud"

# Retrieve the token from environment variables using `os.getenv()`
TOKEN = os.getenv("ACCESS_TOKEN")

# Set headers with the retrieved token
headers = {
    "Authorization": f"Bearer {TOKEN}",
}


def query(payload):
    """Sends a POST request to the API endpoint with the given payload."""
    response = httpx.post(API_URL, headers=headers, json=payload)
    return response.json()  # Use json() method for efficient data access


def main():
    """
    Creates the Streamlit app for uploading audio (optional),
    and sending it to an API for music generation.
    """
    st.title("SonglabAI")

    genres = ["Pop", "Rock", "Jazz", "Electronic", "Hip-Hop", "Classical", "Lofi", "Chillpop", "Bossa Nova"]
    selected_genre = st.selectbox("Select Genre", genres)
    prompt = st.text_input("Enter Music Description")
    
    # bpm
    bpm = st.number_input("Enter Speed in BPM", min_value=60)
    
    # Input form
    duration = st.number_input("Duration (in seconds)", min_value=20)
    
    uploaded_audio = None

    # Optional audio upload
    if st.checkbox("Upload Melody"):
        uploaded_file = st.file_uploader("Choose an audio file", type=["mp3", "wav"])
        if uploaded_file is not None:
            try:
                audio_bytes = uploaded_file.read()
                audio, sample_rate = librosa.load(audio_bytes, sr=None)  # Avoid resampling
                if sample_rate:
                    uploaded_audio = audio[:int(sample_rate * 15)]  # Extract first 15 seconds
                else:
                    st.warning("Could not determine sample rate. Skipping audio processing.")
            except Exception as e:
                st.error(f"Error processing audio: {e}")

    if st.button("Generate Music"):
        if not prompt:
            st.error("Please enter a music description.")
        elif selected_genre:
            st.text("Preparing request...")

            # Prepare the payload based on presence/absence of uploaded audio
            payload = {
                "inputs": f"{prompt} {selected_genre} {bpm} bpm",  # Combine prompt and genre
            }
            if uploaded_audio is not None:
                payload["audio_data"] = uploaded_audio.tolist()  # Convert to list for JSON

            try:
                # Send the request and handle the response
                response = query(payload)

                if "error" in response:
                    st.error(f"Error from API: {response['error']}")
                else:
                    generated_audio_data = response.get("generated_audio")
                    if generated_audio_data is not None:
                        st.success("Audio generation complete.")
                        st.audio(generated_audio_data, format="audio/wav")
                    else:
                        st.warning(
                            "Unexpected response structure. Missing 'generated_audio'."
                        )
            except httpx.HTTPError as e:
                st.error(f"HTTP Error: {e}")
            except Exception as e:
                st.error(f"Unexpected error: {e}")


if __name__ == "__main__":
    main()