import os
import math
import glob
import torch
import torchaudio
import numpy as np
import pyloudnorm as pyln
import pandas as pd

from tqdm import tqdm
from typing import List, Optional
from suno_utils.utils.text import read_jsonl


def compute_band_energy(
    waveform: torch.Tensor, sample_rate: int, n_fft: int = 4096
) -> float:
    bands = {"bass": (20, 250), "mid": (250, 2500), "high": (2500, 20000)}

    # Convert to mono if stereo
    if waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)

    # Split into frames
    frame_length = n_fft
    hop_length = n_fft // 2  # 50% overlap
    frames = waveform.unfold(1, frame_length, hop_length)
    # Compute FFT for each frame
    spectrum = torch.fft.rfft(frames.squeeze(0))  # [num_frames, n_fft//2 + 1]
    freqs = torch.fft.rfftfreq(n_fft, d=1 / sample_rate)  # [n_fft//2 + 1]

    # Compute magnitudes for each frame
    magnitudes = torch.abs(spectrum)  # [num_frames, n_fft//2 + 1]

    # Compute centroid for each frame
    numerator = torch.sum(
        freqs.view(1, -1) * magnitudes, dim=1
    )  # Sum over frequencies for each frame
    denominator = torch.sum(magnitudes, dim=1)

    results = []

    # Compute mean centroid across all frames
    centroid = torch.mean(numerator / (denominator + 1e-8))
    results.append(float(centroid))

    for band_name, (low_freq, high_freq) in bands.items():
        # Create frequency mask
        mask = (freqs >= low_freq) & (freqs <= high_freq)

        # They should now have the same size
        band_energy = torch.mean(magnitudes * mask)
        results.append(float(band_energy))

    return results


def calculate_stereo_width(waveform):
    # Split into left and right channels
    left = waveform[0]
    right = waveform[1]

    # Compute mid/side representation
    mid = (left + right) / 2
    side = (left - right) / 2

    # Compute RMS energy of mid and side channels
    mid_energy = torch.sqrt(torch.mean(mid**2))
    side_energy = torch.sqrt(torch.mean(side**2))

    # Compute stereo width based on mid/side ratio
    # Normalize to range 0-1 using sigmoid-like function
    width_ratio = (side_energy / (mid_energy + 1e-8)).item()
    stereo_width = 2 * (1 / (1 + np.exp(-width_ratio)) - 0.5)

    return stereo_width.item()


def calculate_crest_factor(signal, window_size):
    """
    Calculate RMS values for windows of audio data using PyTorch.

    Parameters:
    signal (torch.Tensor): Audio signal tensor
    window_size (int): Size of the window for RMS calculation

    Returns:
    torch.Tensor: Tensor of RMS values
    """
    # Ensure input is a tensor
    if not isinstance(signal, torch.Tensor):
        signal = torch.tensor(signal, dtype=torch.float32)

    # convert to mono if stereo
    if signal.shape[0] > 1:
        signal = signal.mean(dim=0)

    # Calculate pad size
    pad_size = window_size - (len(signal) % window_size)
    if pad_size < window_size:
        # Use constant padding (default value is 0)
        padded_signal = torch.nn.functional.pad(signal, (0, pad_size))
    else:
        padded_signal = signal

    # Reshape signal into windows using unfold
    # unfold(dimension, size, step) creates overlapping windows
    # here we use step=size to create non-overlapping windows
    windows = padded_signal.unfold(0, window_size, window_size)

    # Calculate RMS for each window
    # torch.mean along dim=1 averages across the window
    # keepdim=False reduces the dimension
    window_rms = torch.sqrt(torch.mean(windows**2, dim=1))
    rms = torch.mean(window_rms).item()

    # get the peak value
    peak = torch.max(torch.abs(signal)).item()

    # compute crest factor
    crest_factor = peak / (rms + 1e-8)

    return np.log(crest_factor + 1e-8)


import torch
import torchaudio


