import gradio as gr
import requests
from pydub import AudioSegment
import shazamio
import os
import time

# Environment variables
token = os.environ.get('TOKEN')
client_id = os.environ.get('SOUNDCLOUD_CLIENT_ID')
client_secret = os.environ.get('SOUNDCLOUD_CLIENT_SECRET')

# Function to download audio
def download_audio(streaming_url, output_path, headers):
    try:
        response = requests.get(streaming_url, headers=headers, stream=True)
        response.raise_for_status()
        with open(output_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        return "Audio downloaded successfully."
    except Exception as e:
        return f"Error downloading audio: {e}"

# Function to identify track using Shazam
async def identify_track(shazam, audio_chunk):
    try:
        temp_file = 'temp_chunk.wav'
        audio_chunk.export(temp_file, format='wav')
        
        # Ensure temp_file exists and is not a directory
        if not os.path.isfile(temp_file):
            raise FileNotFoundError(f"Temporary file '{temp_file}' does not exist or is not a valid file.")
        
        result = await shazam.recognize_file(temp_file)
        os.remove(temp_file)  # Clean up temporary file
        if result and 'track' in result:
            track_data = result['track']
            return {
                'title': track_data['title'],
                'subtitle': track_data['subtitle']
            }
        else:
            return None
    except Exception as e:
        return f"Error identifying track: {e}"

# Main processing function
async def process_dj_set(track_url, progress=gr.Progress()):
    try:
        status = "Resolving SoundCloud URL..."
        headers = {'Authorization': f'Bearer {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.get('stream_url')
        
        if not streaming_url:
            return "Streaming URL not found.", ""
        
        status = "Downloading audio..."
        download_status = download_audio(streaming_url, 'track.wav', headers)
        
        if "Error" in download_status:
            return download_status, ""
        
        status = "Processing audio..."
        
        # Ensure 'track.wav' exists and is not a directory
        if not os.path.isfile('track.wav'):
            raise FileNotFoundError("Downloaded audio file 'track.wav' does not exist or is not valid.")
        
        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 isinstance(track_info, str) and "Error" in track_info:  # Check for errors in identification
                return track_info, ""
            
            if track_info:
                timestamp = time.strftime("%M:%S", time.gmtime(start_time / 1000))
                tracklist.append(f"{timestamp} - {track_info['title']} by {track_info['subtitle']}")
        
        status = "Generating output..."
        tracklist_output = "\n".join(tracklist)
        
        return tracklist_output, tracklist_output
    except Exception as e:
        return f"An error occurred during processing: {e}", ""

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_input = 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_output = 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_input])

    run_button.click(process_dj_set, inputs=track_url_input, outputs=[result_output, download_button])

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