import argparse
import json
import subprocess
import sys
from pathlib import Path


def run(cmd):
    print("Running:", " ".join(cmd))
    result = subprocess.run(cmd, check=False)
    if result.returncode != 0:
        raise SystemExit(result.returncode)


def flag_present(flag_name):
    return flag_name in sys.argv


def usable_repo_id(repo_id):
    if not repo_id:
        return ""
    placeholders = ("your-username/", "your-user/", "username/")
    return "" if repo_id.startswith(placeholders) else repo_id


def apply_config_defaults(args):
    config_path = Path("training_config.json")
    if not config_path.exists():
        return args

    with config_path.open("r", encoding="utf-8") as f:
        cfg = json.load(f)

    if not flag_present("--model-name"):
        args.model_name = cfg.get("model_name", args.model_name)
    if not flag_present("--dataset-size"):
        args.dataset_size = cfg.get("dataset_size", args.dataset_size)
    if not flag_present("--train-file"):
        args.train_file = cfg.get("train_file", args.train_file)
    if not flag_present("--output-dir"):
        args.output_dir = cfg.get("output_dir", args.output_dir)
    if not flag_present("--hf-repo"):
        args.hf_repo = usable_repo_id(cfg.get("hf_repo_id", args.hf_repo))
    if not flag_present("--epochs"):
        args.epochs = cfg.get("epochs", args.epochs)
    if not flag_present("--batch-size"):
        args.batch_size = cfg.get("batch_size", args.batch_size)
    if not flag_present("--learning-rate"):
        args.learning_rate = cfg.get("learning_rate", args.learning_rate)
    if not flag_present("--max-length"):
        args.max_length = cfg.get("max_length", args.max_length)
    if not flag_present("--max-train-samples"):
        args.max_train_samples = cfg.get("max_train_samples", args.max_train_samples)
    if not flag_present("--use-4bit"):
        args.use_4bit = cfg.get("use_4bit", args.use_4bit)
    return args


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-size", type=int, default=8000)
    parser.add_argument("--train-file", type=str, default="train.json")
    parser.add_argument("--output-dir", type=str, default="model")
    parser.add_argument("--model-name", type=str, default="Qwen/Qwen2.5-Coder-0.5B-Instruct")
    parser.add_argument("--epochs", type=float, default=1)
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--learning-rate", type=float, default=2e-4)
    parser.add_argument("--max-length", type=int, default=512)
    parser.add_argument("--max-train-samples", type=int, default=0)
    parser.add_argument("--use-4bit", action="store_true")
    parser.add_argument("--hf-repo", type=str, default="")
    parser.add_argument("--skip-generate", action="store_true")
    parser.add_argument("--skip-train", action="store_true")
    parser.add_argument("--skip-upload", action="store_true")
    parser.add_argument("--skip-benchmark", action="store_true")
    args = parser.parse_args()
    args = apply_config_defaults(args)

    if not (5000 <= args.dataset_size <= 10000):
        raise ValueError("dataset-size must be between 5000 and 10000")

    if not args.skip_generate:
        run([sys.executable, "generate_dataset.py", "--size", str(args.dataset_size), "--out", args.train_file])

    if not args.skip_train:
        train_cmd = [
            sys.executable,
            "finetune_coding_llm_colab.py",
            "--dataset-size",
            str(args.dataset_size),
            "--train-file",
            args.train_file,
            "--output-dir",
            args.output_dir,
            "--model-name",
            args.model_name,
            "--epochs",
            str(args.epochs),
            "--batch-size",
            str(args.batch_size),
            "--learning-rate",
            str(args.learning_rate),
            "--max-length",
            str(args.max_length),
            "--max-train-samples",
            str(args.max_train_samples),
            "--skip-dataset-gen",
        ]
        if args.use_4bit:
            train_cmd.append("--use-4bit")
        run(train_cmd)
    else:
        print("Skipping training stage (--skip-train).")

    if not args.skip_upload:
        if not args.hf_repo:
            raise ValueError("Pass --hf-repo when upload is enabled, or use --skip-upload")
        run([sys.executable, "upload_to_hf.py", "--model-dir", args.output_dir, "--repo-id", args.hf_repo])

    if not args.skip_benchmark:
        run([sys.executable, "generate_benchmark.py"])
    else:
        print("Skipping benchmark chart stage (--skip-benchmark).")

    print("Pipeline completed.")


if __name__ == "__main__":
    main()