import numpy as np
import torch
import argparse
import json
from pathlib import Path
from typing import List, Optional, Union, Dict, Any
from suno_utils.tasks.mert_25 import encode, preload_models
from suno_utils.audio import Audio
from tqdm import tqdm
import os


def encode_audio_to_npz(
    audio_path: Union[str, Path],
    save_path: Union[str, Path],
    device: Optional[str] = None,
) -> None:
    """
    Encode an audio file using MERT and save the codes to an NPZ file.

    Args:
        audio_path: Path to the input audio file
        save_path: Path to save the output NPZ file
        device: CUDA device string (e.g., 'cuda:0'). If None, will use default device.
    """
    import tempfile
    import shutil

    # Convert Path to str for Audio.from_file
    audio_path = str(audio_path)
    save_path = str(save_path)

    # Load audio file
    audio = Audio.from_file(audio_path)

    # Encode using MERT (it will use the device specified during preload_models)
    # we only keep the first one!
    codes = encode(audio, do_clustering=True, device=device)[:, 0]

    # Ensure codes is a numpy array
    if not isinstance(codes, np.ndarray):
        codes = np.array(codes)

        # Save to a temporary file first, then move atomically
    save_dir = os.path.dirname(save_path)

    # Generate a unique temporary filename without creating the file
    tmp_path = os.path.join(save_dir, f".tmp_{os.getpid()}_{id(codes)}.npz")

    try:
        # Save to temporary file
        np.savez_compressed(tmp_path, codes=codes)

        # Atomically move to final location
        # os.rename is atomic on the same filesystem
        os.rename(tmp_path, save_path)

    except Exception as e:
        # Clean up temporary file if something went wrong
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        raise e


def process_jsonl_chunk(
    jsonl_path: Union[str, Path],
    output_dir: Union[str, Path],
    chunk_id: int,
    total_chunks: int = 8,
    local_gpu_id: Optional[int] = None,
) -> None:
    """
    Process a chunk of the JSONL file based on chunk_id.

    Args:
        jsonl_path: Path to the JSONL file
        output_dir: Directory to save NPZ files
        chunk_id: Global chunk ID across all nodes
        total_chunks: Total number of chunks to split the data
        local_gpu_id: Local GPU ID on this node (0-7). If None, uses chunk_id % 8
    """
    # Determine local GPU ID
    if local_gpu_id is None:
        local_gpu_id = chunk_id % 8

    # When CUDA_VISIBLE_DEVICES is set, PyTorch sees only the visible GPU(s)
    # and indexes them starting from 0
    if "CUDA_VISIBLE_DEVICES" in os.environ:
        # If CUDA_VISIBLE_DEVICES is set, use device 0 (the only visible device)
        device = "cuda:0"
        device_id = 0
    else:
        # If not set, use the actual GPU ID
        device = f"cuda:{local_gpu_id}"
        device_id = local_gpu_id

    # Check if CUDA is available
    if not torch.cuda.is_available():
        print(f"Warning: CUDA not available, using CPU instead")
        device = "cpu"
    else:
        try:
            torch.cuda.set_device(device_id)
        except RuntimeError as e:
            print(f"Warning: Failed to set CUDA device {device_id}: {e}")
            device = "cpu"

    # Preload models onto the specified device
    print(f"Loading models on {device} for chunk {chunk_id}...")
    preload_models(
        checkpoint_filepath="/app/suno/data/dpo/models/mert_25.pt",
        centroids_filepath="/app/suno/data/dpo/models/mert_25_2x4k.npy",
        device=device,
    )

    # Create output directory if it doesn't exist
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load JSONL file
    jsonl_path = Path(jsonl_path)
    print(f"Loading JSONL file: {jsonl_path}")

    entries: List[Dict[str, Any]] = []
    try:
        with open(jsonl_path, "r") as f:
            for line in f:
                line = line.strip()
                if line:  # Skip empty lines
                    try:
                        entry = json.loads(line)
                        entries.append(entry)
                    except json.JSONDecodeError as e:
                        print(f"Warning: Failed to parse JSON line: {e}")
                        continue
    except FileNotFoundError:
        print(f"Error: JSONL file not found: {jsonl_path}")
        return
    except Exception as e:
        print(f"Error reading JSONL file: {e}")
        return

    # Calculate chunk boundaries
    total_entries = len(entries)
    if total_entries == 0:
        print("Warning: No entries found in JSONL file")
        return

    chunk_size = total_entries // total_chunks
    start_idx = chunk_id * chunk_size

    # Handle the last chunk to include any remaining entries
    if chunk_id == total_chunks - 1:
        end_idx = total_entries
    else:
        end_idx = (chunk_id + 1) * chunk_size

    chunk_entries = entries[start_idx:end_idx]
    print(
        f"Processing chunk {chunk_id} (GPU {local_gpu_id}): entries {start_idx} to {end_idx-1} (total: {len(chunk_entries)})"
    )

    # Process each entry in the chunk
    successful = 0
    failed = 0

    for entry in tqdm(chunk_entries, desc=f"Chunk {chunk_id}"):
        try:
            # Get the local file path and ID from the entry
            audio_path = entry.get("local_filepath")
            entry_id = entry.get("id")

            if not audio_path:
                print(f"Warning: Missing local_filepath in entry: {entry}")
                failed += 1
                continue

            if not entry_id:
                print(f"Warning: Missing id in entry: {entry}")
                failed += 1
                continue

            # Ensure paths are strings
            audio_path = str(audio_path)
            entry_id = str(entry_id)

            # Check if audio file exists
            if not Path(audio_path).exists():
                print(f"Warning: Audio file not found: {audio_path}")
                failed += 1
                continue

            # Generate output filename using the ID
            output_filename = f"{entry_id}.npz"
            output_path = output_dir / output_filename

            # Skip if already processed
            # if output_path.exists():
            #     successful += 1
            #     continue

            # Process the audio file
            encode_audio_to_npz(audio_path, output_path, device=device)
            successful += 1

        except Exception as e:
            entry_id = entry.get("id", "unknown")
            print(f"Error processing entry {entry_id}: {e}")
            import traceback

            traceback.print_exc()
            failed += 1

    print(f"\nChunk {chunk_id} completed: {successful} successful, {failed} failed")


