from __future__ import annotations

"""
prepare_recordings_contributor.py

Purpose
-------
Scan contributor WAV files, calculate overlap with astronomical night, and
create Audacity label files marking ASTRO_NIGHT sections.

This is a lightweight contributor version of the main NFC prepare_recordings
workflow. It does NOT use the project database.

Expected folder structure
-------------------------
nocturnal_bird_calls_contributor/
  0_scripts/
    prepare_recordings_contributor.py
    sites.csv
  1_raw_audio/
  2_labels/
  3_clips/

Expected filename format
------------------------
    SITE_YYYYMMDD_HHMMSS.wav

Example:
    DELC_20260321_025900.wav

Expected sites.csv format
-------------------------
site_code,site_name,latitude,longitude,timezone
DELC,Delaney Circuit,-27.5200,153.1100,Australia/Brisbane

Outputs
-------
For each WAV file:
- matching Audacity label file in 2_labels
- summary CSV in 2_labels/astro_night_summary.csv

Typical usage
-------------
From 0_scripts:
    python prepare_recordings_contributor.py

Overwrite existing label files:
    python prepare_recordings_contributor.py --overwrite-labels
"""

import argparse
import csv
import re
import sys
import wave
from dataclasses import dataclass
from datetime import date, datetime, timedelta
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from zoneinfo import ZoneInfo

from astral import LocationInfo
from astral.sun import sun


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

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


def default_raw_audio_root() -> Path:
    return get_contributor_root() / "1_raw_audio"


def default_label_output_root() -> Path:
    return get_contributor_root() / "2_labels"


def default_summary_csv() -> Path:
    return default_label_output_root() / "astro_night_summary.csv"


def default_sites_csv() -> Path:
    return Path(__file__).resolve().parent / "sites.csv"


def path_relative_to_root(path: Path, root: Path) -> str:
    return path.resolve().relative_to(root.resolve()).as_posix()


def normalised_relative_subdir(
    wav_path: Path,
    raw_audio_root: Path,
    site_code: str,
) -> Path:
    """
    Return a clean relative subdirectory for mirrored outputs.
    """
    relative_subdir = wav_path.parent.resolve().relative_to(raw_audio_root.resolve())

    if relative_subdir == Path("."):
        return Path(site_code)

    parts = list(relative_subdir.parts)
    if parts and parts[0].upper() == site_code.upper():
        return Path(*parts)

    return relative_subdir


# ------------------------------------------------------------
# Filename pattern
# ------------------------------------------------------------

FILENAME_RE = re.compile(
    r"^(?P<site>[A-Z0-9]+)_(?P<ymd>\d{8})_(?P<hms>\d{6})\.wav$",
    re.IGNORECASE,
)


# ------------------------------------------------------------
# Data classes
# ------------------------------------------------------------

@dataclass
class SiteInfo:
    site_code: str
    site_name: str
    latitude: float
    longitude: float
    timezone: str


@dataclass
class RecordingInfo:
    wav_path: Path
    recording_id: str
    site_code: str
    recording_date: str
    start_datetime: str
    start_dt: datetime
    end_dt: datetime
    duration_s: float
    sample_rate_hz: int
    channels: int
    relative_path: str


# ------------------------------------------------------------
# Site loading
# ------------------------------------------------------------

