import numpy as np
import soundfile as sf
from pathlib import Path
from suno_utils.utils.text import read_jsonl
from typing import List, Dict, Optional, Any, Tuple
from joblib import Parallel, delayed
import time
import json
from datetime import datetime


def _fast_trim_mono(
    x: np.ndarray,  # shape: (samples,), float32/64 in [-1, 1]
    sr: int,  # sample rate (Hz)
    thresh_db_rel: float = -35,  # keep where RMS > max_RMS + thresh (dB)
    win_ms: float = 20.0,  # moving RMS window size (ms)
    pad_ms: float = 20.0,  # pad around kept regions (ms)
    min_keep_ms: float = 40.0,  # drop kept bits shorter than this (ms)
) -> Tuple[np.ndarray, List[Tuple[int, int]]]:
    """
    Ultra-fast silence trimmer for mono audio. No convolutions, all O(n).
    Returns (trimmed_audio, kept_spans) with kept_spans in original sample indices.
    """
    assert x.ndim == 1, "Expected mono waveform of shape (samples,)"
    n = x.size
    if n == 0:
        return x[:0], []

    # --- Moving RMS via cumulative sums (box filter), O(n) ---
    # Compute moving average of power over a window, then sqrt.
    win = max(1, int(round(sr * win_ms / 1000.0)))
    if win > n:
        win = n

    # power and cumulative sum (use float64 for numeric safety)
    sq = x.astype(np.float64) ** 2
    csum = np.empty(n + 1, dtype=np.float64)
    csum[0] = 0.0
    np.cumsum(sq, out=csum[1:])  # csum[k] = sum_{i<k} sq[i]

    # moving average (valid positions)
    # ma_valid[t] = mean of sq[t : t+win]
    ma_valid = (csum[win:] - csum[:-win]) / win  # length n - win + 1

    # Center-align to original length by padding equally on both sides
    left = win // 2
    right = n - (ma_valid.size + left)
    rms = np.sqrt(np.pad(ma_valid, (left, right), mode="edge"))

    # --- Threshold relative to max ---
    eps = 1e-12
    rel_db = 20.0 * np.log10(np.maximum(rms, eps) / (np.max(rms) + eps))
    mask = rel_db > thresh_db_rel  # True = keep

    # --- Turn mask into spans, expand by pad, merge, drop short ---
    pad = max(0, int(round(sr * pad_ms / 1000.0)))
    min_keep = max(1, int(round(sr * min_keep_ms / 1000.0)))

    # Find rising/falling edges
    m = mask.astype(np.int8)
    edges = np.flatnonzero(np.diff(m, prepend=0, append=0))
    # edges come in pairs [start0, end0, start1, end1, ...]
    starts = edges[::2]
    ends = edges[1::2]

    if starts.size == 0:
        return x[:0], []

    # Expand by pad and clamp
    starts = np.maximum(0, starts - pad)
    ends = np.minimum(n, ends + pad)

    # Merge overlaps and drop short spans
    spans: List[Tuple[int, int]] = []
    s_prev = int(starts[0])
    e_prev = int(ends[0])
    for s, e in zip(starts[1:], ends[1:]):
        s = int(s)
        e = int(e)
        if s <= e_prev:  # overlap/adjacent -> merge
            e_prev = max(e_prev, e)
        else:
            if (e_prev - s_prev) >= min_keep:
                spans.append((s_prev, e_prev))
            s_prev, e_prev = s, e
    # last span
    if (e_prev - s_prev) >= min_keep:
        spans.append((s_prev, e_prev))

    if not spans:
        return x[:0], []

    # --- Concatenate kept spans (one pass) ---
    parts = [x[a:b] for (a, b) in spans]
    y = np.concatenate(parts, axis=0).astype(x.dtype)
    return y, spans


def _build_artist_vox_mappings(
    metas: List[Dict],
) -> Tuple[Dict[str, List[int]], Dict[str, Optional[Dict[str, Any]]]]:
    """
    Build mappings from artist IDs to metadata indices and vox stem paths with duration.

    Args:
        metas: List of metadata dictionaries

    Returns:
        Tuple of (artist_id_to_meta_idx, artist_id_to_vox_paths)
        - artist_id_to_meta_idx: Maps artist_id to list of meta indices
        - artist_id_to_vox_paths: Maps artist_id to dict with 'path' and 'duration_s' (or None)
    """
    # Build mapping of artist_ids to meta indices
    artist_id_to_meta_idx = {}
    for idx, meta in enumerate(metas):
        if "artist_ids" in meta and meta["artist_ids"] is not None:
            if len(meta["artist_ids"]) == 1:  # only do if there is a single artist
                for artist_id in meta["artist_ids"]:
                    if artist_id not in artist_id_to_meta_idx:
                        artist_id_to_meta_idx[artist_id] = []
                    artist_id_to_meta_idx[artist_id].append(idx)

    # Build mapping of artist_ids to vox stem paths with duration
    artist_id_to_vox_paths = {}
    for artist_id, meta_indices in artist_id_to_meta_idx.items():
        for meta_idx in meta_indices:
            meta = metas[meta_idx]
            if (
                meta.get("stems") is not None
                and meta.get("stems").get("Vocals") is not None
            ):
                # Extract duration_s from the parent metadata
                duration_s = meta.get("duration_s", None)
                if artist_id not in artist_id_to_vox_paths:
                    artist_id_to_vox_paths[artist_id] = []
                artist_id_to_vox_paths[artist_id].append(
                    {
                        "path": meta["stems"]["Vocals"],
                        "duration_s": duration_s,
                    }
                )

    return artist_id_to_vox_paths


