# -*- coding: utf-8 -*-
"""
Benchmark: how many 30s chunks/second can we load & encode
"""

from __future__ import annotations
import os
import sys
import time
import soxr
import torch
import random
import tempfile
import numpy as np

from typing import Any, Dict, Iterable, List, Optional, Tuple
from contextlib import contextmanager
from suno_utils.audio import Audio
from suno_utils.utils.opusfile import OpusFile

from torch.utils.data import IterableDataset, DataLoader
from tokenizers import Tokenizer
from tqdm import tqdm

# If these are in a local repo, keep your path append — but do it once, at top.
sys.path.append("/home/christian/code/neon/sunoDiff")

# Replace with your actual audio module imports
# from your_audio_lib import Audio, OpusFile

from helpers import (
    dist_barrier,
    download_s3_file,
    get_filename,
    read_jsonl,
    write_jsonl,
    print_with_time_master,
    respell_random_words_in_text,
)

# test some of the local data
CODEC_FILEPATH = "s3://suno-data/minz/models/dac_vae_tuned_25hz.pth"
SEMANTIC_FILEPATH = "s3://suno-data/georg/models/semantic/mert_25.pt"
SEMANTIC_CLUSTERS_FILEPATH = "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy"

from suno_utils.tasks.dac_vae_fixed_25hz import (
    preload_models as preload_codec_models,
    decode as codec_decode,
    encode as codec_encode,
)

from suno_utils.tasks.mert_25 import (
    preload_models as preload_semantic_models,
    encode as semantic_encode,
)


# -----------------------------
# Text prep
# -----------------------------
def augment_text_training(tags: List[str], text: str) -> str:
    """Combine tags and text for training; tags may be empty."""
    tags = tags or []
    text = text or ""
    return (",".join(tags) + " " + text).strip()


def prepare_text_inference(tags: List[str], text: str) -> str:
    """Combine tags and text for inference; tags may be empty."""
    tags = tags or []
    text = text or ""
    return (",".join(tags) + " " + text).strip()


# -----------------------------
# Tokenizer
# -----------------------------
@contextmanager
def _download_from_s3_if_needed(maybe_s3_filepath: str):
    """If path is S3, download to a temp file and yield local path."""
    temp_dir = None
    tmp_filepath = maybe_s3_filepath
    if maybe_s3_filepath.startswith("s3://"):
        temp_dir = tempfile.TemporaryDirectory()
        filename = get_filename(maybe_s3_filepath, keep_ext=True)
        tmp_filepath = os.path.join(temp_dir.name, filename)
        download_s3_file(maybe_s3_filepath, tmp_filepath)
    try:
        yield tmp_filepath
    finally:
        if temp_dir is not None:
            temp_dir.cleanup()


def load_tokenizer(
    tokenizer_filepath: str = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json",
    pad_token: str = "[PAD]",
    extra_specials: Optional[List[str]] = None,
) -> Tokenizer:
    """Load and minimally configure a Tokenizer."""
    extra_specials = extra_specials or ["\n"]
    with _download_from_s3_if_needed(tokenizer_filepath) as tmp_fp:
        tok = Tokenizer.from_file(tmp_fp)
    tok.add_special_tokens(extra_specials)

    # Ensure PAD exists; fall back to adding if missing
    pad_id = tok.token_to_id(pad_token)
    if pad_id is None:
        tok.add_special_tokens([pad_token])
        pad_id = tok.token_to_id(pad_token)

    # Attach convenient attributes
    tok.pad_token = pad_token
    tok.pad_idx = pad_id
    return tok


# -----------------------------
# Audio slicing + meta handling
# -----------------------------
SAMPLE_RATE = 48_000


def _resample_to_mert(arr):
    assert arr.ndim == 2
    assert arr.shape[0] == 2
    out_arr = soxr.resample(arr.T, 48_000, 24_000).mean(axis=1).astype(np.float32)
    audio = Audio.from_array_float(out_arr, sample_rate=24000, max_allowed_val=12)
    return audio


def _choose_window(
    duration_s: float,
    chunk_s: float,
    ctx_s: float,
    use_ctx: bool,
) -> Tuple[float, float, float, float, bool]:
    """
    Choose a random window from the audio file.

    Returns:
        full_start, full_end, target_start, target_end, use_ctx
    """
    total_needed = chunk_s + ctx_s
    has_enough_for_context = duration_s >= total_needed

    # Only use context if requested AND we have enough audio
    use_ctx = use_ctx and has_enough_for_context

    if use_ctx:
        # Use both target and context: read 60s, context first 30s, target last 30s
        max_start = max(0.0, duration_s - total_needed)
        full_start = random.uniform(0.0, max_start)
        full_end = full_start + total_needed
        target_start = full_start + ctx_s
        target_end = full_start + total_needed
    else:
        # Use only target: read 30s
        max_start = max(0.0, duration_s - chunk_s)
        full_start = random.uniform(0.0, max_start)
        full_end = full_start + chunk_s
        target_start = full_start
        target_end = full_start + chunk_s

    return full_start, full_end, target_start, target_end, use_ctx


