import argparse

from .generator import generate


def _parse_int_list(value):
    if not value:
        return None
    items = [item.strip() for item in value.split(",") if item.strip()]
    return [int(item) for item in items]


def _parse_str_list(value):
    if not value:
        return None
    return [item.strip() for item in value.split(",") if item.strip()]


def build_parser():
    parser = argparse.ArgumentParser(
        prog="pstorch",
        description="Generate a PyTorch project template.",
    )
    parser.add_argument("name", help="Model class name for the project")
    parser.add_argument("--arch", default="mlp", help="Architecture (mlp)")
    parser.add_argument("--input-dim", type=int, default=784)
    parser.add_argument("--output-dim", type=int, default=10)
    parser.add_argument(
        "--hidden",
        help="Comma-separated hidden sizes (e.g. 256,128)",
    )
    parser.add_argument("--activation", default="relu")
    parser.add_argument(
        "--metrics",
        help="Comma-separated metrics (accuracy,loss,f1,precision,recall)",
    )
    parser.add_argument("--prod", action="store_true", help="Use MLflow logging")
    parser.add_argument("--lr", type=float, default=0.001)
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--batch-size", type=int, default=32)
    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()

    generate(
        name=args.name,
        arch=args.arch,
        input_dim=args.input_dim,
        output_dim=args.output_dim,
        hidden=_parse_int_list(args.hidden),
        activation=args.activation,
        metrics=_parse_str_list(args.metrics),
        prod=args.prod,
        lr=args.lr,
        epochs=args.epochs,
        batch_size=args.batch_size,
    )


if __name__ == "__main__":
    main()
