import io
import json
import os
import time
import uuid
from datetime import datetime
from io import BytesIO
from tempfile import NamedTemporaryFile
from xmlrpc.client import Binary
import jwt

import numpy as np
import requests
import scipy
import soundfile as sf
import streamlit as st
import streamlit_vertical_slider as svs
from pydub import AudioSegment
from scipy.signal import butter, sosfilt
from streamlit import session_state as st_state
from woocommerce import API
from wordpress_xmlrpc import Client
from wordpress_xmlrpc.compat import xmlrpc_client
from wordpress_xmlrpc.methods import media


# Try to get API_URL from environment variables, if not found set to a default value
try:
    API_URL = os.environ["API_URL"]
except KeyError:
    st.error("API_URL environment variable is not set.")
    st.stop()

# Try to get the Bearer token from environment variables, if not found set to a default value
try:
    BEARER_TOKEN = os.environ["BEARER_TOKEN"]
except KeyError:
    st.error("BEARER_TOKEN environment variable is not set.")
    st.stop()

print("API_URL:", os.environ["API_URL"])
print("BEARER_TOKEN:", os.environ["BEARER_TOKEN"])

page_bg_img = '''
<style>
.stApp {
background-image: url("https://songlabai.com/wp-content/uploads/2024/03/4.png");
background-size: cover;
}
</style>
'''
st.markdown(page_bg_img, unsafe_allow_html=True)

def save_to_wordpress(channel1, channel2, sample_rate):

    channel1 = np.array(channel1).astype(np.float32)
    channel2 = np.array(channel2).astype(np.float32)
    stereo_audio = np.column_stack((channel1, channel2))

    # Define your WordPress site URL and authentication credentials
    wordpress_url = "https://songlabai.com/xmlrpc.php"
    woocommerce_url = "https://songlabai.com"
    consumer_key = "ck_93d516ba12289a6fd0eced56bbc0b05ecbf98735"
    consumer_secret = "cs_9d5eb716d631db408a4c47796b5d18b0313d8559"
    username = "admin_h2ibbgql"
    password = "um^VdaNK0H8Vw2*KNJlYABkh"

    # Authenticate with WordPress XML-RPC API
    wav_bytes = BytesIO()
    sf.write(wav_bytes, stereo_audio, samplerate=sample_rate, format="WAV")
    title = f"generated_audio_{datetime.now().timestamp()}.wav"
    file_data = {
        "name": title,
        "type": "audio/x-wav",  # Change the MIME type according to your file type
        "bits": xmlrpc_client.Binary(wav_bytes.getvalue()),
    }
    wp_client = Client(wordpress_url, username, password)
    for _ in range(4):
        try:
            # Upload the file to WordPress Media Library

            media_response = wp_client.call(media.UploadFile(file_data))

            # Handle the response
            if media_response:
                print(
                    "File successfully uploaded to WordPress with attachment ID:",
                    media_response,
                )

                # Create product data for WooCommerce
                product_data = {
                    "status": "pending",
                    "name": title,
                    "type": "simple",
                    "regular_price": "1.00",  # Set the price as needed
                    "sku": str(uuid.uuid4()),
                    "downloadable": True,
                    "download_limit": -1,
                    "download_expiry": -1,
                }

                # Authenticate with WooCommerce API
                wc_api = API(
                    url=woocommerce_url,
                    consumer_key=consumer_key,
                    consumer_secret=consumer_secret,
                    version="wc/v3",
                )

                # Create the product
                response = wc_api.post("products", product_data)

                # Handle the response
                if response.status_code == 201:
                    print(
                        "Product successfully created in WooCommerce:", response.json()
                    )
                    # Update product to add downloadable file URL
                    product_update_data = {
                        "downloads": [
                            {
                                "name": media_response["title"],
                                "file": media_response["link"],
                            }
                        ]
                    }
                    product_id = response.json().get("id")
                    response = wc_api.put(f"products/{product_id}", product_update_data)

                    if response.status_code == 200:
                        print(
                            "Downloadable file URL added to product:", response.json()
                        )
                        return (
                            response.json()["permalink"],
                            response.json()["permalink"].split("p=")[-1],
                        )
                    else:
                        print(
                            "Error adding downloadable file URL to product:",
                            response.text,
                        )
                else:
                    print("Error creating product in WooCommerce:", response.text)
            else:
                print("Error uploading file to WordPress.")
            break
        except Exception as e:
            print("Error:", e)


