import os
import gc
import time
import json
import glob
import random
import torch
import funcy
import queue
import threading
import torchaudio
import numpy as np

from suno_utils.utils.s3 import read_from_s3
from suno_utils.utils.text import read_jsonl
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from typing import List, Iterator, Dict, Any, Optional
from dataclasses import dataclass
from tqdm import tqdm
from dac.model.dac4 import DAC


# load the semantic model
from suno_utils.tasks.mert_25 import (
    preload_models as preload_semantic_models,
    encode as encode_semantic,
    EMBEDDING_RATE as SEMANTIC_HZ,
)

print("loading semantic model...")
semantic_model_filepath = "s3://suno-data/georg/models/semantic/mert_25.pt"
semantic_clusters_filepath = "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy"
_ = preload_semantic_models(
    checkpoint_filepath=semantic_model_filepath,
    centroids_filepath=semantic_clusters_filepath,
    device="cuda",
)


def corrupt_audio(audio, sr):
    audio_corrupt = audio.clone()
    # 10% chance to add highpass
    if np.random.uniform() < 0.25:
        hpf_cutoff = np.random.uniform(20, 1000)
        audio_corrupt = torchaudio.functional.highpass_biquad(
            audio_corrupt, sr, hpf_cutoff
        )
    # 10% chance to add lowpass
    if np.random.uniform() < 0.25:
        lpf_cutoff = np.random.uniform(1000, 20000)
        audio_corrupt = torchaudio.functional.lowpass_biquad(
            audio_corrupt, sr, lpf_cutoff
        )

    # 10% chance to make mono
    if np.random.uniform() < 0.7:
        audio_corrupt = audio_corrupt.mean(dim=0, keepdim=True)
        audio_corrupt = audio_corrupt.repeat(2, 1)

    # 10% chance to add noise
    if np.random.uniform() < 0.2:
        noise_level_db = -np.random.uniform(42, 96)
        noise_signal = torch.randn_like(audio_corrupt) * 10 ** (noise_level_db / 20)
        # apply filter to noise signal
        lpf_cutoff = np.random.uniform(1000, 20000)
        noise_signal = torchaudio.functional.lowpass_biquad(
            noise_signal, sr, lpf_cutoff
        )
        audio_corrupt = audio_corrupt + noise_signal

    # 5% chance to add contrast
    if np.random.uniform() < 0.05:
        contrast = np.random.uniform(0, 100.0)
        audio_corrupt = torchaudio.functional.contrast(
            audio_corrupt, enhancement_amount=contrast
        )

    # 10% add codec artifacts
    if np.random.uniform() < 1.0:
        compression = np.random.choice([32, 64, 128])
        audio_corrupt = torchaudio.functional.apply_codec(
            audio_corrupt, sr, "mp3", compression
        )

    # 10% chance to add distortion
    if np.random.uniform() < 0.3:
        drive_db = np.random.uniform(6, 24)
    else:
        drive_db = 0

    drive_lin = 10 ** (drive_db / 20)
    audio_corrupt *= drive_lin

    if np.random.uniform() < 0.5:
        audio_corrupt = torch.tanh(audio_corrupt)
    else:
        audio_corrupt = audio_corrupt.clamp(-1, 1)

    return audio_corrupt


def load_audio(meta):
    """Just load the audio file - I/O bound operation"""
    try:
        meta_id = meta["id"]

        # skip if npz file exists
        # npz_filepath = os.path.join(out_dir, f"{meta_id}.npz")
        # if os.path.exists(npz_filepath):
        #     return None

        start_time = time.time()
        audio, sr = read_from_s3(meta["audio_filepath"], read_f=torchaudio.load)
        print(f"loaded {meta['id']} in {time.time() - start_time:.2f} seconds")

        # crop to max 60s
        audio = audio[:, : int(48000 * 60.01)]

        # ensure stereo
        if audio.shape[0] == 1:
            audio = audio.repeat(2, 1)
        elif audio.shape[0] > 2:
            audio = audio[:2, :]

        # resample to 48k
        audio = torchaudio.functional.resample(audio, sr, 48000)

        # crop to first 30 seconds
        audio_48k = audio[:, : int(48000 * 30.01)]

        # Apply highpass filter
        audio_48k_corrupted = corrupt_audio(audio_48k, 48000)

        # Normalize for VAE
        audio_48k_normalized = audio_48k  # / audio_48k.abs().max().clamp(min=1e-5)
        audio_48k_corrupted_normalized = (
            audio_48k_corrupted  # / audio_48k_corrupted.abs().max().clamp(min=1e-5)
        )

        # Create 24k versions for semantic
        audio_24k = torchaudio.functional.resample(audio_48k, 48000, 24000)
        audio_24k_corrupted = torchaudio.functional.resample(
            audio_48k_corrupted, 48000, 24000
        )

        audio_24k_mono = audio_24k.mean(dim=0).unsqueeze(0)
        audio_24k_corrupted_mono = audio_24k_corrupted.mean(dim=0).unsqueeze(0)

        return {
            "id": meta["id"],
            "audio_48k": audio_48k_normalized,
            "audio_48k_corrupted": audio_48k_corrupted_normalized,
            "audio_24k_mono": audio_24k_mono,
            "audio_24k_corrupted_mono": audio_24k_corrupted_mono,
        }

    except Exception as e:
        print(f"Error loading {meta['id']}: {str(e)}")
        return None


