from __future__ import annotations

"""
extract_annotation_clips_contributor.py

Purpose
-------
Extract raw audio clips from WAV + Audacity label files (no database required),
and keep the contributor clip library fully synchronised with the current
contents of 2_labels.

Expected structure
------------------
nocturnal_bird_calls_contributor/
  0_scripts/
  1_raw_audio/
  2_labels/
  3_clips/

Label format (Audacity)
-----------------------
start_s    end_s    label

Expected label convention:
    confidence|call_type

Examples:
    12.345    12.600    high|UB-PEW-01
    25.000    25.200    medium|Sacred Kingfisher
    40.500    40.650    low|
    55.000    55.100    zero|

Interpretation
--------------
- Left of "|" = certainty that the sound is a bird:
    high / medium / low / zero
- Right of "|" = call type classification
- If right side is blank:
    * zero|    -> "unknown nonbird call type"
    * otherwise -> "unknown bird call type"

ASTRO_NIGHT rows are ignored.

Audio processing
----------------
No audio enhancement is applied.

Extracted clips are written as raw buffered excerpts from the source WAV.
Any noise reduction, filtering, amplification, or other cleanup should be done
manually on selected clips intended for website use.

Sync behaviour
--------------
This script treats 3_clips as a rebuildable derivative of 2_labels.

By default, it:
- creates any missing clip WAV/PNG pairs implied by current labels
- updates clips if --overwrite is used
- removes orphan clip WAV/PNG files not implied by current labels
- removes orphan or empty clip folders

If a label file exists but the matching WAV cannot be found, this is reported
clearly and those expected clips are NOT added to the synced output set.

Clip naming
-----------
Clip filenames are deterministically derived from annotation start time in
milliseconds.

If multiple annotations share the same rounded start time within a recording,
suffixes are assigned by sorting on start time, end time, and label text.

Examples:
    DELC_20260321_025900_ann00012345.wav
    DELC_20260321_025900_ann00012345a.wav
    DELC_20260321_025900_ann00012345b.wav
"""

import argparse
import re
import shutil
from collections import defaultdict
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import soundfile as sf


# ------------------------------------------------------------
# Paths
# ------------------------------------------------------------

def root() -> Path:
    return Path(__file__).resolve().parent.parent


RAW = root() / "1_raw_audio"
LABELS = root() / "2_labels"
CLIPS = root() / "3_clips"


# ------------------------------------------------------------
# Label parsing
# ------------------------------------------------------------

def parse_label_file(path: Path) -> list[tuple[float, float, str]]:
    rows: list[tuple[float, float, str]] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            parts = line.rstrip("\n").split("\t")
            if len(parts) < 3:
                continue

            start_text, end_text, label_text = parts[0], parts[1], parts[2]

            try:
                start_s = float(start_text)
                end_s = float(end_text)
            except ValueError:
                continue

            label = label_text.strip()
            if not label:
                continue

            rows.append((start_s, end_s, label))

    return rows


def is_astro_night_label(label: str) -> bool:
    return label.strip().upper() == "ASTRO_NIGHT"


def parse_annotation_label(raw_label: str) -> tuple[str, str]:
    """
    Parse labels of the form:
        confidence|call_type

    Examples:
        "high|UB-PEW-01" -> ("high", "UB-PEW-01")
        "medium|Sacred Kingfisher" -> ("medium", "Sacred Kingfisher")
        "low|" -> ("low", "")
        "zero|" -> ("zero", "")
        "YFHE" -> ("", "YFHE")   # fallback if no pipe present
    """
    text = raw_label.strip()

    if "|" in text:
        left, right = text.split("|", 1)
        return left.strip().lower(), right.strip()

    return "", text.strip()


def sanitise_filename_component(text: str) -> str:
    """
    Make a string safe for Windows file/folder names.
    """
    text = text.strip()
    text = re.sub(r'[<>:"/\\|?*]', "_", text)
    text = text.rstrip(" .")
    return text or "UNLABELLED"


def call_type_folder_name(raw_label: str) -> str:
    confidence, call_type = parse_annotation_label(raw_label)

    if call_type:
        folder = call_type
    else:
        if confidence == "zero":
            folder = "unknown nonbird call type"
        else:
            folder = "unknown bird call type"

    return sanitise_filename_component(folder)


