"""Upload audio files to the HF dataset repo.

Run this script LOCALLY on your Windows machine:

    pip install huggingface_hub pandas
    python upload_audio.py

It reads Barhatnaya.csv, finds matching MP3 files, and uploads them
to Niko-NN/gold-store-dialogs/audio/ on the Hub.
"""

from __future__ import annotations

import csv
import sys
from pathlib import Path

from huggingface_hub import HfApi, login

# ── Configuration ─────────────────────────────────────────────────────────
DATASET_REPO = "Niko-NN/gold-store-dialogs"
AUDIO_DIR = Path(r"C:\Users\Nick\Desktop\2")
CSV_PATH = Path(r"C:\IDEAI\test\razmetka\Barhatnaya.csv")
AUDIO_EXTENSIONS = [".mp3", ".wav", ".ogg", ".flac", ".mp4"]
# ──────────────────────────────────────────────────────────────────────────


def _has_dialog(row: dict) -> bool:
    """Return True if the row has a dialog (start/end is not just '-')."""
    start = (row.get("start") or "").strip()
    end = (row.get("end") or "").strip()
    for val in (start, end):
        if val and val not in ("-", ""):
            return True
    return False


def read_file_ids(csv_path: Path, dialogs_only: bool = True) -> list[str]:
    """Read video/file IDs from the CSV.

    If dialogs_only=True, only return files that have a dialog
    (start or end is not '-').
    """
    all_ids: list[str] = []
    dialog_ids: list[str] = []
    with csv_path.open("r", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        for row in reader:
            file_id = (row.get("video") or "").strip()
            if not file_id:
                continue
            all_ids.append(file_id)
            if _has_dialog(row):
                dialog_ids.append(file_id)

    print(f"CSV: всего {len(all_ids)} файлов, с диалогами: {len(dialog_ids)}")
    return dialog_ids if dialogs_only else all_ids


def find_audio_file(audio_dir: Path, file_id: str) -> Path | None:
    """Find the audio file matching file_id in audio_dir."""
    for ext in AUDIO_EXTENSIONS:
        candidate = audio_dir / f"{file_id}{ext}"
        if candidate.exists():
            return candidate
    for f in audio_dir.iterdir():
        if f.stem == file_id and f.suffix.lower() in AUDIO_EXTENSIONS:
            return f
    return None


def main():
    if not AUDIO_DIR.exists():
        print(f"ОШИБКА: папка с аудио не найдена: {AUDIO_DIR}")
        sys.exit(1)
    if not CSV_PATH.exists():
        print(f"ОШИБКА: CSV не найден: {CSV_PATH}")
        sys.exit(1)

    login()

    api = HfApi()
    file_ids = read_file_ids(CSV_PATH)

    found: list[tuple[str, Path]] = []
    missing: list[str] = []

    for fid in file_ids:
        path = find_audio_file(AUDIO_DIR, fid)
        if path:
            found.append((fid, path))
        else:
            missing.append(fid)

    print(f"\nНайдено аудио: {len(found)}")
    if missing:
        print(f"Не найдено: {len(missing)}")
        for m in missing[:10]:
            print(f"  - {m}")
        if len(missing) > 10:
            print(f"  ... и ещё {len(missing) - 10}")

    if not found:
        print("Нет файлов для загрузки.")
        sys.exit(0)

    # Check what's already uploaded
    try:
        existing = set(
            f.replace("audio/", "")
            for f in api.list_repo_files(DATASET_REPO, repo_type="dataset")
            if f.startswith("audio/")
        )
    except Exception:
        existing = set()

    to_upload = [(fid, p) for fid, p in found if p.name not in existing]
    skipped = len(found) - len(to_upload)
    if skipped:
        print(f"Уже на Hub: {skipped}, осталось загрузить: {len(to_upload)}")

    if not to_upload:
        print("Все файлы уже загружены.")
        sys.exit(0)

    print(f"\nЗагрузка {len(to_upload)} файлов в {DATASET_REPO}/audio/ ...")

    for i, (fid, path) in enumerate(to_upload, 1):
        remote_path = f"audio/{path.name}"
        size_mb = path.stat().st_size / (1024 * 1024)
        print(f"  [{i}/{len(to_upload)}] {path.name} ({size_mb:.1f} MB) -> {remote_path}")
        api.upload_file(
            path_or_fileobj=str(path),
            path_in_repo=remote_path,
            repo_id=DATASET_REPO,
            repo_type="dataset",
        )

    print(f"\nГотово! Загружено {len(found)} файлов.")


if __name__ == "__main__":
    main()