/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * ternarize-split: Extract FP16 model into per-layer workdir with initial ternary.
 *
 * Usage:
 *   ternarize-split --model model.safetensors --config config.json --workdir ./work
 */

#include "../ternarize/ternarize.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static void usage(const char *prog) {
    fprintf(stderr, "Usage: %s --model <safetensors> --config <config.json> --workdir <dir>\n", prog);
    fprintf(stderr, "\nOptions:\n");
    fprintf(stderr, "  --model    Path to model .safetensors file\n");
    fprintf(stderr, "  --config   Path to HuggingFace config.json\n");
    fprintf(stderr, "  --workdir  Output work directory\n");
    fprintf(stderr, "  --group    Group size for per-group scales (default: 128)\n");
}

int main(int argc, char **argv) {
    ft_config cfg = ft_config_defaults();
    const char *model_path = NULL;
    const char *config_path = NULL;
    const char *workdir = NULL;

    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--model") == 0 && i + 1 < argc) {
            model_path = argv[++i];
        } else if (strcmp(argv[i], "--config") == 0 && i + 1 < argc) {
            config_path = argv[++i];
        } else if (strcmp(argv[i], "--workdir") == 0 && i + 1 < argc) {
            workdir = argv[++i];
        } else if (strcmp(argv[i], "--group") == 0 && i + 1 < argc) {
            cfg.group_size = atoi(argv[++i]);
        } else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
            usage(argv[0]);
            return 0;
        } else {
            fprintf(stderr, "Unknown option: %s\n", argv[i]);
            usage(argv[0]);
            return 1;
        }
    }

    if (!model_path || !config_path || !workdir) {
        usage(argv[0]);
        return 1;
    }

    /* Parse model config */
    if (!ft_config_from_json(&cfg, config_path)) {
        fprintf(stderr, "ERROR: failed to parse %s\n", config_path);
        return 1;
    }

    snprintf(cfg.model_path, sizeof(cfg.model_path), "%s", model_path);
    snprintf(cfg.workdir, sizeof(cfg.workdir), "%s", workdir);

    fprintf(stderr, "ternarize-split\n");
    fprintf(stderr, "  Model: %s\n", model_path);
    fprintf(stderr, "  Config: %s (%s)\n", config_path, ft_arch_name(cfg.arch));
    fprintf(stderr, "  Workdir: %s\n", workdir);
    fprintf(stderr, "  Layers: %d, Hidden: %d, FFN: %d\n",
            cfg.n_layers, cfg.n_embd, cfg.ffn_dim);
    fprintf(stderr, "  Group size: %d\n", cfg.group_size);

    if (!ft_split(&cfg)) {
        fprintf(stderr, "ERROR: split failed\n");
        return 1;
    }

    return 0;
}