"""PLUS batch: обработка папки целиком, с возобновлением и CSV-сводкой.

Для каждой пары «аудио + MIDI» прогоняет пайплайн, пишет результат рядом с
исходником и собирает сводку. Уже обработанные файлы пропускаются, поэтому
прерванный прогон можно продолжить.

CLI:
  python plus_batch.py --input-dir takes --out-dir takes_plus \
      --stages cleanup,quantize,hands,pedal,dynamics,report --audio-ext .wav
"""
from __future__ import annotations

import argparse
import csv
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from plus_core import dump_json, load_midi, save_midi, step  # noqa: E402
from plus_transcribe import run_pipeline  # noqa: E402

__version__ = "1.0"


def find_pairs(folder: str, audio_ext: str) -> list[tuple[str, str | None]]:
    """MIDI и, если есть, одноимённое аудио рядом."""
    pairs = []
    for name in sorted(os.listdir(folder)):
        if not name.lower().endswith((".mid", ".midi")):
            continue
        midi_path = os.path.join(folder, name)
        base = os.path.splitext(midi_path)[0]
        audio = None
        for ext in audio_ext.split(","):
            candidate = base + ext.strip()
            if os.path.isfile(candidate):
                audio = candidate
                break
        pairs.append((midi_path, audio))
    return pairs


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="PLUS batch")
    parser.add_argument("--input-dir", required=True)
    parser.add_argument("--out-dir", required=True)
    parser.add_argument("--stages", default="cleanup,quantize,hands,pedal,dynamics,report")
    parser.add_argument("--audio-ext", default=".wav,.mp3,.flac,.m4a,.ogg")
    parser.add_argument("--force", action="store_true", help="перезаписывать готовое")
    args = parser.parse_args(argv)

    os.makedirs(args.out_dir, exist_ok=True)
    pairs = find_pairs(args.input_dir, args.audio_ext)
    if not pairs:
        raise SystemExit("в " + args.input_dir + " нет MIDI")

    stages = [s.strip() for s in args.stages.split(",") if s.strip()]
    csv_path = os.path.join(args.out_dir, "summary.csv")
    rows: list[dict] = []
    if os.path.isfile(csv_path):
        with open(csv_path, newline="", encoding="utf-8") as handle:
            rows = list(csv.DictReader(handle))
    by_file = {row["file"]: row for row in rows}
    started = time.time()
    for index, (midi_path, audio) in enumerate(pairs, 1):
        base = os.path.splitext(os.path.basename(midi_path))[0]
        out_midi = os.path.join(args.out_dir, base + "_plus.mid")
        out_json = os.path.join(args.out_dir, base + "_report.json")
        if os.path.isfile(out_midi) and os.path.isfile(out_json) and not args.force:
            step(index, len(pairs), "%s: уже готово, пропускаю" % base)
            continue
        step(index, len(pairs), "%s%s" % (base, " + аудио" if audio else " (без аудио)"))
        try:
            piece = load_midi(midi_path)
            options = {"audio": audio}
            report = run_pipeline(piece, stages, options)
            save_midi(piece, out_midi)
            dump_json(report, out_json)
            by_file[base] = {"file": base,
                             "notes_in": report["stages"].get("report", {}).get(
                                 "notes", len(piece.notes)),
                             "notes_out": len(piece.notes), "status": "ok"}
        except Exception as exc:  # один битый файл не должен ронять пачку
            by_file[base] = {"file": base, "notes_in": "", "notes_out": "",
                             "status": "%s: %s" % (type(exc).__name__, exc)}

    rows = sorted(by_file.values(), key=lambda row: row["file"])
    with open(csv_path, "w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=["file", "notes_in", "notes_out", "status"])
        writer.writeheader()
        writer.writerows(rows)
    ok = sum(1 for row in rows if row["status"] == "ok")
    print("готово: %d из %d за %.1f с | сводка: %s"
          % (ok, len(pairs), time.time() - started, csv_path), flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())