headers = {
    "Authorization": f"Bearer {BEARER_TOKEN}",
    "Content-Type": "application/json",
}

# Streamlit app title
st.title("Songlab AI")

# Initialize session state variables
if "vocal_audio" not in st_state:
    st_state.vocal_audio = None
if "vocal_sample_rate" not in st_state:
    st_state.vocal_sample_rate = None

if "audio" not in st_state:
    st_state.audio = None
if "audio_pydub" not in st_state:
    st_state.audio_pydub = None
if "audio_sample_rate" not in st_state:
    st_state.audio_sample_rate = None

if "augmented_audio" not in st_state:
    st_state.augmented_audio = None
if "augmented_audio_pydub" not in st_state:
    st_state.augmented_audio_pydub = None
if "augmented_audio_sample_rate" not in st_state:
    st_state.augmented_audio_sample_rate = None


genres = [
    "Pop",
    "Rock",
    "Hip Hop",
    "Jazz",
    "Blues",
    "Country",
    "Classical",
    "Electronic",
    "Reggae",
    "Folk",
    "R&B",
    "Metal",
    "Punk",
    "Indie",
    "Dance",
    "World",
    "Gospel",
    "Soul",
    "Funk",
    "Ambient",
    "Techno",
    "Disco",
    "House",
    "Trance",
    "Dubstep",
]

genre = st.selectbox("Select Genre:", genres)
energy_levels = ["Low", "Medium", "High"]
energy_level = st.radio("Energy Level:", energy_levels, horizontal=True)
description = st.text_input("Description:", "")
tempo = st.slider("Tempo (in bpm):", min_value=40, max_value=100, value=60, step=5)

# Duration input
duration = st.slider(
    "Duration (in seconds):", min_value=15, max_value=300, value=30, step=1
)

prompt2 = "Classical ,Low, 40, slow calm"


def convert_audio_segment_to_float_array(audio_pydub):
    """
    Convert a pydub AudioSegment to a NumPy array of type float32.

    Args:
    audio_pydub (AudioSegment): The AudioSegment object to be converted.

    Returns:
    np.ndarray: A NumPy array containing the audio data as float32.
    """
    # Get the raw audio data as a sequence of samples
    samples = audio_pydub.get_array_of_samples()

    # Convert the samples to a NumPy array and normalize to float32
    audio_array = np.array(samples).astype(np.float32)

    # Normalize the audio array to range between -1.0 and 1.0
    max_val = 2**15  # Assuming 16-bit audio, modify this if using different bit depths
    audio_array /= max_val
    return audio_array

def time_post_request(api_url, headers=None, payload=None):
    """
    Times the execution of a POST request.

    Parameters:
    - api_url (str): The URL to which the POST request is sent.
    - headers (dict): The headers to include in the POST request.
    - payload (dict): The payload to include in the POST request.

    Returns:
    - response (requests.Response): The response object returned by the POST request.
    - execution_time (float): The time it took to execute the POST request.
    """
    start_time = time.time()
    response = requests.post(api_url, headers=headers, json=payload)
    end_time = time.time()

    execution_time = end_time - start_time
    print(f"Execution time: {execution_time} seconds")

    return response