def main():
    parser = argparse.ArgumentParser(
        description="Encode audio files from JSONL using MERT and save to NPZ"
    )
    parser.add_argument(
        "--jsonl_path",
        type=str,
        default="/app2/suno/data/dpo/sft/sft_metas_tr_v11.jsonl",
        help="Path to the JSONL file containing audio metadata",
    )
    parser.add_argument(
        "--output_dir", type=str, required=True, help="Directory to save NPZ files"
    )
    parser.add_argument(
        "--chunk_id",
        type=int,
        required=True,
        help="Global chunk ID across all nodes",
    )
    parser.add_argument(
        "--total_chunks",
        type=int,
        default=128,
        help="Total number of chunks to split the data",
    )
    parser.add_argument(
        "--local_gpu_id",
        type=int,
        default=None,
        help="Local GPU ID on this node (0-7)",
    )

    args = parser.parse_args()

    # Process the chunk
    process_jsonl_chunk(
        args.jsonl_path,
        args.output_dir,
        args.chunk_id,
        args.total_chunks,
        args.local_gpu_id,
    )


if __name__ == "__main__":
    main()


# Simple test function
def test_encode_audio():
    """
    Simple test case for encoding audio.
    This creates a synthetic audio file for testing.
    """
    import tempfile
    import os

    # Set device
    device = "cuda:0" if torch.cuda.is_available() else "cpu"

    # Preload models
    print("Loading models for test...")
    preload_models(
        checkpoint_filepath="/app/suno/data/dpo/models/mert_25.pt",
        centroids_filepath="/app/suno/data/dpo/models/mert_25_2x4k.npy",
        device=device,
    )

    # Create a test audio file (silence for testing)
    test_audio = Audio.from_silence(duration_s=5.0, sample_rate=24000)

    with tempfile.TemporaryDirectory() as tmpdir:
        # Save test audio
        audio_path = os.path.join(tmpdir, "test_audio.wav")
        test_audio.write_wav(audio_path)

        # Encode and save
        npz_path = os.path.join(tmpdir, "test_encoded.npz")
        encode_audio_to_npz(audio_path, npz_path, device=device)

        # Verify the file was created
        assert os.path.exists(npz_path), "NPZ file was not created"

        # Load and verify the encoded data
        loaded_data = np.load(npz_path, allow_pickle=True)
        codes = loaded_data["codes"]

        # Basic checks
        assert isinstance(codes, np.ndarray), "Codes should be a numpy array"
        assert codes.ndim == 1, f"Codes should be 1D, but got shape {codes.shape}"
        assert codes.shape[0] > 0, "Codes should have non-zero length"

        print(f"Test passed! Encoded shape: {codes.shape}")
        print(f"Codes dtype: {codes.dtype}")