import os
import wave
import numpy as np
import subprocess
import tempfile


def filter_audio_list(
    input_path="train_list.txt",
    output_path="train_list_filtered.txt",
    min_duration_sec=1.25,
    max_duration_sec=10.0,
    min_amplitude_threshold=0.02,
    max_amplitude_threshold=0.94,
    ffmpeg_path="ffmpeg",
):
    """
    Filter lines in input_path according to duration and amplitude checks.
    Removes duplicate paths: if a given .wav path appears more than once in the
    input file, ALL occurrences are removed from the output.
    """
    if not os.path.exists(input_path):
        print(f"Error: Input file not found at '{input_path}'")
        return

    kept_lines = []
    removed_count = 0
    total_count = 0

    print(f"Starting to process '{input_path}'...")
    print(f" - Minimum duration: {min_duration_sec}s")
    print(f" - Maximum duration: {max_duration_sec}s")
    print(f" - Minimum amplitude: {min_amplitude_threshold}")
    print(f" - Maximum amplitude: {max_amplitude_threshold}")
    print(
        f" - Target format: 44.1kHz, 16-bit PCM, mono "
        f"(will convert using {ffmpeg_path})"
    )

    # Read input lines (preserve full line content for output)
    with open(input_path, "r", encoding="utf-8") as f:
        raw_lines = f.readlines()

    # Build canonical path for each non-empty line so we can detect duplicates
    base_dir = os.path.dirname(os.path.abspath(input_path)) or os.getcwd()
    canon_for_line = [None] * len(raw_lines)
    canon_to_indices = {}

    for idx, raw in enumerate(raw_lines):
        ln = raw.strip()
        if not ln:
            continue
        audio_rel = ln.split("|")[0].strip()
        # Match original behavior of prefixing "../" relative to the input file
        candidate = os.path.join(base_dir, "..", audio_rel)
        canonical = os.path.realpath(os.path.abspath(candidate))
        canon_for_line[idx] = canonical
        canon_to_indices.setdefault(canonical, []).append(idx)

    # Identify duplicate canonical paths (appear more than once)
    duplicates_canonical = {p for p, idxs in canon_to_indices.items() if len(idxs) > 1}
    if duplicates_canonical:
        print(
            f"  [Deduplication] Found {len(duplicates_canonical)} duplicate "
            "path(s). All occurrences will be removed."
        )
        for p in sorted(duplicates_canonical):
            print(f"    - {p} (occurrences={len(canon_to_indices[p])})")

    # Process lines, skipping duplicates early
    for idx, raw in enumerate(raw_lines):
        total_count += 1
        line = raw.strip()
        if not line:
            continue

        canonical = canon_for_line[idx]
        if canonical in duplicates_canonical:
            # remove duplicate occurrences (do not keep any of them)
            print(f"  [Removing] Duplicate entry: {line.split('|')[0].strip()}")
            removed_count += 1
            continue

        # Use the canonical absolute path for filesystem operations
        audio_path = canonical

        try:
            if not os.path.exists(audio_path):
                print(f"  [Warning] File not found, skipping: {audio_path}")
                removed_count += 1
                continue

            # Try to open as WAV; if not readable, convert with ffmpeg
            try:
                wav_file = wave.open(audio_path, "rb")
            except wave.Error:
                print(
                    f"  [Converting] Not a readable WAV: converting {audio_path} -> "
                    "44.1kHz 16-bit mono WAV"
                )
                tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
                tmp.close()
                cmd = [
                    ffmpeg_path,
                    "-y",
                    "-i",
                    audio_path,
                    "-ar",
                    "44100",
                    "-ac",
                    "1",
                    "-acodec",
                    "pcm_s16le",
                    tmp.name,
                ]
                proc = subprocess.run(
                    cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
                )
                if proc.returncode != 0:
                    stderr = proc.stderr.decode("utf-8", errors="replace")
                    print(
                        f"  [Error] ffmpeg conversion failed for {audio_path}: "
                        f"{stderr}"
                    )
                    removed_count += 1
                    try:
                        os.unlink(tmp.name)
                    except Exception:
                        pass
                    continue

                try:
                    os.replace(tmp.name, audio_path)
                except Exception as e:
                    print(
                        f"  [Error] Could not replace original with converted file: "
                        f"{e}"
                    )
                    removed_count += 1
                    try:
                        os.unlink(tmp.name)
                    except Exception:
                        pass
                    continue

                try:
                    wav_file = wave.open(audio_path, "rb")
                except wave.Error as e:
                    print(
                        f"  [Warning] Still cannot open after conversion: "
                        f"{audio_path} ({e})"
                    )
                    removed_count += 1
                    continue

            # Read header info then close
            try:
                num_frames = wav_file.getnframes()
                framerate = wav_file.getframerate()
                sampwidth = wav_file.getsampwidth()
                n_channels = wav_file.getnchannels()
            finally:
                try:
                    wav_file.close()
                except Exception:
                    pass

            # Convert to target format if needed
            needs_conversion = (
                framerate != 44100 or sampwidth != 2 or n_channels != 1
            )
            if needs_conversion:
                print(
                    f"  [Converting] {audio_path} (sr={framerate}, sw={sampwidth}, "
                    f"ch={n_channels}) -> 44.1kHz 16-bit mono"
                )
                tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
                tmp.close()
                cmd = [
                    ffmpeg_path,
                    "-y",
                    "-i",
                    audio_path,
                    "-ar",
                    "44100",
                    "-ac",
                    "1",
                    "-acodec",
                    "pcm_s16le",
                    tmp.name,
                ]
                proc = subprocess.run(
                    cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
                )
                if proc.returncode != 0:
                    stderr = proc.stderr.decode("utf-8", errors="replace")
                    print(
                        f"  [Error] ffmpeg conversion failed for {audio_path}: "
                        f"{stderr}"
                    )
                    removed_count += 1
                    try:
                        os.unlink(tmp.name)
                    except Exception:
                        pass
                    continue

                try:
                    os.replace(tmp.name, audio_path)
                except Exception as e:
                    print(
                        f"  [Error] Could not replace original with converted file: "
                        f"{e}"
                    )
                    removed_count += 1
                    try:
                        os.unlink(tmp.name)
                    except Exception:
                        pass
                    continue

            # Open the (possibly converted) WAV and perform checks
            try:
                with wave.open(audio_path, "rb") as wf:
                    num_frames = wf.getnframes()
                    framerate = wf.getframerate()
                    sampwidth = wf.getsampwidth()
                    n_channels = wf.getnchannels()

                    # Duration checks
                    duration = num_frames / float(framerate)
                    if duration < min_duration_sec:
                        print(
                            f"  [Removing] Short duration: {audio_path} "
                            f"({duration:.2f}s)"
                        )
                        removed_count += 1
                        continue

                    if duration > max_duration_sec:
                        print(
                            f"  [Removing] Too long: {audio_path} "
                            f"({duration:.2f}s > {max_duration_sec}s)"
                        )
                        removed_count += 1
                        continue

                    # Amplitude check (16-bit PCM)
                    if sampwidth == 2:
                        frames = wf.readframes(num_frames)
                        if not frames:
                            print(f"  [Removing] Empty audio: {audio_path}")
                            removed_count += 1
                            continue

                        samples = np.frombuffer(frames, dtype=np.int16)
                        if samples.size == 0:
                            print(f"  [Removing] Empty samples: {audio_path}")
                            removed_count += 1
                            continue

                        normalized_samples = samples.astype(np.float32) / 32768.0
                        max_abs_amplitude = float(
                            np.max(np.abs(normalized_samples))
                        )

                        if max_abs_amplitude > max_amplitude_threshold:
                            print(
                                f"  [Removing] High amplitude: {audio_path} "
                                f"(Peak: {max_abs_amplitude:.2f})"
                            )
                            removed_count += 1
                            continue

                        if max_abs_amplitude < min_amplitude_threshold:
                            print(
                                f"  [Removing] Low amplitude: {audio_path} "
                                f"(Peak: {max_abs_amplitude:.2f})"
                            )
                            removed_count += 1
                            continue
                    else:
                        print(
                            f"  [Info] Skipping amplitude check for {audio_path} "
                            f"(sample width: {sampwidth})"
                        )

            except wave.Error as e:
                print(
                    f"  [Warning] Could not read WAV after conversion/open: "
                    f"{audio_path} ({e})"
                )
                removed_count += 1
                continue

            # Passed all checks
            kept_lines.append(line + "\n")

        except wave.Error as e:
            print(
                f"  [Warning] Could not read WAV file, skipping: {audio_path} ({e})"
            )
            removed_count += 1
        except FileNotFoundError as e:
            print(f"  [Warning] File disappeared, skipping: {audio_path} ({e})")
            removed_count += 1
        except Exception as e:
            print(f"  [Error] An unexpected error occurred on line: {line}")
            print(f"          {e}")
            removed_count += 1

    with open(output_path, "w", encoding="utf-8") as f:
        f.writelines(kept_lines)

    print("\n--- Processing Complete ---")
    print(f"Total entries processed: {total_count}")
    print(f"Entries kept: {len(kept_lines)}")
    print(f"Entries removed (duration/amplitude/errors/dupes): {removed_count}")
    print(f"Filtered list saved to: '{output_path}'")


if __name__ == "__main__":
    # Ensure numpy and ffmpeg are available.
    filter_audio_list(
        input_path="train_3_filtered.txt",
        output_path="train_3_filtered_1.txt",
        min_duration_sec=1.25,
        max_duration_sec=10.0,
        min_amplitude_threshold=0.02,
        max_amplitude_threshold=0.93,
        ffmpeg_path="ffmpeg",
    )