import gradio as gr
import requests
from pydub import AudioSegment
import shazamio
import os
import time
token = os.environ.get('TOKEN')
# Replace with your actual client ID and client secret
client_id = os.environ.get('SOUNDCLOUD_CLIENT_ID')
client_secret = os.environ.get('SOUNDCLOUD_CLIENT_SECRET')

def download_audio(streaming_url, output_path, headers):
    response = requests.get(streaming_url, headers=headers, stream=True)
    with open(output_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)

async def identify_track(shazam, audio_chunk):
    temp_file = 'temp_chunk.wav'
    audio_chunk.export(temp_file, format='wav')
    result = await shazam.recognize_file(temp_file)
    if result and 'track' in result:
        track_data = result['track']
        return {
            'title': track_data['title'],
            'subtitle': track_data['subtitle']
        }
    else:
        return None

async def process_dj_set(track_url, progress=gr.Progress()):
    headers = {
        'Authorization': 'Bearer ' + f'{token}'
    }
    resolve_response = requests.get(f'https://api.soundcloud.com/resolve.json?url={track_url}', headers=headers)
    if resolve_response.status_code != 200:
        return "Failed to resolve track URL.", ""
    track_data = resolve_response.json()
    streaming_url = track_data['stream_url']
    download_audio(streaming_url, 'track.wav', headers)
    audio = AudioSegment.from_wav('track.wav')
    chunk_duration = 30000  # 30 seconds
    overlap = 10000  # 10 seconds
    chunks = []
    start = 0
    while start + chunk_duration < len(audio):
        end = start + chunk_duration
        chunk = audio[start:end]
        chunks.append((start, chunk))
        start += chunk_duration - overlap
    shazam = shazamio.Shazam()
    tracklist = []
    for start_time, chunk in chunks:
        progress(0.1)
        track_info = await identify_track(shazam, chunk)
        if track_info:
            timestamp = time.strftime("%M:%S", time.gmtime(start_time / 1000))
            tracklist.append(f"{timestamp} - {track_info['title']} by {track_info['subtitle']}")
    tracklist_output = "\n".join(tracklist)
    download_content = tracklist_output
    return tracklist_output, download_content

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# SoundCloud DJ Set Track Identifier")

        with gr.Row():
            track_url = gr.Text(
                label="SoundCloud DJ Set URL",
                show_label=False,
                max_lines=1,
                placeholder="Enter SoundCloud DJ set URL",
                container=False,
            )

            run_button = gr.Button("Process", scale=0, variant="primary")

        result = gr.Textbox(label="Tracklist", show_label=False)

        with gr.Accordion("Download Tracklist", open=False):
            download_button = gr.File(label="Download")

        gr.Examples(examples=["https://soundcloud.com/your-track-url"], inputs=[track_url])

    run_button.click(process_dj_set, inputs=track_url, outputs=[result, download_button])

if __name__ == "__main__":
    demo.launch()