import os
import pandas as pd
import librosa
import requests
import io
import soundfile as sf
import numpy as np
from tqdm import tqdm


class AudioPreprocessor:
    def __init__(self, target_sr=16000, chunk_duration=3, stride_duration=1):
        self.target_sr = target_sr
        self.chunk_duration = chunk_duration
        self.stride_duration = stride_duration
        self.samples_per_chunk = int(target_sr * chunk_duration)
        self.samples_per_stride = int(target_sr * stride_duration)

    def download_audio(self, url):
        """Downloads audio from URL into memory."""
        try:
            response = requests.get(url, timeout=15)
            response.raise_for_status()
            return io.BytesIO(response.content)
        except Exception as e:
            print(f"Failed to download {url}: {e}")
            return None

    def process_and_split(self, audio_bytes, trim_silence=True, peak_norm=True):
        """Converts to mono, resamples, optionally trims silence, and splits using a rolling window."""
        try:
            audio_bytes.seek(0)
            y, sr = librosa.load(audio_bytes, sr=self.target_sr, mono=True)

            if len(y) == 0:
                print("Warning: Downloaded audio is empty.")
                return []

            # Trim leading/trailing silence (gentle threshold to avoid cutting speech)
            if trim_silence and len(y) > self.samples_per_chunk:
                y_trim, _ = librosa.effects.trim(
                    y, top_db=25, frame_length=2048, hop_length=512
                )
                if len(y_trim) >= self.samples_per_chunk:
                    y = y_trim

            chunks = []
            # OVERLAPPING WINDOW: Step forward by samples_per_stride (1 second) instead of full chunk
            for i in range(0, len(y), self.samples_per_stride):
                chunk = y[i : i + self.samples_per_chunk].copy()

                # Keep chunks that are at least 2.5 seconds long.
                # (Shorter trailing pieces are safely dropped because the overlap caught them earlier)
                min_required_samples = int(self.target_sr * 2.5)
                if len(chunk) < min_required_samples:
                    continue

                # Peak-normalize chunk so level is consistent across chunks
                if peak_norm:
                    peak = np.abs(chunk).max()
                    if peak > 1e-8:
                        chunk = chunk / peak
                chunks.append(chunk)
            return chunks
        except Exception as e:
            print(f"Error processing audio: {e}")
            return []

    def run_pipeline(self, csv_path, output_dir):
        """Main execution loop."""
        df = pd.read_csv(csv_path)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        processed_data = []

        print(f"Starting processing of {len(df)} files...")
        for _, row in tqdm(df.iterrows(), total=len(df)):
            audio_io = self.download_audio(row["audio_url"])
            if audio_io:
                chunks = self.process_and_split(audio_io)

                for idx, chunk in enumerate(chunks):
                    chunk_filename = f"{row['dp_id']}_chunk{idx}.wav"
                    chunk_path = os.path.join(output_dir, chunk_filename)

                    # Save chunk as .wav for the training pipeline
                    sf.write(chunk_path, chunk, self.target_sr)

                    processed_data.append(
                        {
                            "chunk_path": chunk_path,
                            "original_id": row["dp_id"],
                            "label": row["nativity_status"],
                            "language": row["language"],
                        }
                    )

        # Save a new manifest for the ML model
        manifest_df = pd.DataFrame(processed_data)
        manifest_df.to_csv("manifests/processed_manifest.csv", index=False)
        print(f"Done! Saved {len(manifest_df)} chunks to {output_dir}")


if __name__ == "__main__":
    # Usage
    preprocessor = AudioPreprocessor(chunk_duration=3, stride_duration=1)
    # Ensure your csv has columns: 'audio_url', 'dp_id', 'nativity_status', 'language'
    preprocessor.run_pipeline("manifests/input_data.csv", "data/processed_audio_chunks")