# ------------------------------------------------------------
# WAV lookup
# ------------------------------------------------------------

def recording_id_from_label_file(label_path: Path) -> str:
    recording_id = label_path.stem
    for suffix in ["_night", "_labels"]:
        if recording_id.endswith(suffix):
            recording_id = recording_id[: -len(suffix)]
    return recording_id


def find_matching_wav(recording_id: str) -> Path | None:
    matches = list(RAW.rglob(f"{recording_id}.wav"))
    if not matches:
        return None
    if len(matches) > 1:
        matches = sorted(matches)
    return matches[0]


# ------------------------------------------------------------
# Stable clip naming
# ------------------------------------------------------------

def start_ms_code(start_s: float) -> int:
    return int(round(start_s * 1000.0))


def suffix_from_index(idx: int) -> str:
    """
    0 -> ''
    1 -> 'a'
    2 -> 'b'
    ...
    26 -> 'z'
    27 -> 'aa'
    """
    if idx == 0:
        return ""

    idx -= 1
    letters: list[str] = []
    while True:
        letters.append(chr(ord("a") + (idx % 26)))
        idx = idx // 26 - 1
        if idx < 0:
            break
    return "".join(reversed(letters))


def build_clip_base_name(recording_id: str, start_s: float, duplicate_index: int) -> str:
    ms = start_ms_code(start_s)
    suffix = suffix_from_index(duplicate_index)
    return f"{recording_id}_ann{ms:08d}{suffix}"


def build_annotation_entries(
    recording_id: str,
    rows: list[tuple[float, float, str]],
) -> list[dict]:
    """
    Convert parsed label rows into deterministic annotation entries with stable
    clip base names.

    Duplicate start times within a recording are handled by assigning suffixes
    a, b, c ... after sorting by end time and label text.
    """
    non_astro_rows = [
        {"start_s": start_s, "end_s": end_s, "label": label}
        for start_s, end_s, label in rows
        if not is_astro_night_label(label)
    ]

    grouped: dict[int, list[dict]] = defaultdict(list)
    for row in non_astro_rows:
        grouped[start_ms_code(row["start_s"])].append(row)

    entries: list[dict] = []
    for ms_code in sorted(grouped.keys()):
        group = grouped[ms_code]
        group_sorted = sorted(
            group,
            key=lambda r: (
                r["start_s"],
                r["end_s"],
                r["label"].lower(),
            ),
        )
        for duplicate_index, row in enumerate(group_sorted):
            row = dict(row)
            row["base_name"] = build_clip_base_name(
                recording_id=recording_id,
                start_s=row["start_s"],
                duplicate_index=duplicate_index,
            )
            entries.append(row)

    return entries


# ------------------------------------------------------------
# Audio utils
# ------------------------------------------------------------

def to_mono_for_display(audio: np.ndarray) -> np.ndarray:
    if audio.ndim == 1:
        return audio
    return audio.mean(axis=1)


def extract_clip(
    audio: np.ndarray,
    sr: int,
    start_s: float,
    end_s: float,
    pre: float,
    post: float,
) -> np.ndarray:
    start_idx = max(0, int(round((start_s - pre) * sr)))
    end_idx = min(len(audio), int(round((end_s + post) * sr)))
    return audio[start_idx:end_idx]


# ------------------------------------------------------------
# Spectrogram
# ------------------------------------------------------------

def save_spec(audio: np.ndarray, sr: int, out_path: Path) -> None:
    display_audio = to_mono_for_display(audio)

    plt.figure(figsize=(8, 3))
    plt.specgram(display_audio, Fs=sr, NFFT=1024, noverlap=512, cmap="gray_r")
    plt.ylim(0, 10000)
    plt.xlabel("Time (s)")
    plt.ylabel("Frequency (Hz)")
    plt.tight_layout()
    plt.savefig(out_path, dpi=150)
    plt.close()


# ------------------------------------------------------------
# Sync helpers
# ------------------------------------------------------------

def remove_empty_dirs(root_dir: Path) -> int:
    removed = 0
    for path in sorted(
        [p for p in root_dir.rglob("*") if p.is_dir()],
        key=lambda p: len(p.parts),
        reverse=True,
    ):
        try:
            path.rmdir()
            removed += 1
        except OSError:
            pass
    return removed