def process_batch(batch_data, codec_model, out_dir, num_vae_tokens: int):
    """Process a batch through models and save results"""
    if not batch_data:
        return

    # Stack for VAE processing
    original_batch = torch.stack([data["audio_48k"] for data in batch_data]).cuda()
    corrupted_batch = torch.stack(
        [data["audio_48k_corrupted"] for data in batch_data]
    ).cuda()

    # Process through VAE
    with torch.no_grad():
        vae_original = codec_model.encode(original_batch)["z"].cpu()
        vae_corrupted = codec_model.encode(corrupted_batch)["z"].cpu()

    # Clear GPU memory
    del original_batch, corrupted_batch
    torch.cuda.empty_cache()

    # Process semantic codes
    semantic_original = encode_semantic([d["audio_24k_mono"] for d in batch_data])[0]
    semantic_corrupted = encode_semantic(
        [d["audio_24k_corrupted_mono"] for d in batch_data]
    )[0]

    # Save results
    for idx, data in enumerate(batch_data):
        vae_orig = vae_original[idx, :, :num_vae_tokens]
        vae_corrupt = vae_corrupted[idx, :, :num_vae_tokens]

        np.savez(
            os.path.join(out_dir, f"{data['id']}.npz"),
            vae_latents_original=vae_orig,
            vae_latents_corrupted=vae_corrupt,
            semantic_codes=semantic_original.astype(np.uint16)[:, 0],
            semantic_codes_corrupted=semantic_corrupted.astype(np.uint16)[:, 0],
        )
        # print(f"saved {len(batch_data)} files")

    # Clear memory
    del vae_original, vae_corrupted, semantic_original, semantic_corrupted
    torch.cuda.empty_cache()


def process_meta_chunk(metas: List[dict], results_queue: queue.Queue):
    """
    Downloads all audio files for a chunk of metas.
    This runs in its own thread for each chunk.
    """
    chunk_results = []
    for meta in metas:
        try:
            result = load_audio(meta)
            if result is not None:
                chunk_results.append(result)
        except Exception as e:
            print(f"Error processing {meta['id']}: {str(e)}")

    if chunk_results:
        results_queue.put(chunk_results)


def process_chunks_parallel(
    metas: List[dict],
    codec_model,
    out_dir: str,
    num_threads: int = 16,
    chunk_size: int = 100,
    batch_size: int = 8,
):
    """
    Main processing function that coordinates threaded downloads and GPU processing.
    """
    os.makedirs(out_dir, exist_ok=True)
    results_queue = queue.Queue()

    # Split metas into chunks
    meta_chunks = [metas[i : i + chunk_size] for i in range(0, len(metas), chunk_size)]

    # Keep track of active threads and completed chunks
    active_threads = []
    chunks_completed = 0
    total_chunks = len(meta_chunks)

    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        chunk_idx = 0

        while chunks_completed < total_chunks:
            # Start new download threads if we have chunks remaining and thread slots available
            while chunk_idx < len(meta_chunks) and len(active_threads) < num_threads:
                thread = executor.submit(
                    process_meta_chunk, meta_chunks[chunk_idx], results_queue
                )
                active_threads.append(thread)
                chunk_idx += 1

            # Check for completed threads
            active_threads = [t for t in active_threads if not t.done()]

            # Process any available results without blocking
            try:
                while True:  # Process all available results
                    chunk_results = results_queue.get_nowait()
                    print(f"\nProcessing chunk of {len(chunk_results)} files on GPU")

                    # Process in batches
                    for i in range(0, len(chunk_results), batch_size):
                        batch = chunk_results[i : i + batch_size]
                        process_batch(batch, codec_model, out_dir, num_vae_tokens=750)

                    chunks_completed += 1
                    print(f"Completed {chunks_completed}/{total_chunks} chunks")

                    # Clear GPU memory after processing chunk
                    torch.cuda.empty_cache()

            except queue.Empty:
                # No results available right now, continue with thread management
                pass

            # Small sleep to prevent tight loop
            time.sleep(0.1)

    # Process any remaining results in the queue
    while not results_queue.empty():
        chunk_results = results_queue.get()
        print(f"\nProcessing final chunk of {len(chunk_results)} files on GPU")

        for i in range(0, len(chunk_results), batch_size):
            batch = chunk_results[i : i + batch_size]
            process_batch(batch, codec_model, out_dir, num_vae_tokens=750)

        chunks_completed += 1
        print(f"Completed {chunks_completed}/{total_chunks} chunks")
        torch.cuda.empty_cache()