def measure_silence_percentage(
    waveform: torch.Tensor,
    sample_rate: int,
    silence_threshold_db: float = -60,
    window_size_ms: int = 100,
) -> float:
    """
    Measures the percentage of audio that could be considered silent.

    Args:
        file_path: Path to audio file
        silence_threshold_db: RMS threshold in dB below which audio is considered silent
        window_size_ms: Size of analysis window in milliseconds

    Returns:
        Percentage (0-100) of audio that is below the silence threshold
    """

    # Convert to mono if stereo
    if waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)

    # Calculate window size in samples
    window_size = int(sample_rate * window_size_ms / 1000)

    # Unfold the waveform into windows
    windows = waveform.unfold(1, window_size, window_size)

    # Calculate RMS for each window
    rms = torch.sqrt(torch.mean(windows**2, dim=2))
    db = 20 * torch.log10(rms + 1e-10)

    # Calculate percentage of windows below threshold
    silence_percentage = 100 * torch.mean((db < silence_threshold_db).float()).item()

    return silence_percentage


def measure_spectral_flatness(audio_tensor):
    """
    Calculate spectral flatness (Wiener entropy) of the signal.
    Returns value between 0 (pure tone) and 1 (white noise).

    Parameters:
    audio_tensor: Input audio tensor of shape [..., samples]

    Returns:
    torch.Tensor: Spectral flatness value between 0 and 1
    """
    # Get spectrum magnitude
    audio_tensor = audio_tensor.mean(dim=0)
    spectrum = torch.abs(torch.fft.rfft(audio_tensor, dim=-1))

    # Add small epsilon to avoid log(0)
    epsilon = 1e-8
    spectrum = spectrum + epsilon

    # Calculate geometric mean and arithmetic mean
    log_spectrum = torch.log(spectrum)
    geometric_mean = torch.exp(torch.mean(log_spectrum, dim=-1))
    arithmetic_mean = torch.mean(spectrum, dim=-1)

    # Compute flatness
    flatness = geometric_mean / (arithmetic_mean + 1e-8)

    return flatness.item()


def download_audio(s3_filepath: str, example_id: str, tmp_dir: str):
    filename = os.path.basename(s3_filepath)
    out_filepath = os.path.join(tmp_dir, f"{example_id}-{filename}")
    # only download the file if its not already downloaded
    if not os.path.isfile(out_filepath):
        os.system(f"aws s3 cp {s3_filepath} {out_filepath} > /dev/null 2>&1")
    return out_filepath


import subprocess
from subprocess import TimeoutExpired
import os


def download_audio_with_timeout(
    s3_filepath: str, example_id: str, tmp_dir: str, timeout: int = 300
):
    filename = os.path.basename(s3_filepath)
    # check if the file extension is mp3, mp4, m4a, wav, webm, or ogg
    if not filename.endswith((".mp3", ".mp4", ".m4a", ".wav", ".webm", ".ogg")):
        raise ValueError(f"Unsupported file extension: {filename}")
    out_filepath = os.path.join(tmp_dir, f"{example_id}-{filename}")

    if not os.path.isfile(out_filepath):
        try:
            subprocess.run(
                ["aws", "s3", "cp", s3_filepath, out_filepath],
                capture_output=True,
                timeout=timeout,
                check=True,
            )
        except TimeoutExpired:
            if os.path.exists(out_filepath):
                os.remove(out_filepath)
            raise

    return out_filepath