def generate_audio(genre, energy_level, tempo, description, duration):
    count_down = 300
    prompt = f"Genre: {genre}, Energy Level: {energy_level}, Tempo: {tempo}, Description: {description},"
    payload = {"inputs": {"prompt": prompt, "duration": duration}}
    with st.spinner("Generating audio ..."):
        response = time_post_request(API_URL, headers, payload)
    placeholder1 = st.empty()
    if response.status_code != 200:
        # Temp Jay Start
        print(f"Request Headers: {response.request.headers}")
        print(f"Request Body: {response.request.body}")
        print(f"Response Code: {response.status_code}")
        print(f"Response Body: {response.text}")
        # Temp Jay End

        
        print(str(response.content))
        count_down_placeholder = st.empty()
        for seconds in range(count_down):
            count_down_placeholder.info(
                f"The server is currently loading. Retrying in ⏳ {count_down - seconds} seconds."
            )
            time.sleep(1)
        count_down_placeholder.empty()
        with st.spinner("Generating audio ..."):
            response = time_post_request(API_URL, headers, payload)
        if response.status_code != 200:
            count_down_placeholder = st.empty()
            for seconds in range(count_down):
                count_down_placeholder.info(
                    f"2nd attempt failed. trying again one more time in ⏳ {count_down - seconds} seconds."
                )
                time.sleep(1)
            count_down_placeholder.empty()
            with st.spinner("Generating audio ..."):
                response = time_post_request(API_URL, headers, payload)
            if response.status_code != 200:
                placeholder1.error(
                    "Failed to generate audio after multiple attempts. Please try again later."
                )
            else:
                placeholder1.success(f"✔️ Audio Generated, Loading...")
                load_and_play_generated_audio(response)
                placeholder1.empty()
        else:
            placeholder1.success(f"✔️ Audio Generated, Loading...")
            load_and_play_generated_audio(response)
            placeholder1.empty()
    else:
        placeholder1.success(f"✔️ Audio Generated, Loading...")
        load_and_play_generated_audio(response)
        placeholder1.empty()


def load_and_play_generated_audio(response):
    
    response_eval = json.loads(response.content)
    channel1 =  response_eval[0]["generated_audio"][0]
    channel2 =  response_eval[0]["generated_audio"][1]
    with NamedTemporaryFile() as f1:
        scipy.io.wavfile.write(
            f1.name, rate=32000, data=np.array(channel1).astype(np.float32)
        )
        channel1_temp = AudioSegment.from_file(f1.name)
        with NamedTemporaryFile() as f2:
            scipy.io.wavfile.write(
                f2.name, rate=32000, data=np.array(channel2).astype(np.float32)
            )
            channel2_temp = AudioSegment.from_file(f2.name)
            audio_pydub = AudioSegment.from_mono_audiosegments(channel1_temp, channel2_temp)
            st_state.audio_pydub = audio_pydub
            st_state.audio_pydub_sample_rate = audio_pydub.frame_rate
    
    st_state.audio = np.array(st_state.audio_pydub.get_array_of_samples()).astype(np.float32)
    sample_rate = response_eval[0]["sample_rate"]
    perm_link, product_code = save_to_wordpress(
        channel1, channel2, 32000
    )
    
    st_state.audio_sample_rate = sample_rate
    st_state.augmented_audio_pydub = st_state.audio_pydub
    st.audio(st_state.audio_pydub.export().read())
    col_btn, col_text = st.columns([2,4])
    
    with col_btn:
        st.link_button("Publish your Song", "https://songlabai.com/contact_us/")
        st.link_button("Download Song", "https://songlabai.com/download-music/")
        
    with col_text:
        st.write(
            f"To Publish, please contact the admin by sending the following link: {perm_link}"
        )
        st.write(f"To download use the following product code: {product_code}")


if st.button("Generate Audio"):
    if genre and energy_level and description and tempo:
        generate_audio(genre, energy_level, tempo, description, duration)
        
    if description == "":
        st.info("Description field is required.")

# Post-processing options
st.header("Post-processing Options")

vocal_file = st.file_uploader(
    "Upload Vocal File", type=["mp3", "wav", "ogg", "flac", "aac"]
)
if vocal_file:
    st_state.vocal_audio = vocal_file.read()
    # st.audio(st_state.vocal_audio, format="audio/wav")

