"""Configure a dataset for Matcha-only or joint fine-tuning."""

import argparse
import json
import shutil
from pathlib import Path
from typing import Optional


def configure_matcha_ft(
    config_path: Path,
    matcha_only: bool = True,
    learning_rate: Optional[float] = None,
    epochs: Optional[int] = None,
    save_every_steps: Optional[int] = None,
) -> dict:
    if not config_path.is_file():
        raise FileNotFoundError(f"Config file not found: {config_path}")
    with config_path.open(encoding="utf-8") as stream:
        config = json.load(stream)

    train = config.setdefault("train", {})
    model = config.setdefault("model", {})
    train["matcha_only"] = matcha_only
    model["use_matcha"] = True
    model["matcha_use_diff_attention"] = True
    if learning_rate is not None:
        if learning_rate <= 0:
            raise ValueError("Learning rate must be positive")
        train["learning_rate"] = learning_rate
    if epochs is not None:
        if epochs < 1:
            raise ValueError("Epochs must be at least 1")
        train["epochs"] = epochs
    if save_every_steps is not None:
        if save_every_steps < 1:
            raise ValueError("Save interval must be at least 1")
        train["eval_interval"] = save_every_steps

    backup_path = config_path.with_suffix(config_path.suffix + ".bak")
    if not backup_path.exists():
        shutil.copy2(config_path, backup_path)
    with config_path.open("w", encoding="utf-8") as stream:
        json.dump(config, stream, indent=2, ensure_ascii=False)
        stream.write("\n")
    return config


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Configure Matcha Flow fine-tuning for an existing dataset"
    )
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument(
        "--joint",
        action="store_true",
        help="Jointly fine-tune the full model instead of Matcha-only training",
    )
    parser.add_argument("--learning-rate", type=float)
    parser.add_argument("--epochs", type=int)
    parser.add_argument("--save-every-steps", type=int)
    args = parser.parse_args()

    config = configure_matcha_ft(
        config_path=args.config,
        matcha_only=not args.joint,
        learning_rate=args.learning_rate,
        epochs=args.epochs,
        save_every_steps=args.save_every_steps,
    )
    mode = "joint fine-tuning" if args.joint else "Matcha-only fine-tuning"
    print(f"Configured {args.config} for {mode}")
    print(f"learning_rate={config['train']['learning_rate']}")
    print(f"epochs={config['train']['epochs']}")
    print(f"save_every_steps={config['train']['eval_interval']}")


if __name__ == "__main__":
    main()