def _extract_audio_segments(
    path: str,
    duration_s: float,
    audio_chunk_s: float,
    audio_ctx_s: float,
) -> Tuple["Audio", Optional["Audio"], float, float]:
    """Extract target and context audio segments from the file.

    Returns:
        audio_target, audio_ctx, target_start, target_end
    """

    def read_opus_window(start_s: float, total_len_s: float) -> "Audio":
        buf_size = int(SAMPLE_RATE * total_len_s)
        float_arr = OpusFile(path=path).read(
            buf_size=buf_size,
            float_samples=True,
            from_position=int(SAMPLE_RATE * max(0.0, start_s)),
        )
        return Audio.from_array_float(
            float_arr, sample_rate=SAMPLE_RATE, max_allowed_val=12
        )

    # Handle short files
    if duration_s < audio_chunk_s:
        window = read_opus_window(0.0, duration_s)
        audio_target = window.pad_to_length(audio_chunk_s)
        return audio_target, None, 0.0, duration_s

    # Choose window and extract segments
    use_ctx = random.random() < 0.5
    full_start, full_end, target_start, target_end, use_ctx = _choose_window(
        duration_s=duration_s,
        chunk_s=audio_chunk_s,
        ctx_s=audio_ctx_s,
        use_ctx=use_ctx,
    )

    window = read_opus_window(full_start, full_end - full_start)

    if use_ctx:
        # When using context: first 30s is context, last 30s is target
        audio_ctx = window.get_segment(from_s=0.0, to_s=audio_ctx_s)
        audio_target = window.get_segment(
            from_s=audio_ctx_s, to_s=audio_ctx_s + audio_chunk_s
        ).pad_to_length(audio_chunk_s)
    else:
        # When not using context: only target audio (first 30s of window)
        audio_ctx = None
        audio_target = window.get_segment(from_s=0.0, to_s=audio_chunk_s).pad_to_length(
            audio_chunk_s
        )

    return audio_target, audio_ctx, target_start, target_end


def _extract_lyrics_for_window(
    text_aligned: List[Tuple[float, float, str]],
    window_start: float,
    window_end: float,
) -> str:
    """
    Extract lyrics that fall within the given time window.

    Args:
        text_aligned: list of (start_time, end_time, text) tuples
        window_start: start time of the window
        window_end: end time of the window

    Returns:
        Concatenated text from lyrics fully contained within the window
    """
    if not text_aligned:
        return ""
    parts = [
        text
        for (s, e, text) in text_aligned
        if (s >= window_start) and (e <= window_end)
    ]
    return "".join(parts).strip()


def _process_text(
    meta: Dict[str, Any],
    target_start: float,
    target_end: float,
    tokenizer: Tokenizer,
    cond_text_len: int,
    is_training: bool,
    text_drop_prob: float,
) -> torch.Tensor:
    """Process text data into tokenized tensor."""
    # Extract text for the target window
    tags = meta.get("tags") or []
    if "text_aligned" in meta and meta["text_aligned"]:
        text = _extract_lyrics_for_window(
            meta["text_aligned"], target_start, target_end
        )
    else:
        text = ""

    # Combine tags and text
    if is_training:
        full_text = (
            ""
            if random.random() <= text_drop_prob
            else augment_text_training(tags, text)
        )
    else:
        full_text = prepare_text_inference(tags, text)

    # Tokenize and pad/truncate
    ids = tokenizer.encode(full_text).ids[:cond_text_len]
    if tokenizer.pad_idx is None:
        raise ValueError("Tokenizer must define a pad_idx.")
    if len(ids) < cond_text_len:
        ids += [tokenizer.pad_idx] * (cond_text_len - len(ids))

    return torch.tensor(ids, dtype=torch.long)


