# %%
from suno_utils.audio import Audio

# %%
# spanish = "0f0d1134-4213-489e-b425-3e37eb202d5a.mp3"
# japanese = "0f0b67c3-b8d5-49ad-aa2c-268e3ae5d9c6.mp3"
# english = "0ee2dcc5-c2a6-4ca5-a7aa-1c4791af4008.mp3"
# audio = Audio.from_s3(
#     f"s3://suno-data/datasets/harvest/podcast_episodes/audio/{japanese}"
# ).get_segment(0, 60 * 5)
# audio.write_mp3("speech.mp3")

# %%
# audio.get_segment(0, 5 * 60).play()

# %%


# %%
import os
from openai import OpenAI

api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
    raise ValueError("OPENAI_API_KEY environment variable is required")

client = OpenAI(api_key=api_key)

# audio_file = open("speech.mp3", "rb")
# transcript = client.audio.transcriptions.create(
#     file=audio_file,
#     model="gpt-4o-mini-transcribe",
#     response_format="text",
#     # timestamp_granularities=["word"],
# )

# transcript


# %%
def detect_language(path):
    """Detect the language of an audio file using OpenAI's transcription and chat models."""
    # Load and segment the audio
    audio = Audio.from_file(path)
    duration = audio.duration_s
    audio = audio.get_segment(duration / 2, duration / 2 + 15)
    temp_file = f"/tmp/temp_language_detection_{path.split('/')[-1]}.mp3"
    audio.write_mp3(temp_file)

    try:
        # Transcribe a short segment
        with open(temp_file, "rb") as f:
            transcript = client.audio.transcriptions.create(
                file=f,
                model="whisper-1",
                response_format="verbose_json",
            )

        return transcript.language

    finally:
        # Clean up temp file
        import os

        if os.path.exists(temp_file):
            os.remove(temp_file)


# Test the function
# language = detect_language("/app2/suno/data/podcast/audio/00001797-db90-4f6d-bb63-1e20bc572c12.mp3")
# print(f"Detected language: {language}")


# %%
import json

metas = []
with open("/app2/suno/data/podcast/podcast_episodes.jsonl", "r") as f:
    for line in f:
        data = json.loads(line)
        metas.append(data)

print(len(metas))
metas[0]

# %%
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor
import threading

# Thread-safe print function
print_lock = threading.Lock()


def safe_print(message):
    with print_lock:
        print(message)


def process_meta(meta):
    s3_url = meta["s3_filepath"]
    local_path = f"/app2/suno/data/podcast/audio/{s3_url.split('/')[-1]}"
    try:
        language = detect_language(local_path)
        meta["detected_language"] = language
        # safe_print(f"Episode {meta['id']}: {language}")
        return meta
    except Exception as e:
        safe_print(f"Error processing {meta['id']}: {e}")
        meta["detected_language"] = None
        return meta


# Detect languages for all metas using threads
with ThreadPoolExecutor(max_workers=100) as executor:
    list(
        tqdm(
            executor.map(process_meta, metas),
            total=len(metas),
            desc="Detecting languages",
        )
    )


# save to jsonl
with open("/app2/suno/data/podcast/podcast_episodes_with_language.jsonl", "w") as f:
    for meta in metas:
        f.write(json.dumps(meta) + "\n")

# %%