from __future__ import annotations

"""
summarise_night_annotations_contributor.py

Summarise annotations for a selected astronomical night.

Uses the same label parsing rules as:
extract_annotation_clips_contributor.py
"""

import argparse
import csv
import re
import sys
from collections import Counter
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 root() -> Path:
    return Path(__file__).resolve().parent.parent


RAW = root() / "1_raw_audio"
LABELS = root() / "2_labels"
OUT = root() / "4_summaries"


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


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

import re

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
    start_dt: datetime
    end_dt: datetime


# ------------------------------------------------------------
# Label parsing (MATCHES CLIP SCRIPT)
# ------------------------------------------------------------

def parse_annotation_label(raw_label: str) -> tuple[str, str]:
    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:
    text = text.strip()
    text = re.sub(r'[<>:"/\\|?*]', "_", text)
    text = text.rstrip(" .")
    return text or "UNLABELLED"


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

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

    return sanitise_filename_component(name)


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


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

def load_sites(csv_path: Path) -> Dict[str, SiteInfo]:
    sites = {}
    with csv_path.open("r", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        for row in reader:
            code = row["site_code"].strip().upper()
            sites[code] = SiteInfo(
                site_code=code,
                site_name=row["site_name"],
                latitude=float(row["latitude"]),
                longitude=float(row["longitude"]),
                timezone=row["timezone"],
            )
    return sites


# ------------------------------------------------------------
# WAV parsing
# ------------------------------------------------------------

def parse_recording(wav: Path, sites: Dict[str, SiteInfo]) -> RecordingInfo:
    m = FILENAME_RE.match(wav.name)
    if not m:
        raise ValueError(f"Bad filename: {wav.name}")

    site = m.group("site").upper()
    ymd = m.group("ymd")
    hms = m.group("hms")

    tz = ZoneInfo(sites[site].timezone)

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

    # rough duration estimate from filename only not needed → we’ll ignore
    end = start + timedelta(hours=12)

    return RecordingInfo(wav, wav.stem, site, start, end)


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

def astro_interval(site: SiteInfo, d: date):
    obs = LocationInfo(
        name=site.site_name,
        region="",
        timezone=site.timezone,
        latitude=site.latitude,
        longitude=site.longitude,
    ).observer

    tz = ZoneInfo(site.timezone)

    s1 = sun(obs, date=d, tzinfo=tz, dawn_dusk_depression=18)
    s2 = sun(obs, date=d + timedelta(days=1), tzinfo=tz, dawn_dusk_depression=18)

    return s1["dusk"], s2["dawn"]


def overlaps(rec: RecordingInfo, site: SiteInfo, d: date) -> bool:
    ns, ne = astro_interval(site, d)
    return max(rec.start_dt, ns) < min(rec.end_dt, ne)


# ------------------------------------------------------------
# Labels
# ------------------------------------------------------------

def find_label(recording_id: str) -> Optional[Path]:
    for p in LABELS.rglob(f"{recording_id}_*.txt"):
        return p
    return None


def read_labels(path: Path):
    rows = []
    with path.open() as f:
        for line in f:
            parts = line.strip().split("\t")
            if len(parts) >= 3:
                rows.append(parts[2])
    return rows


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

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--night", help="YYYY-MM-DD")
    parser.add_argument("--sites", type=Path, default=default_sites_csv())
    args = parser.parse_args()

    night = args.night or input("Night (YYYY-MM-DD): ").strip()
    night_date = datetime.strptime(night, "%Y-%m-%d").date()

    sites = load_sites(args.sites)

    counts = Counter()

    for wav in RAW.rglob("*.wav"):
        try:
            rec = parse_recording(wav, sites)
            site = sites[rec.site_code]

            if not overlaps(rec, site, night_date):
                continue

            label_file = find_label(rec.recording_id)
            if not label_file:
                continue

            labels = read_labels(label_file)

            for lab in labels:
                if is_astro(lab):
                    continue

                confidence, _ = parse_annotation_label(lab)

                # Skip confirmed non-bird sounds
                if confidence == "zero":
                    continue

                name = call_type_name(lab)
                counts[(rec.site_code, name)] += 1

        except Exception as e:
            print(f"Error: {wav.name}: {e}")

    OUT.mkdir(exist_ok=True)

    out_csv = OUT / f"summary_{night}.csv"

    with out_csv.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["date", "site", "call_type", "n"])

        for (site, call), n in sorted(counts.items()):
            writer.writerow([night, site, call, n])

    print("\nSummary\n" + "-" * 40)

    for (site, call), n in sorted(counts.items()):
        print(f"{night} | {site} | {call} | {n}")

    print(f"\nSaved to {out_csv}")


if __name__ == "__main__":
    main()