def load_sites_from_csv(csv_path: Path) -> Dict[str, SiteInfo]:
    if not csv_path.exists():
        raise FileNotFoundError(f"sites.csv not found: {csv_path}")

    sites: Dict[str, SiteInfo] = {}

    with csv_path.open("r", newline="", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        required = {"site_code", "site_name", "latitude", "longitude", "timezone"}
        missing = required - set(reader.fieldnames or [])
        if missing:
            raise ValueError(
                f"sites.csv is missing required columns: {', '.join(sorted(missing))}"
            )

        for row in reader:
            site_code = str(row["site_code"]).strip().upper()
            if not site_code:
                raise ValueError("sites.csv contains a blank site_code")

            sites[site_code] = SiteInfo(
                site_code=site_code,
                site_name=str(row["site_name"]).strip(),
                latitude=float(row["latitude"]),
                longitude=float(row["longitude"]),
                timezone=str(row["timezone"]).strip(),
            )

    if not sites:
        raise ValueError("No valid site rows found in sites.csv")

    return sites


# ------------------------------------------------------------
# WAV + filename parsing
# ------------------------------------------------------------

def get_wav_metadata(wav_path: Path) -> tuple[float, int, int]:
    with wave.open(str(wav_path), "rb") as wf:
        frames = wf.getnframes()
        rate = wf.getframerate()
        channels = wf.getnchannels()

    if rate <= 0:
        raise ValueError(f"Invalid sample rate in {wav_path}")

    duration_s = frames / float(rate)
    return duration_s, rate, channels


def parse_recording_from_filename(
    wav_path: Path,
    raw_audio_root: Path,
    sites: Dict[str, SiteInfo],
) -> RecordingInfo:
    match = FILENAME_RE.match(wav_path.name)
    if not match:
        raise ValueError(
            f"Filename does not match expected pattern SITE_YYYYMMDD_HHMMSS.wav: {wav_path.name}"
        )

    site_code = match.group("site").upper()
    ymd = match.group("ymd")
    hms = match.group("hms")

    if site_code not in sites:
        raise ValueError(
            f"Unknown site code in filename: {site_code} (file: {wav_path.name}). "
            f"Add it to sites.csv before running this script."
        )

    site = sites[site_code]
    tz = ZoneInfo(site.timezone)

    start_dt = datetime.strptime(f"{ymd}_{hms}", "%Y%m%d_%H%M%S").replace(tzinfo=tz)

    duration_s, sample_rate_hz, channels = get_wav_metadata(wav_path)
    end_dt = start_dt + timedelta(seconds=duration_s)

    recording_date = f"{ymd[:4]}-{ymd[4:6]}-{ymd[6:8]}"
    start_datetime = f"{ymd[:4]}-{ymd[4:6]}-{ymd[6:8]} {hms[:2]}:{hms[2:4]}:{hms[4:6]}"

    return RecordingInfo(
        wav_path=wav_path,
        recording_id=wav_path.stem,
        site_code=site_code,
        recording_date=recording_date,
        start_datetime=start_datetime,
        start_dt=start_dt,
        end_dt=end_dt,
        duration_s=duration_s,
        sample_rate_hz=sample_rate_hz,
        channels=channels,
        relative_path=path_relative_to_root(wav_path, raw_audio_root),
    )


# ------------------------------------------------------------
# Astronomical night
# ------------------------------------------------------------

def astronomical_night_interval_for_date(
    site: SiteInfo,
    local_date: date,
) -> Tuple[datetime, datetime]:
    observer = LocationInfo(
        name=site.site_name or site.site_code,
        region="",
        timezone=site.timezone,
        latitude=site.latitude,
        longitude=site.longitude,
    ).observer

    tz = ZoneInfo(site.timezone)

    s_today = sun(observer, date=local_date, tzinfo=tz, dawn_dusk_depression=18)
    s_next = sun(
        observer,
        date=local_date + timedelta(days=1),
        tzinfo=tz,
        dawn_dusk_depression=18,
    )

    dusk_dt = s_today["dusk"]
    dawn_next_dt = s_next["dawn"]

    return dusk_dt, dawn_next_dt


def intersect_intervals(
    a_start: datetime,
    a_end: datetime,
    b_start: datetime,
    b_end: datetime,
) -> Optional[Tuple[datetime, datetime]]:
    start = max(a_start, b_start)
    end = min(a_end, b_end)
    if end > start:
        return start, end
    return None


def get_candidate_night_intervals(
    recording: RecordingInfo,
    site: SiteInfo,
) -> List[Tuple[datetime, datetime]]:
    intervals: List[Tuple[datetime, datetime]] = []
    start_date = recording.start_dt.date() - timedelta(days=1)
    end_date = recording.end_dt.date()

    current = start_date
    while current <= end_date:
        intervals.append(astronomical_night_interval_for_date(site, current))
        current += timedelta(days=1)

    return intervals


def overlaps_as_offsets(
    recording: RecordingInfo,
    site: SiteInfo,
) -> List[Tuple[float, float]]:
    overlaps: List[Tuple[float, float]] = []

    for night_start, night_end in get_candidate_night_intervals(recording, site):
        overlap = intersect_intervals(
            recording.start_dt,
            recording.end_dt,
            night_start,
            night_end,
        )
        if overlap is None:
            continue

        overlap_start_dt, overlap_end_dt = overlap
        start_offset_s = (overlap_start_dt - recording.start_dt).total_seconds()
        end_offset_s = (overlap_end_dt - recording.start_dt).total_seconds()

        start_offset_s = max(0.0, min(start_offset_s, recording.duration_s))
        end_offset_s = max(0.0, min(end_offset_s, recording.duration_s))

        if end_offset_s > start_offset_s:
            overlaps.append((start_offset_s, end_offset_s))

    overlaps.sort()

    merged: List[Tuple[float, float]] = []
    for start_s, end_s in overlaps:
        if not merged:
            merged.append((start_s, end_s))
            continue

        prev_start, prev_end = merged[-1]
        if start_s <= prev_end:
            merged[-1] = (prev_start, max(prev_end, end_s))
        else:
            merged.append((start_s, end_s))

    return merged


# ------------------------------------------------------------
# Label writing
# ------------------------------------------------------------

def write_audacity_label_file(
    label_path: Path,
    intervals: List[Tuple[float, float]],
) -> None:
    label_path.parent.mkdir(parents=True, exist_ok=True)
    with label_path.open("w", encoding="utf-8", newline="\n") as f:
        for start_s, end_s in intervals:
            f.write(f"{start_s:.3f}\t{end_s:.3f}\tASTRO_NIGHT\n")


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

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Create ASTRO_NIGHT Audacity label files for contributor WAV recordings."
    )
    parser.add_argument(
        "--raw-audio-root",
        type=Path,
        default=default_raw_audio_root(),
        help=f"Root folder containing WAV files (default: {default_raw_audio_root()})",
    )
    parser.add_argument(
        "--label-output-root",
        type=Path,
        default=default_label_output_root(),
        help=f"Folder where ASTRO_NIGHT label files will be written (default: {default_label_output_root()})",
    )
    parser.add_argument(
        "--summary-csv",
        type=Path,
        default=default_summary_csv(),
        help=f"Summary CSV path (default: {default_summary_csv()})",
    )
    parser.add_argument(
        "--sites-csv",
        type=Path,
        default=default_sites_csv(),
        help=f"CSV containing site metadata (default: {default_sites_csv()})",
    )
    parser.add_argument(
        "--file-glob",
        default="*.wav",
        help="Glob pattern for audio files (default: *.wav)",
    )
    parser.add_argument(
        "--overwrite-labels",
        action="store_true",
        help="Overwrite existing ASTRO_NIGHT label files",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()

    if not args.raw_audio_root.exists():
        print(f"Raw audio root not found: {args.raw_audio_root}", file=sys.stderr)
        return 1

    try:
        sites = load_sites_from_csv(args.sites_csv)
    except Exception as e:
        print(f"Error loading sites.csv: {e}", file=sys.stderr)
        return 1

    wav_files = sorted(args.raw_audio_root.rglob(args.file_glob))
    if not wav_files:
        print("No WAV files found.", file=sys.stderr)
        return 1

    created_labels = 0
    skipped_labels = 0
    overwritten_labels = 0
    had_errors = False

    summary_rows: List[dict] = []

    for wav_path in wav_files:
        try:
            rec = parse_recording_from_filename(
                wav_path=wav_path,
                raw_audio_root=args.raw_audio_root,
                sites=sites,
            )

            relative_subdir = normalised_relative_subdir(
                wav_path=wav_path,
                raw_audio_root=args.raw_audio_root,
                site_code=rec.site_code,
            )
            label_dir = args.label_output_root / relative_subdir
            label_path = label_dir / f"{rec.recording_id}_night.txt"

            site = sites[rec.site_code]
            intervals = overlaps_as_offsets(rec, site)

            if label_path.exists() and not args.overwrite_labels:
                label_status = "exists"
                skipped_labels += 1
            else:
                if intervals:
                    write_audacity_label_file(label_path, intervals)
                    if label_path.exists() and args.overwrite_labels:
                        label_status = "overwritten"
                        overwritten_labels += 1
                    else:
                        label_status = "created"
                        created_labels += 1
                else:
                    label_status = "no_overlap"
                    skipped_labels += 1

            total_overlap_s = sum(end - start for start, end in intervals) if intervals else 0.0
            interval_text = "; ".join(f"{s:.3f}-{e:.3f}" for s, e in intervals)

            summary_rows.append(
                {
                    "recording_id": rec.recording_id,
                    "wav_path": path_relative_to_root(wav_path, get_contributor_root()),
                    "site_code": rec.site_code,
                    "relative_path": rec.relative_path,
                    "start_datetime_local": rec.start_dt.isoformat(),
                    "end_datetime_local": rec.end_dt.isoformat(),
                    "duration_s": f"{rec.duration_s:.3f}",
                    "sample_rate_hz": rec.sample_rate_hz,
                    "channels": rec.channels,
                    "astro_night_total_overlap_s": f"{total_overlap_s:.3f}",
                    "n_label_intervals": len(intervals),
                    "label_file": path_relative_to_root(label_path, get_contributor_root()) if label_path.exists() else "",
                    "label_status": label_status,
                    "intervals_s": interval_text,
                    "error": "",
                }
            )

            print(f"{rec.recording_id}: label {label_status}")

        except Exception as e:
            had_errors = True
            summary_rows.append(
                {
                    "recording_id": wav_path.stem,
                    "wav_path": path_relative_to_root(wav_path, get_contributor_root()),
                    "site_code": "",
                    "relative_path": "",
                    "start_datetime_local": "",
                    "end_datetime_local": "",
                    "duration_s": "",
                    "sample_rate_hz": "",
                    "channels": "",
                    "astro_night_total_overlap_s": "",
                    "n_label_intervals": "",
                    "label_file": "",
                    "label_status": "error",
                    "intervals_s": "",
                    "error": str(e),
                }
            )
            print(f"ERROR processing {wav_path.name}: {e}", file=sys.stderr)

    args.summary_csv.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "recording_id",
        "wav_path",
        "site_code",
        "relative_path",
        "start_datetime_local",
        "end_datetime_local",
        "duration_s",
        "sample_rate_hz",
        "channels",
        "astro_night_total_overlap_s",
        "n_label_intervals",
        "label_file",
        "label_status",
        "intervals_s",
        "error",
    ]
    with args.summary_csv.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(summary_rows)

    print("\nSummary:")
    print(f"  New ASTRO_NIGHT labels created: {created_labels}")
    print(f"  Labels overwritten: {overwritten_labels}")
    print(f"  Labels skipped: {skipped_labels}")
    print(f"  Summary CSV written to: {args.summary_csv}")

    return 1 if had_errors else 0


if __name__ == "__main__":
    raise SystemExit(main())