def load_meta(
    meta: Dict[str, Any],
    tokenizer: Tokenizer,
    cond_text_len: int,
    audio_chunk_s: float,
    audio_ctx_s: float,
    is_training: bool,
    text_drop_prob: float = 0.1,
):
    """Load audio and text data from a meta entry.

    Assumes meta['local_filepath'] is an OPUS file.
    Guarantees:
      - audio_target is exactly `audio_chunk_s` seconds (pads with silence if needed)
      - If the file is shorter than `audio_chunk_s`, audio_ctx is None
    """
    duration_s = float(meta["duration_s"])
    path = meta["local_filepath"]

    # Extract audio segments
    audio_target, audio_ctx, target_start, target_end = _extract_audio_segments(
        path=path,
        duration_s=duration_s,
        audio_chunk_s=audio_chunk_s,
        audio_ctx_s=audio_ctx_s,
    )

    # Process text data
    text_codes = _process_text(
        meta=meta,
        target_start=target_start,
        target_end=target_end,
        tokenizer=tokenizer,
        cond_text_len=cond_text_len,
        is_training=is_training,
        text_drop_prob=text_drop_prob,
    )

    # Create 24kHz mono version for semantic encoding
    audio_target_24k = _resample_to_mert(audio_target.array_float)

    return audio_target, audio_ctx, text_codes, audio_target_24k


# -----------------------------
# Dataset & collate
# -----------------------------
class DynamicDataset(IterableDataset):
    def __init__(
        self,
        metas: Iterable[Dict[str, Any]],
        *,
        audio_chunk_s: float = 30.0,
        audio_ctx_s: float = 30.0,
        cond_text_len: int = 1536,
        text_drop_prob: float = 0.1,
        is_training: bool = False,
        tokenizer: Optional[Tokenizer] = None,
        tokenizer_path: str = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json",
    ):
        self.metas = list(metas)
        self.audio_chunk_s = float(audio_chunk_s)
        self.audio_ctx_s = float(audio_ctx_s)
        self.cond_text_len = int(cond_text_len)
        self.is_training = bool(is_training)
        self.text_drop_prob = float(text_drop_prob)
        self.tokenizer = tokenizer or load_tokenizer(tokenizer_path)

    def __iter__(self):
        # IterableDataset must not rely on shuffle=True in DataLoader.
        for meta in self.metas:
            try:
                audio_target, audio_ctx, text_codes, audio_target_24k = load_meta(
                    meta=meta,
                    tokenizer=self.tokenizer,
                    cond_text_len=self.cond_text_len,
                    audio_chunk_s=self.audio_chunk_s,
                    audio_ctx_s=self.audio_ctx_s,
                    is_training=self.is_training,
                    text_drop_prob=self.text_drop_prob,
                )
                yield audio_target, audio_ctx, text_codes, audio_target_24k
            except Exception as e:
                print(f"[WARN] Error loading meta {meta.get('id', '<unknown>')}: {e}")
                continue


def collate_fn(batch: List[Tuple[Any, Optional[Any], torch.Tensor, Any]]):
    audio_target_list = [item[0] for item in batch]
    audio_ctx_list = [item[1] for item in batch]
    text_codes_list = [item[2] for item in batch]
    audio_target_24k_list = [item[3] for item in batch]
    # If your encoder expects tensors, you can add padding/stacking here instead.
    return audio_target_list, audio_ctx_list, text_codes_list, audio_target_24k_list


def _as_tensor_batch(x):
    """Convert codec output to a torch tensor (batch-first)."""
    if isinstance(x, torch.Tensor):
        return x.detach().cpu()
    try:
        import numpy as np

        if isinstance(x, np.ndarray):
            return torch.from_numpy(x)
    except Exception:
        pass
    if isinstance(x, (list, tuple)):
        elems = [_as_tensor_batch(e) for e in x]
        return torch.stack(elems, dim=0)
    return torch.as_tensor(x)