# Mixing
mix_vocals = st.checkbox("Mix Vocals")

if mix_vocals and st_state.vocal_audio is not None:
    with NamedTemporaryFile() as f:
        f.write(st_state.vocal_audio)
        st_state.vocal_audio = AudioSegment.from_file(f.name)
    st_state.augmented_audio_pydub = st_state.augmented_audio_pydub.overlay(
        st_state.vocal_audio, position=100
    )
    # st.audio(st_state.augmented_audio_pydub.export().read())
    st_state.augmented_audio = convert_audio_segment_to_float_array(
        st_state.augmented_audio_pydub
    )
    st_state.augmented_audio_sample_rate = st_state.augmented_audio_pydub.frame_rate
elif not mix_vocals and st_state.vocal_audio is not None:
    st_state.augmented_audio_pydub = st_state.audio_pydub
    st_state.augmented_audio_sample_rate = st_state.audio_pydub.frame_rate
# Mastering
st.header("Mastering")
st.markdown("")
# Volume Balance, Compression Ratio, and Reverb Amount
vol_col, pitch_shift,buttons_col = st.columns([2, 2, 2.5,])
with buttons_col:
    with st.container(height=371, border=True):
        st.markdown("")
        apply_stereo = st.button("Apply Stereo Effect")
        st.markdown("")
        reverse = st.button("Apply Audio Reverse ")
        st.markdown("")
        reset_post_processing = st.button("Undo All Post-processings")
        st.markdown("")

with vol_col:
    with st.container(border=True):
        volume_balance = svs.vertical_slider(
            "Volume Balance",
            min_value=-10.0,
            max_value=10.0,
            default_value=0.0,
            step=0.1,
            slider_color="green",
            track_color="lightgray",
            thumb_color="red",
            thumb_shape="pill",
        )
        vol_button = st.button("Apply Vol-Balance")


# Pitch shifting
with pitch_shift:
    with st.container(border=True):
        pitch_semitones = svs.vertical_slider(
            label="Pitch (semitones)",
            min_value=-12,
            max_value=12,
            default_value=0,
            step=1,
            slider_color="red",
            track_color="lightgray",
            thumb_color="red",
            thumb_shape="pill",
        )
        pitch_shift_button = st.button(
            "Apply Pitch Shift",
        )


if st_state.augmented_audio_pydub is not None:
    if vol_button:
        st_state.augmented_audio_pydub = st_state.augmented_audio_pydub + volume_balance

    if apply_stereo:
        st_state.augmented_audio_pydub = st_state.augmented_audio_pydub.pan(
            -0.5
        ).overlay(st_state.augmented_audio_pydub.pan(0.5))
    if reverse:
        st_state.augmented_audio_pydub = st_state.augmented_audio_pydub.reverse()
    if pitch_shift_button:
        st_state.augmented_audio_pydub = st_state.augmented_audio_pydub._spawn(
            st_state.augmented_audio_pydub.raw_data,
            overrides={
                "frame_rate": int(
                    st_state.augmented_audio_pydub.frame_rate
                    * (2 ** (pitch_semitones / 12.0))
                )
            },
        )

# Display the final audio
if st_state.augmented_audio_pydub is not None:
    st.audio(st_state.augmented_audio_pydub.export().read())
    # sample_rate = st_state.augmented_audio_sample_rate
    # st.audio(st_state.augmented_audio, format="audio/wav", sample_rate=sample_rate*2, start_time=0)

st.link_button(
    label="⬇️ Download/Save",
    url="https://songlabai.com/subcribe/",
    type="primary",
    use_container_width=True,
)
# m = st.markdown("""
# <style>
# div.stButton > button:first-child {
# # color: rgb(204, 49, 49);
# # background-color: ;
# # border-radius: 5%;
# backgroud-color: #00ff00;
# }
# # </style>""", unsafe_allow_html=True)

if reset_post_processing and st_state.audio_pydub is not None:
    st_state.augmented_audio_pydub = st_state.audio_pydub