class AudioProductionFeatures(torch.utils.data.Dataset):
    def __init__(
        self,
        metas: List[dict],
        sample_rate: int = 48000,
        target_loudness: float = -16.0,
        tmp_dir: str = "/tmp/cjs",
    ):
        self.metas = metas
        self.sample_rate = sample_rate
        self.target_loudness = target_loudness
        self.meter = pyln.Meter(self.sample_rate)

        # we need to filter metas to only use ones with s3_filepath
        # but, we have to be careful to maintain the index of the original metadata
        filtered_metas = []
        for meta_idx, meta in enumerate(metas):
            if "audio_filepath" in meta:
                meta["orig_idx"] = meta_idx
                filtered_metas.append(meta)
            elif "s3_filepath" in meta:
                meta["orig_idx"] = meta_idx
                meta["audio_filepath"] = meta["s3_filepath"]
                filtered_metas.append(meta)

        num_filtered = len(filtered_metas)
        num_original = len(metas)
        percent_remaining = (num_filtered / num_original) * 100
        print(
            f"{num_filtered}/{num_original} ({percent_remaining:0.2f}%) examples have s3_filepath."
        )
        self.metas = filtered_metas
        self.tmp_dir = tmp_dir

        os.makedirs(self.tmp_dir, exist_ok=True)

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

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

        # check if the file exists in the local directory
        local_filepath_dir = "/app2/suno/data/raw_audio_opus_v0"
        local_filepath = os.path.join(local_filepath_dir, meta["id"] + ".opus")

        if os.path.exists(local_filepath):
            filepath = local_filepath
        else:
            # check for filepath ("audio_filepath")
            if "audio_filepath" in meta:
                filepath = meta["audio_filepath"]
            elif "s3_filepath" in meta:
                filepath = meta["s3_filepath"]
            else:
                raise ValueError(f"No filepath found in meta: {meta}")

            try:
                filepath = download_audio_with_timeout(
                    meta["audio_filepath"], meta["id"], self.tmp_dir
                )
            except Exception as e:
                print(f"Error: {e}")
                print(f"Failed to download: {filepath}")
                return meta["id"], torch.tensor([np.nan] * 6)

        try:
            audio, sample_rate = torchaudio.load(filepath)
            # crop audio to max of 8min
            audio = audio[:, : int(8 * 60 * sample_rate)]

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

        except Exception as e:
            print(f"Error: {e}")
            print(f"Failed to load: {filepath}")
            return meta["id"], torch.tensor([np.nan] * 6)

        # delete audio file
        os.remove(filepath)

        # resample if necessary
        if sample_rate != self.sample_rate:
            audio = torchaudio.functional.resample(audio, sample_rate, self.sample_rate)

        # loudness normalization
        try:
            loudness = self.meter.integrated_loudness(audio.permute(1, 0).numpy())
        except Exception as e:
            print(f"Error: {e}")
            print(f"Failed to normalize loudness: {filepath}")
            return meta["id"], torch.tensor([np.nan] * 6)
        # check if loudness is -inf
        if loudness == -np.inf:
            loudness = -80.0
        loudness_diff = self.target_loudness - loudness
        # limit the loudness difference to +/- 20 db
        loudness_diff = np.clip(loudness_diff, -20, 20)
        audio *= 10 ** (loudness_diff / 20.0)
        # this can cause problems

        # compute features
        with torch.no_grad():
            results_spectral = compute_band_energy(audio, self.sample_rate)
            results_crest_factor = calculate_crest_factor(audio, 1024)
            results_stereo_width = calculate_stereo_width(audio)
            results_spectral_flatness = measure_spectral_flatness(audio)
            results_silence_percentage = measure_silence_percentage(
                audio, self.sample_rate
            )

        results = results_spectral + [
            results_crest_factor,
            results_stereo_width,
            results_spectral_flatness,
            results_silence_percentage,
            loudness,
        ]

        return meta["id"], torch.tensor(results)


def collate_fn(batch):
    return batch


def save_checkpoint(
    data: list, output_filepath: str, columns: list, tmp_suffix: str = ".tmp"
):
    """
    Safely save data to CSV with temporary file to prevent corruption.

    Args:
        data: List of rows to save
        output_filepath: Path to save the CSV
        columns: Column names for the DataFrame
        tmp_suffix: Suffix for temporary file
    """
    tmp_filepath = output_filepath + tmp_suffix
    df = pd.DataFrame(data, columns=columns)

    # Save to temporary file first
    df.to_csv(tmp_filepath, index=False)

    # If successful, rename to final filename
    os.replace(tmp_filepath, output_filepath)


