# merge lora into base, export fp32 onnx, quantize to int8.
# writes runs/<run-id>/merged/, onnx/model.onnx, onnx/int8/model.onnx.

import argparse
from pathlib import Path

from cleanup.config import load_train_config
from cleanup.export.merge import merge_adapter
from cleanup.export.quantize import quantize_int8
from cleanup.export.to_onnx import export_onnx


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="configs/train.yaml")
    parser.add_argument("--runs-dir", default="runs")
    parser.add_argument("--run-id", required=True)
    parser.add_argument("--skip-int8", action="store_true")
    parser.add_argument("--skip-onnx", action="store_true", help="merge only, no onnx export")
    args = parser.parse_args()

    cfg = load_train_config(args.config)
    run_dir = Path(args.runs_dir) / args.run_id
    adapter_dir = run_dir / "model"
    merged_dir = run_dir / "merged"
    onnx_dir = run_dir / "onnx"
    int8_dir = onnx_dir / "int8"

    merge_adapter(cfg, adapter_dir, merged_dir)
    if args.skip_onnx:
        print("[export] skipping onnx per --skip-onnx")
        return
    fp32_onnx = export_onnx(merged_dir, onnx_dir)
    if not args.skip_int8:
        quantize_int8(fp32_onnx, int8_dir)
    print(f"next: make benchmark RUN_ID={args.run_id}    (LOCAL cpu only)")


if __name__ == "__main__":
    main()