from suno_utils.audio import Audio
from typing import List, Tuple, Dict, Any
import tempfile
import bisect
import numpy as np
import torch
from torch.nn import functional as F
import boto3
RNN_DOWNBEAT_MERT_MODEL_PATH = "s3://suno-data/m4burns/beat_this_orig.pt"
UPLOADS_KEY_PREFIX = "studio/uploads"
UPLOADS_BUCKET = "suno-data-uploads"
SAMPLE_RATE = 16000
WINDOW_SIZE = 512
from dotenv import load_dotenv
import os
load_dotenv() 

S3_CLIENT = boto3.client("s3")
# Download youtube meta json file  s3://suno-data/datasets/bundles/v0/youtube/metas.jsonl 
import json
def get_json_file():
    with tempfile.NamedTemporaryFile(suffix=".jsonl") as f:
        S3_CLIENT.download_fileobj("suno-data", Key="datasets/bundles/v0/youtube/metas.jsonl", Fileobj=f)
        # Flush and rewind to the beginning of the file
        f.flush()
        f.seek(0)
        metas = []
        for raw in f:
            line = raw.decode('utf-8').strip()
            if not line:
                continue
            try:
                metas.append(json.loads(line))
            except json.JSONDecodeError as e:
                # 根据需要记录或跳过错误行
                print(f"跳过解析错误行: {e}")
        return metas

def get_genius_file():
    #s3://suno-data/datasets/bundles/v0/youtube/audio/4a6be17f-fe1f-4b40-9bd9-21f8a8d5b442.mp3
    with tempfile.NamedTemporaryFile(suffix=".jsonl") as f:
        S3_CLIENT.download_fileobj("suno-data", Key="datasets/bundles/v0/genius/metas.jsonl", Fileobj=f)
        # Flush and rewind to the beginning of the file
        f.flush()
        f.seek(0)
        metas = []
        for raw in f:
            line = raw.decode('utf-8').strip()
            if not line:
                continue
            try:
                metas.append(json.loads(line))
            except json.JSONDecodeError as e:
                # 根据需要记录或跳过错误行
                print(f"跳过解析错误行: {e}")
        return metas
    
def get_discogs_json_file():
    #s3://suno-data/datasets/bundles/v3/discogs/metas.jsonl
    with tempfile.NamedTemporaryFile(suffix=".jsonl") as f:
        S3_CLIENT.download_fileobj("suno-data", Key="datasets/bundles/v3/discogs/metas.jsonl", Fileobj=f)
        # Flush and rewind to the beginning of the file
        f.flush()
        f.seek(0)
        metas = []
        for raw in f:
            line = raw.decode('utf-8').strip()
            if not line:
                continue
            try:
                metas.append(json.loads(line))
            except json.JSONDecodeError as e:
                print(f"skip parsing error line: {e}")
        return metas

import botocore.exceptions
def get_audio(gen_id):
    key = f"{UPLOADS_KEY_PREFIX}/{gen_id}.opus"
    with tempfile.NamedTemporaryFile(suffix=".opus") as f:
        try:
            S3_CLIENT.download_fileobj(Bucket=UPLOADS_BUCKET, Key=key, Fileobj=f)
            f.flush()
            return Audio.from_file(f.name, n_channels=1, sample_rate=SAMPLE_RATE)
        except (S3_CLIENT.exceptions.NoSuchKey, S3_CLIENT.exceptions.ClientError):
            print(f"Opus not found for {gen_id}, trying to download mp3...")
            with tempfile.NamedTemporaryFile(suffix=".mp3") as f:
                try:
                    S3_CLIENT.download_fileobj(UPLOADS_BUCKET, f"studio/uploads/{gen_id}.mp3", f)
                    prompt_audio = Audio.from_file(f.name, n_channels=1, sample_rate=SAMPLE_RATE)
                    return prompt_audio
                except botocore.exceptions.ClientError:
                    print("No .mp3 file found for", audio_id)
                    return None
        except Exception as e:
            print(f"Error getting opus for {gen_id}: {e}")
            with tempfile.NamedTemporaryFile(suffix=".mp3") as f:
                try:
                    S3_CLIENT.download_fileobj(UPLOADS_BUCKET, f"studio/uploads/{gen_id}.mp3", f)
                    prompt_audio = Audio.from_file(f.name, n_channels=1, sample_rate=SAMPLE_RATE)
                    return prompt_audio
                except botocore.exceptions.ClientError:
                    print("No .mp3 file found for", audio_id)
                    return None
            raise e