import asyncio
from shazamio import Shazam
import json
import time
import ffmpeg

mix = "The Magician b2b A-Trak - Live @ Elsewhere, Brooklyn [1739480907].opus"
output_file = ".".join(mix.split(".")[:-1]) + ".txt"
segment_length = 15
start_from = 39


def convert_seconds_to_hhmmss(seconds):
    time_struct = time.gmtime(seconds)
    return time.strftime("%H:%M:%S", time_struct)

def get_audio_length(filename):
    probe = ffmpeg.probe(filename)
    duration = float(probe['format']['duration'])
    return duration

def extract_audio_segment(input_file, start_time, duration, output_file):
    (
        ffmpeg
        .input(input_file, ss=start_time, t=duration)
        .output(output_file, acodec="copy")
        .run(overwrite_output=True, quiet=True)
    )

async def main():
    with open(output_file, "a") as of:
        shazam = Shazam()

        audio_length = get_audio_length(mix)
        print(f"Got total mix length: {audio_length}")

        total_segments = int(audio_length // segment_length)
        print(f"Splitting into {total_segments} segments of {segment_length}s")
        n_digits = len(str(total_segments))

        probe = ffmpeg.probe(mix)
        audio_codec = next(stream for stream in probe['streams'] if stream['codec_type'] == 'audio')['codec_name']
        print(f"Got audio codec: {audio_codec}" + "\n")

        segment = 0


        for s in range(0, int(audio_length), segment_length):
            if segment >= start_from:
                line = convert_seconds_to_hhmmss(s) + "\t"
                extract_audio_segment(mix, s, segment_length, "sample.opus")

                out = await shazam.recognize('sample.opus')
                if out.get("track"):
                    name = out["track"]["share"]["subject"]

                    line += name
                    print(f"[{str(int(s/segment_length)).zfill(n_digits)}/{str(total_segments).zfill(n_digits)}]  " + line)

                    of.write(line + "\n")
                    of.flush()

                time.sleep(3)
            segment += 1

        of.close()

print(f"Processing {mix}...")

loop = asyncio.get_event_loop()
loop.run_until_complete(main())