def _preprocess_vox_meta(
    filepath: str,
    dry_run: bool = False,
    output_dir: Optional[str] = None,
    skip_existing: bool = True,
) -> Dict:
    """
    Load an Opus audio file, convert to mono, trim silence, and save as processed Opus file.

    Args:
        filepath: Path to the input Opus file
        dry_run: If True, process the file but don't write output
        output_dir: Optional custom directory to save output files (for testing)
        skip_existing: If True, skip processing if output file already exists

    Returns:
        Dict with processing results
    """
    try:
        # Validate input file exists
        if not Path(filepath).exists():
            raise FileNotFoundError(f"Input file does not exist: {filepath}")

        # Load the Opus audio file
        audio_data, sample_rate = sf.read(filepath, dtype="float32")

        # Validate audio data
        if len(audio_data) == 0:
            raise ValueError("Audio file is empty")

        # Convert to mono if stereo
        if audio_data.ndim > 1:
            audio_data = np.mean(audio_data, axis=1)

        # Trim silence using the fast trim function
        trimmed_audio, kept_spans = _fast_trim_mono(
            audio_data,
            sr=sample_rate,
            thresh_db_rel=-35.0,
            win_ms=20.0,
            pad_ms=20.0,
            min_keep_ms=40.0,
        )

        # Create output path with _trimmed suffix
        input_path = Path(filepath)
        if output_dir:
            # Use custom output directory for testing
            output_dir_path = Path(output_dir)
            output_path = (
                output_dir_path / f"{input_path.stem}_trimmed{input_path.suffix}"
            )
        else:
            # Use original directory
            output_path = (
                input_path.parent / f"{input_path.stem}_trimmed{input_path.suffix}"
            )

        # Check if output file already exists (resumability)
        if skip_existing and output_path.exists():
            # Get file info for existing file
            existing_size = output_path.stat().st_size
            return {
                "input_path": filepath,
                "output_path": str(output_path),
                "original_duration_s": 0,  # We don't know without processing
                "trimmed_duration_s": 0,  # We don't know without processing
                "kept_spans": [],
                "success": True,
                "dry_run": dry_run,
                "output_dir": output_dir,
                "skipped": True,
                "existing_size_bytes": existing_size,
            }

        # Save as Opus file (unless dry run)
        if not dry_run:
            # Ensure output directory exists
            output_path.parent.mkdir(parents=True, exist_ok=True)

            # Save as Opus file
            sf.write(
                str(output_path),
                trimmed_audio,
                sample_rate,
                format="OGG",  # Opus files use OGG container
                subtype="OPUS",
            )
        else:
            if output_dir:
                print(
                    f"Dry run: {filepath} -> {output_path} (custom dir: {output_dir})"
                )
            else:
                print(f"Dry run: {filepath} -> {output_path}")

        return {
            "input_path": filepath,
            "output_path": str(output_path),
            "original_duration_s": len(audio_data) / sample_rate,
            "trimmed_duration_s": len(trimmed_audio) / sample_rate,
            "kept_spans": kept_spans,
            "success": True,
            "dry_run": dry_run,
            "output_dir": output_dir,
            "skipped": False,
        }

    except Exception as e:
        return {"input_path": filepath, "error": str(e), "success": False}


def _process_single_file(
    vox_meta: Dict, dry_run: bool, output_dir: Optional[str], skip_existing: bool = True
) -> Dict:
    """Wrapper function for parallel processing of a single file."""
    return _preprocess_vox_meta(
        vox_meta["path"],
        dry_run=dry_run,
        output_dir=output_dir,
        skip_existing=skip_existing,
    )