def benchmark(
    metas: Iterable[Dict[str, Any]],
    *,
    batch_size: int = 2,
    num_workers: int = 1,
    num_eval_batches: int = 100,
    codec_encode_fn=None,
    audio_chunk_s: float = 30.02,
    audio_ctx_s: float = 30.02,
    cond_text_len: int = 1536,
    is_training: bool = False,
    text_drop_prob: float = 0.1,
    n_vae_tokens: int = 750,
    vae_dim: int = 128,
):
    """
    codec_encode_fn: function(audio_list, normalize_volume=False) -> (B, T, D)
      Expected output shape: (B, n_vae_tokens, vae_dim)

    Builds:
      - vae_target: (B, n_vae_tokens, vae_dim)
      - audio_ctx_vae: (B, n_vae_tokens, vae_dim) (zeros where ctx=None)
      - audio_ctx_mask: (B, 1, 1)  (1.0 where ctx exists, 0.0 otherwise)
    """
    assert codec_encode_fn is not None, "Please pass codec_encode_fn=your_encoder"

    dataset = DynamicDataset(
        metas=metas,
        audio_chunk_s=audio_chunk_s,
        audio_ctx_s=audio_ctx_s,
        cond_text_len=cond_text_len,
        is_training=is_training,
        text_drop_prob=text_drop_prob,
    )
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,  # IterableDataset must not be shuffled here
        collate_fn=collate_fn,
        num_workers=num_workers,
        pin_memory=False,
    )

    seen = 0
    t0 = time.time()

    for i, (
        audio_target_list,
        audio_ctx_list,
        text_codes_list,
        audio_target_24k_list,
    ) in enumerate(tqdm(loader, desc="Processing batches")):
        # ----------------------------
        # Encode audio targets (always present)
        # ----------------------------
        vae_target = _as_tensor_batch(
            codec_encode_fn(audio_target_list, normalize_volume=False)
        )

        # Encode semantic targets (24kHz mono audio)
        semantic_codes = _as_tensor_batch(semantic_encode(audio_target_24k_list))

        if vae_target.ndim != 3:
            raise ValueError(
                f"codec_encode(target) must return (B, T, D); got {vae_target.shape}"
            )

        B, T, D = vae_target.shape
        if T != n_vae_tokens or D != vae_dim:
            raise ValueError(
                f"codec_encode(target) shape mismatch: expected (B, {n_vae_tokens}, {vae_dim}), "
                f"got (B, {T}, {D})"
            )

        device, dtype = vae_target.device, vae_target.dtype

        # ----------------------------
        # Encode audio context (may be None)
        # ----------------------------
        audio_ctx_vae = torch.zeros(
            (B, n_vae_tokens, vae_dim), dtype=dtype, device=device
        )
        audio_ctx_mask = torch.zeros((B, n_vae_tokens, 1), dtype=dtype, device=device)

        ctx_indices = [idx for idx, a in enumerate(audio_ctx_list) if a is not None]
        if ctx_indices:
            ctx_batch = [audio_ctx_list[idx] for idx in ctx_indices]
            encoded_ctx = _as_tensor_batch(
                codec_encode_fn(ctx_batch, normalize_volume=False)
            )

            if encoded_ctx.ndim != 3:
                raise ValueError(
                    f"codec_encode(ctx) must return (b, T, D); got {encoded_ctx.shape}"
                )

            b_ctx, T_ctx, D_ctx = encoded_ctx.shape
            if T_ctx != n_vae_tokens or D_ctx != vae_dim:
                raise ValueError(
                    f"codec_encode(ctx) shape mismatch: expected (b, {n_vae_tokens}, {vae_dim}), "
                    f"got (b, {T_ctx}, {D_ctx})"
                )

            # Scatter ctx results into full batch tensors
            for k, global_idx in enumerate(ctx_indices):
                audio_ctx_vae[global_idx] = encoded_ctx[k]
                audio_ctx_mask[global_idx] = 1.0

        # ----------------------------
        # Here you can run downstream model / profiling / etc.
        # ----------------------------
        # Example: dummy use so these aren't optimized out
        _ = (
            vae_target,
            audio_ctx_vae,
            audio_ctx_mask,
            text_codes_list,
            semantic_codes,
        )

        # special case to decode audio
        print(
            vae_target.shape,
            audio_ctx_vae.shape,
            audio_ctx_mask.shape,
            semantic_codes.shape,
        )
        # vae_audio_decoded = codec_decode(vae_target[0])
        # audio_ctx_vae_decoded = codec_decode(audio_ctx_vae[0])
        # combine the audio
        combined_vae = torch.cat([vae_target, audio_ctx_vae], dim=1)
        for n in range(combined_vae.shape[0]):
            combined_vae_decoded = codec_decode(combined_vae[n])
            combined_vae_decoded.write_hq_mp3(f"combined_vae_decoded_{n}.mp3")

        seen += 1
        if seen >= num_eval_batches:
            break

    # ----------------------------
    # Timing summary
    # ----------------------------
    t1 = time.time()
    elapsed = t1 - t0
    time_per_batch = elapsed / max(1, seen)
    chunks_per_batch = batch_size
    chunks_per_second = (
        chunks_per_batch / time_per_batch if time_per_batch > 0 else float("inf")
    )

    print(f"Time taken: {elapsed:.3f}s for {seen} batches")
    print(f"Time per batch: {time_per_batch:.4f}s")
    print(f"~Chunks/sec (30s targets): {chunks_per_second:.2f}")


def main():
    print("Loading models...")
    _ = preload_codec_models(CODEC_FILEPATH)
    _ = preload_semantic_models(SEMANTIC_FILEPATH, SEMANTIC_CLUSTERS_FILEPATH)

    print("Loading metas...")
    filepath = "/app2/suno/data/auk_v0/metas_v8_tr_mini.jsonl"
    metas = read_jsonl(filepath)
    print(len(metas))

    print("Benchmarking...")
    benchmark(metas, codec_encode_fn=codec_encode)


if __name__ == "__main__":
    main()