def main(
    out_dir: str,
    num_threads: int = 16,
    batch_size: int = 8,
    num_samples: int = 10000,
):
    VAE_RATE_HZ = 25

    # load the vae model
    device = "cuda:0"
    # device = "cpu"
    if VAE_RATE_HZ == 100:
        checkpoint_filepath = "s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth"
        # checkpoint_filepath = "/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/dac/weights.pth"
        num_vae_tokens = 3000
    elif VAE_RATE_HZ == 25:
        checkpoint_filepath = "s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth"
        num_vae_tokens = 750
    else:
        raise ValueError(f"VAE_RATE_HZ must be 100 or 25, got {VAE_RATE_HZ}")
    load_f = funcy.partial(torch.load, map_location="cpu")

    if checkpoint_filepath.startswith("s3://"):
        sd = read_from_s3(checkpoint_filepath, read_f=load_f)
    else:
        sd = load_f(checkpoint_filepath)

    sd["metadata"]["kwargs"] = {
        k: v
        for k, v in sd["metadata"]["kwargs"].items()
        if k in DAC.__init__.__code__.co_varnames
    }
    codec_model = DAC(**sd["metadata"]["kwargs"])
    codec_model.load_state_dict(sd["state_dict"])
    codec_model.eval()
    codec_model.to(device)

    base_metas = (
        "/home/christian/code/christian/metadata/genius_hq_metas_filtered.jsonl"
    )
    metas = read_jsonl(base_metas)
    print(f"loaded {len(metas):,} metas")

    # also load the lyric alignments
    genius_alignments_filepath = (
        "/home/tony/Work/tony/hoot/tmp/genius_hq_alignments_t30_v1.jsonl"
    )
    # if os.path.exists(genius_alignments_filepath):
    #    print("loading genius alignments")
    #    genius_alignments = read_jsonl(genius_alignments_filepath, progress=False)
    #    genius_alignments_map = {a[0]: a[1] for a in genius_alignments}
    # else:
    #    genius_alignments_map = {}
    # merge alignments into single map
    # alignments_map = {**genius_alignments_map}
    # print(f"loaded {len(alignments_map):,} alignments")

    # create metas map
    metas_map = {meta["id"]: meta for meta in metas}
    print(f"loaded {len(metas_map):,} metas")

    # Setup
    os.makedirs(out_dir, exist_ok=True)
    random.seed(10)
    random.shuffle(metas)
    subset_metas = metas[:num_samples]

    # subset_metas_with_alignments = []
    # add alignments to metas
    # for meta in subset_metas:
    #    if meta["id"] not in alignments_map:
    #         continue
    #    meta["alignments"] = alignments_map[meta["id"]]
    #    subset_metas_with_alignments.append(meta)

    # Move model to GPU
    codec_model.to("cuda:0")

    # Initialize processor
    process_chunks_parallel(
        metas=subset_metas,
        codec_model=codec_model,
        out_dir=out_dir,
        num_threads=num_threads,
        chunk_size=num_threads * 2,
        batch_size=batch_size,
    )


if __name__ == "__main__":
    out_dir = "/home/christian/data/dpo/genius_hq_corrupt_dpo_25hz_30_v2_npz"
    num_samples = 10000
    num_threads = 64
    batch_size = 8
    main(out_dir, num_threads, batch_size, num_samples)