#!/usr/bin/env bash
set -euo pipefail

if [[ $# -lt 3 ]]; then
  echo "Usage: $0 DATASET_DIR OUTPUT_DIR PIPER_REPO [EPOCHS]" >&2
  exit 2
fi

dataset_dir="$(realpath "$1")"
output_dir="$(mkdir -p "$2" && realpath "$2")"
piper_repo="$(realpath "$3")"
epochs="${4:-2000}"

metadata="$dataset_dir/metadata.csv"
audio_dir="$dataset_dir/wav"
if [[ ! -f "$metadata" || ! -d "$audio_dir" ]]; then
  echo "Dataset must contain metadata.csv and wav/" >&2
  exit 2
fi
if [[ ! -f "$piper_repo/src/piper/train/__main__.py" ]]; then
  echo "Piper training sources not found at $piper_repo" >&2
  exit 2
fi

python3 -m piper.train fit \
  --data.voice_name aletheia_ru \
  --data.csv_path "$metadata" \
  --data.audio_dir "$audio_dir" \
  --data.espeak_voice ru \
  --data.cache_dir "$output_dir/cache" \
  --data.config_path "$output_dir/aletheia_ru.onnx.json" \
  --data.batch_size 16 \
  --data.num_workers 4 \
  --model.sample_rate 22050 \
  --trainer.accelerator gpu \
  --trainer.devices 1 \
  --trainer.precision 16-mixed \
  --trainer.max_epochs "$epochs" \
  --trainer.default_root_dir "$output_dir/checkpoints"

checkpoint="$(find "$output_dir/checkpoints" -type f -name '*.ckpt' -printf '%T@ %p\n' | sort -nr | head -n 1 | cut -d' ' -f2-)"
if [[ -z "$checkpoint" ]]; then
  echo "Training finished without a checkpoint" >&2
  exit 1
fi

python3 -m piper.train.export_onnx \
  --checkpoint "$checkpoint" \
  --output-file "$output_dir/aletheia_ru.onnx"

python3 - "$output_dir" <<'PY'
import hashlib
import json
import pathlib
import sys

root = pathlib.Path(sys.argv[1])
artifacts = {}
for name in ("aletheia_ru.onnx", "aletheia_ru.onnx.json"):
    path = root / name
    digest = hashlib.sha256(path.read_bytes()).hexdigest()
    artifacts[name] = {"bytes": path.stat().st_size, "sha256": digest}
(root / "artifacts.json").write_text(json.dumps(artifacts, indent=2), encoding="utf-8")
print(json.dumps(artifacts, indent=2))
PY