"""Загрузка тестовых частей наборов и сборка манифеста для run.py.

Запуск:
    bench/.venv/Scripts/python.exe bench/fetch.py ruls --limit 300
    bench/.venv/Scripts/python.exe bench/fetch.py golos_test          # весь тест

Берём ТОЛЬКО test-parquet через `https://huggingface.co/api/datasets/<id>/parquet`.
Причина простая: тестовые части — единицы гигабайт, полные наборы — десятки, и
качать train ради замера незачем. Оригинальные архивы (OpenSLR, Яндекс.Диск) в
`sources.py` указаны для того, кто захочет проверить нас по первоисточнику.

Выборка при `--limit` берётся с равномерным шагом по всему тесту, а не первыми N
записями: в этих наборах записи идут группами по диктору и по сессии, первые N дадут
несколько голосов и одну акустику вместо среза набора.

Контрольные суммы скачанных parquet пишутся в `bench/results/checksums-<id>.json` и
публикуются вместе с цифрами: без них замер невозможно повторить на том же материале.
Сами хеши производным материалом не являются, поэтому под ShareAlike не попадают.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from pathlib import Path

import pyarrow.parquet as pq
import requests

sys.path.insert(0, str(Path(__file__).parent))

from sources import BY_ID, Part, Source  # noqa: E402

HF_API = 'https://huggingface.co/api/datasets'
ROOT = Path(__file__).parent
CHUNK = 1 << 20


def parquet_urls(part: Part) -> list[str]:
    """Ссылки на parquet нужного сплита. Отдельный вызов API вместо угадывания путей:
    HF пересобирает конвертированные ветки, и захардкоженный путь однажды отвалится."""
    r = requests.get(f'{HF_API}/{part.hf}/parquet', timeout=60)
    r.raise_for_status()
    tree = r.json()
    config = tree.get(part.config) or next(iter(tree.values()))
    urls = config.get(part.split) or []
    if not urls:
        raise SystemExit(f'{part.hf}: сплита {part.split} нет, есть {sorted(config)}')
    return urls


def download(url: str, dest: Path) -> str:
    """Скачивание с докачкой по факту наличия файла и sha256 на выходе."""
    dest.parent.mkdir(parents=True, exist_ok=True)
    if not dest.exists():
        tmp = dest.with_suffix(dest.suffix + '.part')
        with requests.get(url, stream=True, timeout=600) as r:
            r.raise_for_status()
            total = int(r.headers.get('content-length') or 0)
            got = 0
            with open(tmp, 'wb') as fh:
                for chunk in r.iter_content(CHUNK):
                    fh.write(chunk)
                    got += len(chunk)
                    if total and got % (32 * CHUNK) < CHUNK:
                        print(f'  {dest.name}: {got / 1e6:.0f}/{total / 1e6:.0f} МБ', flush=True)
        tmp.replace(dest)
    h = hashlib.sha256()
    with open(dest, 'rb') as fh:
        for chunk in iter(lambda: fh.read(CHUNK), b''):
            h.update(chunk)
    return h.hexdigest()


def audio_bytes(cell: object) -> tuple[bytes | None, str]:
    """Байты аудио и исходное имя файла из ячейки parquet. Форма ячейки зависит от
    того, как репак сохранял датасет: обычно структура {bytes, path}, иногда просто
    байты. Разбираем оба варианта, вместо того чтобы падать на втором."""
    if isinstance(cell, dict):
        return cell.get('bytes'), str(cell.get('path') or '')
    if isinstance(cell, (bytes, bytearray)):
        return bytes(cell), ''
    return None, ''


def pick(total: int, limit: int) -> list[int]:
    """Индексы выборки равномерным шагом. При limit >= total — все записи."""
    if not limit or limit >= total:
        return list(range(total))
    return [round(i * total / limit) for i in range(limit)]


def fetch_part(src: Source, part: Part, out_dir: Path, limit: int) -> tuple[list[dict], dict]:
    """Один parquet-источник: скачать, разложить аудио, собрать записи манифеста."""
    urls = parquet_urls(part)
    raw_dir = out_dir / 'parquet'
    checks: dict[str, str] = {}
    tables = []
    for i, url in enumerate(urls):
        dest = raw_dir / f'{part.hf.replace("/", "__")}-{part.split}-{i}.parquet'
        print(f'{part.hf}: {url.split("/")[-1]}')
        checks[dest.name] = download(url, dest)
        tables.append(pq.read_table(dest))

    import pyarrow as pa

    table = pa.concat_tables(tables) if len(tables) > 1 else tables[0]
    if part.rows and table.num_rows != part.rows:
        print(f'  ВНИМАНИЕ: записей {table.num_rows}, ожидалось {part.rows} — набор изменился')

    text_col = part.text_field if part.text_field in table.column_names else None
    if text_col is None:
        raise SystemExit(f'{part.hf}: нет колонки {part.text_field}, есть {table.column_names}')

    audio_dir = out_dir / 'audio'
    audio_dir.mkdir(parents=True, exist_ok=True)
    rows: list[dict] = []
    idx = pick(table.num_rows, limit)
    audio = table.column(part.audio_field)
    text = table.column(text_col)
    prefix = part.hf.split('/')[-1]
    for n, i in enumerate(idx):
        data, name = audio_bytes(audio[i].as_py())
        ref = text[i].as_py() or ''
        if not data or not ref.strip():
            continue
        ext = Path(name).suffix or '.wav'
        rid = f'{prefix}-{i}'
        path = audio_dir / f'{rid}{ext}'
        if not path.exists():
            path.write_bytes(data)
        rows.append({'id': rid, 'path': f'audio/{path.name}', 'ref': ref})
        if n and n % 500 == 0:
            print(f'  разложено {n}/{len(idx)}', flush=True)
    return rows, checks


def main() -> int:
    ap = argparse.ArgumentParser(description='Загрузка тестовой части набора')
    ap.add_argument('source', choices=sorted(BY_ID), help='id набора из sources.py')
    ap.add_argument('--limit', type=int, default=0, help='сколько записей взять (0 — все)')
    ap.add_argument('--out', default=None, help='куда положить, по умолчанию bench/data/<id>')
    args = ap.parse_args()

    src = BY_ID[args.source]
    out_dir = Path(args.out) if args.out else ROOT / 'data' / src.id
    print(f'{src.title} — {src.license}, срез {src.slice}')

    manual = [p for p in src.parts if p.manual]
    if manual:
        print('\nАвтоматически этот набор забрать нельзя:\n')
        for p in manual:
            print(p.manual)
        return 3
    if not src.parts:
        print('источник не описан в sources.py')
        return 3

    rows: list[dict] = []
    checks: dict[str, str] = {}
    # при нескольких частях лимит делим между ними, чтобы срез остался пропорциональным
    per_part = (args.limit // len(src.parts)) if args.limit else 0
    for part in src.parts:
        part_rows, part_checks = fetch_part(src, part, out_dir, per_part)
        rows += part_rows
        checks |= part_checks

    manifest = out_dir / 'test.jsonl'
    manifest.write_text(
        '\n'.join(json.dumps(r, ensure_ascii=False) for r in rows) + '\n', encoding='utf-8',
    )
    checks_path = ROOT / 'results' / f'checksums-{src.id}.json'
    checks_path.parent.mkdir(parents=True, exist_ok=True)
    checks_path.write_text(json.dumps({
        'source': src.id, 'title': src.title, 'license': src.license,
        'attribution': src.attribution, 'records_taken': len(rows),
        'sha256': checks,
    }, ensure_ascii=False, indent=2), encoding='utf-8')

    print(f'\nманифест: {manifest} ({len(rows)} записей)')
    print(f'контрольные суммы: {checks_path}')
    print(f'атрибуция при публикации: {src.attribution}')
    return 0


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