def sync_output_structure(
    output_dir: Path,
    expected_call_types: set[str],
    expected_files: set[Path],
) -> tuple[int, int]:
    removed_files = 0
    removed_dirs = 0

    for path in output_dir.rglob("*"):
        if not path.is_file():
            continue

        if path.suffix.lower() not in {".wav", ".png"}:
            continue

        if path not in expected_files:
            path.unlink()
            removed_files += 1
            print(f"Removed orphaned file: {path.relative_to(output_dir).as_posix()}")

    for child in output_dir.iterdir():
        if not child.is_dir():
            continue

        if child.name not in expected_call_types:
            shutil.rmtree(child)
            removed_dirs += 1
            print(f"Removed orphaned folder: {child.relative_to(output_dir).as_posix()}")

    removed_dirs += remove_empty_dirs(output_dir)

    return removed_files, removed_dirs


# ------------------------------------------------------------
# Main
# ------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Extract raw contributor clips from WAV + label files and keep 3_clips synced to 2_labels."
    )
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--pre", type=float, default=2.0)
    parser.add_argument("--post", type=float, default=1.0)
    parser.add_argument(
        "--no-sync",
        action="store_true",
        help="Do not remove orphaned clip files/folders from 3_clips",
    )
    args = parser.parse_args()

    CLIPS.mkdir(exist_ok=True)

    label_files = sorted({
        *LABELS.rglob("*_night.txt"),
        *LABELS.rglob("*_labels.txt"),
    })

    created = 0
    updated = 0
    skipped = 0
    missing_wavs = 0
    failed = 0

    expected_files: set[Path] = set()
    expected_call_types: set[str] = set()

    for lf in label_files:
        recording_id = recording_id_from_label_file(lf)
        wav = find_matching_wav(recording_id)

        if wav is None:
            print(f"WARNING: No matching WAV found for label file: {lf.relative_to(root())}")
            missing_wavs += 1
            continue

        try:
            audio, sr = sf.read(wav)
        except Exception as e:
            print(f"Failed reading WAV for {recording_id}: {e}")
            failed += 1
            continue

        try:
            rows = parse_label_file(lf)
            entries = build_annotation_entries(recording_id=recording_id, rows=rows)
        except Exception as e:
            print(f"Failed reading label file {lf.relative_to(root())}: {e}")
            failed += 1
            continue

        for entry in entries:
            start = entry["start_s"]
            end = entry["end_s"]
            label = entry["label"]
            base_name = entry["base_name"]

            folder_name = call_type_folder_name(label)
            expected_call_types.add(folder_name)

            out_dir = CLIPS / folder_name
            wav_out = out_dir / f"{base_name}.wav"
            png_out = out_dir / f"{base_name}.png"

            expected_files.add(wav_out)
            expected_files.add(png_out)

            files_exist = wav_out.exists() and png_out.exists()

            if files_exist and not args.overwrite:
                skipped += 1
                continue

            try:
                clip = extract_clip(audio, sr, start, end, args.pre, args.post)
                if clip.size == 0:
                    continue

                clip = np.asarray(clip, dtype=np.float32)

                out_dir.mkdir(parents=True, exist_ok=True)
                existed_before = wav_out.exists() or png_out.exists()

                sf.write(wav_out, clip, sr)
                save_spec(clip, sr, png_out)

                if existed_before:
                    updated += 1
                else:
                    created += 1

            except Exception as e:
                print(
                    f"Failed for {recording_id} "
                    f"({base_name}, {label!r}): {e}"
                )
                failed += 1
                continue

    removed_files = 0
    removed_dirs = 0
    if not args.no_sync:
        removed_files, removed_dirs = sync_output_structure(
            output_dir=CLIPS,
            expected_call_types=expected_call_types,
            expected_files=expected_files,
        )

    print()
    print(
        "Done. "
        f"Created {created}, updated {updated}, skipped {skipped}, "
        f"missing WAVs {missing_wavs}, failed {failed}, "
        f"removed files {removed_files}, removed dirs {removed_dirs}"
    )


if __name__ == "__main__":
    main()