def load_existing_data(filepath: str) -> tuple[Optional[pd.DataFrame], set]:
    """
    Load existing data and return DataFrame and set of processed IDs.

    Args:
        filepath: Path to the CSV file

    Returns:
        Tuple of (DataFrame or None, set of processed IDs)
    """
    if os.path.exists(filepath):
        try:
            data = pd.read_csv(filepath)
            print(f"Loaded {len(data)} rows from {filepath}")
            return data, set(data["id"])
        except Exception as e:
            print(f"Error loading existing data: {e}")
            # Backup corrupted file
            if os.path.exists(filepath):
                backup_path = filepath + ".corrupted"
                os.rename(filepath, backup_path)
                print(f"Moved potentially corrupted file to {backup_path}")
            return None, set()
    return None, set()


if __name__ == "__main__":
    num_workers = 220

    # metas_filepath = "/home/christian/code/christian/metadata/genius_hq_metas.jsonl"
    # output_filepath = "/home/christian/code/christian/metadata/genius_hq_audio_production_features_v2.csv"

    # metas_filepath = "/home/christian/code/christian/metadata/youtube_music_metas.jsonl"
    # output_filepath = "/home/christian/code/christian/metadata/youtube_music_audio_production_features_v2.csv"

    # metas_filepath = (
    #    "/home/christian/code/christian/metadata/v4/discogs_subset_metas_v0.jsonl"
    # )

    metas_filepath = "/app/suno/tmp/clean_deezer_v0_metas.jsonl"

    output_filepath = "/home/christian/code/christian/metadata/v4/deezer_audio_production_features_v2.csv"
    checkpoint_interval = 10_000  # Save every 1000 samples

    # Create output directory if it doesn't exist
    output_dir = os.path.dirname(output_filepath)
    os.makedirs(output_dir, exist_ok=True)

    # Load metas
    metas = read_jsonl(metas_filepath)

    # Load existing data and get processed IDs
    existing_data, processed_ids = load_existing_data(output_filepath)

    # Filter metas to only include unprocessed items
    metas = [meta for meta in metas if meta["id"] not in processed_ids]
    print(
        f"Found {len(metas)} unprocessed items out of {len(processed_ids)} existing items"
    )

    # Define columns for consistency
    columns = [
        "id",
        "spectral_centroid",
        "bass",
        "mid",
        "high",
        "crest_factor",
        "stereo_width",
        "spectral_flatness",
        "silence_percentage",
        "loudness",
    ]

    # Initialize dataset and dataloader
    dataset = AudioProductionFeatures(metas, sample_rate=48000, target_loudness=-16.0)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=1,
        shuffle=False,
        num_workers=num_workers,
        persistent_workers=True,
    )

    # Initialize new data list
    new_data = []

    try:
        for meta_id, results in tqdm(dataloader):
            result_list = results[0].tolist()
            new_data.append([meta_id[0]] + result_list)

            # Save checkpoint at regular intervals
            if len(new_data) % checkpoint_interval == 0:
                # If we have existing data, combine it with new data
                if existing_data is not None:
                    combined_data = pd.concat(
                        [existing_data, pd.DataFrame(new_data, columns=columns)],
                        ignore_index=True,
                    )
                    combined_data.to_csv(output_filepath, index=False)
                else:
                    save_checkpoint(new_data, output_filepath, columns)

                print(f"Saved checkpoint with {len(new_data)} new samples")

    except KeyboardInterrupt:
        print("\nInterrupted by user. Saving final checkpoint...")
    except Exception as e:
        print(f"\nEncountered error: {e}")
        print("Attempting to save final checkpoint...")
    finally:
        # Save final checkpoint if there's any unsaved data
        if new_data:
            if existing_data is not None:
                combined_data = pd.concat(
                    [existing_data, pd.DataFrame(new_data, columns=columns)],
                    ignore_index=True,
                )
                combined_data.to_csv(output_filepath, index=False)
            else:
                save_checkpoint(new_data, output_filepath, columns)
            print(f"Saved final checkpoint with {len(new_data)} new samples")