import os
import time
import torch
import torchaudio
import numpy as np
import pyloudnorm as pyln
from tqdm import tqdm
from typing import List

from suno_utils.utils.s3 import read_from_s3
from suno_utils.utils.text import read_jsonl, write_jsonl


# 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 get_hann_window(frame_size: int, device=None) -> torch.Tensor:
    """Create a Hann window."""
    return torch.hann_window(frame_size, device=device)


def frame_audio(x: torch.Tensor, frame_size: int, overlap: float = 0.0) -> torch.Tensor:
    """
    Split an audio tensor into frames. Apply Hann window only when overlap > 0.

    Args:
        x: Input tensor of shape (batch_size, channels, sequence_length)
        frame_size: Number of samples per frame
        overlap: Overlap between consecutive frames

    Returns:
        Framed tensor of shape (batch_size, channels, num_frames, frame_size)
    """
    if not 0.0 <= overlap < 1.0:
        raise ValueError("Overlap must be in [0.0, 1.0)")

    batch_size, channels, seq_len = x.shape
    hop_size = int(frame_size * (1 - overlap))
    num_frames = (seq_len - frame_size) // hop_size + 1

    # Get window if needed
    window = None
    if overlap > 0:
        window = get_hann_window(frame_size, device=x.device)

    # Create output tensor
    output = torch.zeros(
        batch_size, channels, num_frames, frame_size, dtype=x.dtype, device=x.device
    )

    # Fill output tensor with frames
    for i in range(num_frames):
        start_idx = i * hop_size
        frame = x[:, :, start_idx : start_idx + frame_size]
        if window is not None:
            frame = frame * window
        output[:, :, i] = frame

    return output


def reconstruct_audio(
    frames: torch.Tensor, original_length: int, overlap: float = 0.0
) -> torch.Tensor:
    """
    Reconstruct audio signal from frames. Uses overlap-add only when overlap > 0.

    Args:
        frames: Input tensor of shape (batch_size, channels, num_frames, frame_size)
        original_length: Length of the original sequence
        overlap: Overlap used in framing
    """
    batch_size, channels, num_frames, frame_size = frames.shape
    hop_size = int(frame_size * (1 - overlap))

    # For no overlap, we can just reshape
    if overlap == 0:
        # Check if the frames can be directly reshaped
        expected_length = num_frames * frame_size
        if expected_length == original_length:
            return frames.reshape(batch_size, channels, -1)
        else:
            # If not exact match, still do frame-by-frame to handle partial frames
            output = torch.zeros(
                batch_size,
                channels,
                original_length,
                dtype=frames.dtype,
                device=frames.device,
            )
            for i in range(num_frames):
                start_idx = i * frame_size
                end_idx = min(start_idx + frame_size, original_length)
                output[:, :, start_idx:end_idx] = frames[
                    :, :, i, : (end_idx - start_idx)
                ]
            return output

    # For overlap > 0, use overlap-add (Hann window sum to 1)
    output = torch.zeros(
        batch_size, channels, original_length, dtype=frames.dtype, device=frames.device
    )

    for i in range(num_frames):
        start_idx = i * hop_size
        end_idx = start_idx + frame_size
        output[:, :, start_idx:end_idx] += frames[:, :, i]

    return output[:, :, :original_length]


class AudioDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        metas: List[dict],
        npz_out_dir: str = None,
        frame_size: int = 1024,
        overlap: float = 0.0,
    ):
        self.metas = metas
        self.frame_size = frame_size
        self.overlap = overlap
        self.npz_out_dir = npz_out_dir

    def __len__(self):
        return len(self.metas)

    def __getitem__(self, idx):
        meta = self.metas[idx]

        # check if npz exists
        # npz_path = os.path.join(self.npz_out_dir, f"{meta['id']}.npz")
        # if os.path.exists(npz_path):
        #    print(f"skipping {meta['id']} because npz exists")
        #    return None

        # read from s3
        try:
            start_time = time.time()
            audio, sr = read_from_s3(meta["audio_filepath"], read_f=torchaudio.load)

            # crop to DURATION_S
            # audio = audio[:, : int(DURATION_S * sr)]

            # print(f"loaded {meta['id']} in {time.time() - start_time:.2f} seconds")
        except Exception as e:
            print(f"error loading {meta['id']}: {e}")
            return None

        try:
            # resample to 48k
            audio_48k = torchaudio.transforms.Resample(sr, 48000)(audio)

            # loudness normalize
            meter = pyln.Meter(48000)
            input_loudness = meter.integrated_loudness(audio_48k.permute(1, 0).numpy())
            gain_db = TARGET_LUFS - input_loudness
            audio_48k = audio_48k * (10 ** (gain_db / 20.0))

            # resample to 24k (mono)
            audio_24k = torchaudio.transforms.Resample(48000, 24000)(
                audio_48k.mean(dim=0, keepdim=True)
            )
        except Exception as e:
            print(f"error processing {meta['id']}: {e}")
            return None

        return audio_48k, audio_24k, idx


def collate_frames(batch: List[torch.Tensor]) -> List[torch.Tensor]:
    """
    Simple collate function that returns the batch as a list of frame tensors.

    Args:
        batch: List of frame tensors of shape (num_frames, frame_size)

    Returns:
        The same list of frame tensors without any padding or modification
    """
    return batch


if __name__ == "__main__":

    FRAME_SIZE = 1920
    OVERLAP = 0.0
    OUT_DIR = "/app/suno/christian/data/genius_hq_filtered_raw_30s_1920_npz"
    NUM_WORKERS = 64
    BATCH_SIZE = 4
    TARGET_LUFS = -16.0
    DURATION_S = 10.01

    os.makedirs(OUT_DIR, exist_ok=True)

    # load base metas to encode
    base_metas = read_jsonl(
        "/home/christian/code/christian/metadata/genius_hq_metas_filtered.jsonl"
    )
    print(f"Found {len(base_metas)} metas")

    # create dataset
    dataset = AudioDataset(base_metas, npz_out_dir=OUT_DIR)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=BATCH_SIZE,
        shuffle=False,
        collate_fn=collate_frames,
        num_workers=NUM_WORKERS,
    )

    for batch in tqdm(dataloader):
        # save to disk
        for i, result in enumerate(batch):

            if result is None:
                print("skipping")
                continue

            audio_48k, audio_24k, idx = result

            # encode into frames
            frames = frame_audio(audio_48k.unsqueeze(0), FRAME_SIZE, OVERLAP)
            frames = frames.half()  # convert to half

            # encode with semantic
            semantic_codes = encode_semantic([audio_24k.cuda()])[0]
            semantic_codes = semantic_codes.astype(np.uint16)[:, 0]

            meta_id = base_metas[idx]["id"]

            # print(frames.shape)
            # print(semantic_codes.shape)
            # Save results for each item in batch
            np.savez(
                os.path.join(OUT_DIR, f"{meta_id}.npz"),
                frames=frames,
                semantic_codes=semantic_codes,
            )