import streamlit as st
import os
import requests
import io
from PIL import Image
from IPython.display import Audio, display
from freeGPT import Client

api_token = os.environ.get("API_TOKEN")
API_URL = "https://api-inference.huggingface.co/models/facebook/musicgen-small"
API_URL_IMG = "https://api-inference.huggingface.co/models/goofyai/3d_render_style_xl"
headers = {"Authorization": f"Bearer {api_token}"}
st.title("✨ AI Tracks Generator")
st.write("Generate audio, cover and title for your track using ai tools!")
prompt = st.text_input("Enter prompt and generate track", placeholder="Eg. Nice holiday love melody")
generate_btn = st.button("✨ Generate")

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content

def imgquery(payload):
    response = requests.post(API_URL_IMG, headers=headers, json=payload)
    return response.content

def generate_audio(prompt):
    return query({"inputs": prompt})

def generate_image(prompt):
    image_bytes = imgquery({"inputs": prompt})
    return Image.open(io.BytesIO(image_bytes))

if generate_btn:
    audio1 = generate_audio(prompt)
    audio2 = generate_audio(prompt)
    audio3 = generate_audio(prompt)

    title1 = Client.create_completion("gpt3", "generate name of song that have this description. In answer give only short title" + prompt)
    title2 = Client.create_completion("gpt3", "generate name of song that have this description. In answer give only short title" + prompt)
    title3 = Client.create_completion("gpt3", "generate name of song that have this description. In answer give just short title" + prompt)

    cover1 = generate_image("Generate a cover image for this song:" + prompt)
    cover2 = generate_image("Generate a cover for this song:" + prompt)
    cover3 = generate_image("Generate a thumbnail for this song:" + prompt)

    col1, col2, col3 = st.columns(3)

    with col1:
        st.image(cover1)
        st.header(title1)
        st.audio(audio1, format="audio/wav")
        st.download_button("Download Audio 1", audio1)

    with col2:
        st.image(cover2)
        st.header(title2)
        st.audio(audio2, format="audio/wav")
        st.download_button("Download Audio 2", audio2)

    with col3:
        st.image(cover3)
        st.header(title3)
        st.audio(audio3, format="audio/wav")
        st.download_button("Download Audio 3", audio3)