"""SEC-cyBERT training CLI.

Usage:
    uv run main.py dapt --config configs/dapt/modernbert.yaml
    uv run main.py dapt --config configs/dapt/modernbert.yaml \\
        --model-path ../checkpoints/dapt/modernbert-large/final \\
        --data-path ../data/paragraphs/paragraphs-clean.jsonl \\
        --output-dir ../checkpoints/tapt/modernbert-large \\
        --stage tapt
"""

import argparse
import sys


def cmd_dapt(args: argparse.Namespace) -> None:
    from src.common.config import DAPTConfig
    from src.dapt.train import train

    config = DAPTConfig.from_yaml(args.config)
    config.apply_overrides(
        model_path=args.model_path,
        data_path=args.data_path,
        output_dir=args.output_dir,
        stage=args.stage,
    )
    train(config)


def cmd_finetune(args: argparse.Namespace) -> None:
    from src.common.config import FinetuneConfig
    from src.finetune.train import train

    config = FinetuneConfig.from_yaml(args.config)
    config.apply_overrides(
        model_path=args.model_path,
        output_dir=args.output_dir,
        loss_type=args.loss_type,
        class_weighting=args.class_weighting,
        epochs=args.epochs,
        batch_size=args.batch_size,
        learning_rate=args.learning_rate,
    )
    train(config)


def cmd_ablate(args: argparse.Namespace) -> None:
    from src.common.config import FinetuneConfig
    from src.finetune.train import ablate

    config = FinetuneConfig.from_yaml(args.config)
    if args.output_dir:
        config.training.output_dir = args.output_dir
    if args.epochs:
        config.training.num_train_epochs = args.epochs
    ablate(config)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="SEC-cyBERT training pipeline",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # ── dapt / tapt ──
    dapt = sub.add_parser(
        "dapt",
        help="Run DAPT or TAPT pre-training (masked language modeling)",
    )
    dapt.add_argument("--config", required=True, help="Path to YAML config file")
    dapt.add_argument("--model-path", help="Override model name or checkpoint path")
    dapt.add_argument("--data-path", help="Override corpus path (file or directory)")
    dapt.add_argument("--output-dir", help="Override output directory")
    dapt.add_argument("--stage", choices=["dapt", "tapt"], help="Override stage label")
    dapt.set_defaults(func=cmd_dapt)

    # ── finetune ──
    ft = sub.add_parser("finetune", help="Fine-tune classifier (dual-head)")
    ft.add_argument("--config", required=True, help="Path to YAML config file")
    ft.add_argument("--model-path", help="Override model checkpoint path")
    ft.add_argument("--output-dir", help="Override output directory")
    ft.add_argument("--loss-type", choices=["ce", "focal"], help="Override loss type")
    ft.add_argument("--class-weighting", type=lambda x: x.lower() == "true", help="Override class weighting (true/false)")
    ft.add_argument("--epochs", type=int, help="Override number of epochs")
    ft.add_argument("--batch-size", type=int, help="Override batch size")
    ft.add_argument("--learning-rate", type=float, help="Override learning rate")
    ft.set_defaults(func=cmd_finetune)

    # ── ablate ──
    ab = sub.add_parser("ablate", help="Run full ablation grid (3 ckpts x 2 weighting x 2 loss)")
    ab.add_argument("--config", required=True, help="Path to YAML config file")
    ab.add_argument("--output-dir", help="Override base output directory")
    ab.add_argument("--epochs", type=int, help="Override epochs per ablation run (default: config value)")
    ab.set_defaults(func=cmd_ablate)

    # ── eval (placeholder) ──
    ev = sub.add_parser("eval", help="Evaluate a trained model")
    ev.add_argument("--config", required=True, help="Path to YAML config file")
    ev.set_defaults(func=lambda args: print("Evaluation not yet implemented."))

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()