def _save_metadata(
    results: List[Dict],
    output_dir: Optional[str],
    processing_time: float,
    total_files: int,
    processed_count: int,
    skipped_count: int,
    error_count: int,
    n_jobs: int,
) -> str:
    """Save processing metadata to JSON file."""

    # Create metadata summary
    metadata = {
        "processing_info": {
            "timestamp": datetime.now().isoformat(),
            "total_files": total_files,
            "processed_count": processed_count,
            "skipped_count": skipped_count,
            "error_count": error_count,
            "processing_time_seconds": processing_time,
            "n_jobs": n_jobs,
            "avg_time_per_file": (
                processing_time / processed_count if processed_count > 0 else 0
            ),
            "files_per_second": (
                processed_count / processing_time if processing_time > 0 else 0
            ),
        },
        "results": results,
    }

    # Determine output path for metadata
    if output_dir:
        metadata_path = Path(output_dir) / "processing_metadata.json"
    else:
        # Save in current directory if no output_dir specified
        metadata_path = Path("processing_metadata.json")

    # Ensure output directory exists
    metadata_path.parent.mkdir(parents=True, exist_ok=True)

    # Save metadata
    with open(metadata_path, "w") as f:
        json.dump(metadata, f, indent=2)

    return str(metadata_path)


if __name__ == "__main__":

    # Hardcoded parameters
    metas_filepath = "/app2/suno/data/diffusion/v1/metas_diff_v0_tr.jsonl"
    dry_run = False  # Set to False to actually write files
    output_dir = "/app2/suno/data/sft_stems_12_output_v11_vox_trimmed"
    n_jobs = -1  # Use all available cores (-1), or set to specific number like 4
    skip_existing = True  # Set to False to reprocess existing files
    save_metadata = True  # Set to False to skip saving metadata JSON

    print(f"Loading metadata from {metas_filepath}")
    metas = read_jsonl(metas_filepath)

    print("Building artist to vox mappings...")
    artist_id_to_vox_paths = _build_artist_vox_mappings(metas)

    # Count total files to process
    total_files = sum(len(vox_paths) for vox_paths in artist_id_to_vox_paths.values())
    print(
        f"Found {total_files} vocal files to process across {len(artist_id_to_vox_paths)} artists"
    )

    if dry_run:
        print("🔍 DRY RUN MODE - No files will be written")
    else:
        if output_dir:
            print(f"💾 WRITE MODE - Files will be saved to: {output_dir}")
        else:
            print("💾 WRITE MODE - Files will be saved to original directories")

    if skip_existing:
        print("⏭️  RESUMABLE MODE - Skipping existing output files")
    else:
        print("🔄 REPROCESS MODE - Will reprocess all files (including existing ones)")

    if save_metadata:
        print("📄 METADATA SAVING - Will save processing metadata to JSON")

    print(f"🚀 Using {n_jobs} parallel jobs for processing")

    # Flatten all vox_meta entries for parallel processing
    all_vox_metas = []
    for artist_id, vox_paths in artist_id_to_vox_paths.items():
        all_vox_metas.extend(vox_paths)

    print(f"Starting parallel processing of {len(all_vox_metas)} files...")
    start_time = time.time()

    # Process files in parallel with progress tracking
    results = Parallel(n_jobs=n_jobs, verbose=1)(
        delayed(_process_single_file)(vox_meta, dry_run, output_dir, skip_existing)
        for vox_meta in all_vox_metas
    )

    end_time = time.time()
    processing_time = end_time - start_time

    # Count results
    success_count = sum(1 for result in results if result["success"])
    error_count = sum(1 for result in results if not result["success"])
    skipped_count = sum(1 for result in results if result.get("skipped", False))
    processed_count = success_count - skipped_count

    # Print error details
    for result in results:
        if not result["success"]:
            print(f"Error processing {result['input_path']}: {result['error']}")

    # Calculate and display statistics
    avg_time_per_file = processing_time / processed_count if processed_count > 0 else 0
    files_per_second = processed_count / processing_time if processing_time > 0 else 0

    print("\n📊 Processing Statistics:")
    print(f"  Total files: {total_files}")
    print(f"  Successfully processed: {processed_count}")
    print(f"  Skipped (already exist): {skipped_count}")
    print(f"  Errors: {error_count}")
    print(f"  Total time: {processing_time:.2f} seconds")
    if processed_count > 0:
        print(f"  Average time per processed file: {avg_time_per_file:.3f} seconds")
        print(f"  Processing rate: {files_per_second:.2f} files/second")

        estimated_time_for_1000 = (
            (1000 * avg_time_per_file) / n_jobs
            if n_jobs > 0
            else 1000 * avg_time_per_file
        )
        print(
            f"  Estimated time for 1000 new files: {estimated_time_for_1000:.1f} seconds"
        )
    else:
        print("  No new files were processed (all were skipped or failed)")

    # Save metadata if requested
    if save_metadata:
        metadata_path = _save_metadata(
            results=results,
            output_dir=output_dir,
            processing_time=processing_time,
            total_files=total_files,
            processed_count=processed_count,
            skipped_count=skipped_count,
            error_count=error_count,
            n_jobs=n_jobs,
        )
        print(f"\n📄 Metadata saved to: {metadata_path}")