# _xenith - Zsh completion function for Xenith
_xenith() {
    local curcontext="$curcontext" state line
    typeset -A opt_args
    local commands=(
        'new:создать новую модель из корпуса'
        'train:обучить существующую модель'
        'gen:сгенерировать текст'
        'info:показать конфигурацию модели'
        'gradcheck:проверить градиенты численно'
        'bench:замерить скорость форварда и шага обучения'
        'update:автоматическое обновление с гита'
        'tui:графическо-консольный режим (тест)'
    )
    _arguments -C \
        '1: :->command' \
        '*:: :->args' && return 0
    case $state in
        command)
            _describe -t commands 'команды Xenith' commands
            ;;
        args)
            case $line[1] in
                new)
                    _arguments \
                        '--corpus[текст для построения словаря]:файл:_files' \
                        '--out[файл модели]:файл:_files -g "*.xnh"' \
                        '--vocab[размер словаря]' \
                        '--embd[размер эмбеддинга]' \
                        '--layers[число слоёв]' \
                        '--heads[число голов внимания]' \
                        '--ctx[максимальный контекст]' \
                        '--ffn[размер FFN]' \
                        '--rope[база RoPE]' \
                        '--seed[зерно инициализации]' \
                        '--untied[использовать отдельную матрицу выхода]' \
                        '--view-vocab[вывести словарь нейросети]'
                    ;;
                train)
                    _arguments \
                        '--corpus[обучающий текст]:файл:_files' \
                        '--out[куда сохранить]:файл:_files -g "*.xnh"' \
                        '--steps[количество шагов обучения]' \
                        '--batch[размер батча]' \
                        '--block[длина окна контекста]' \
                        '--lr[скорость обучения]' \
                        '--wd[weight decay]' \
                        '--clip[клиппинг нормы градиента]' \
                        '--warmup[шаги прогрева]' \
                        '--seed[зерно генератора]' \
                        '--ckpt-every[сохранять чекпоинт каждые N шагов]' \
                        '--val-every[валидация каждые N шагов]' \
                        '--val-tokens[размер валидационной выборки]' \
                        '--threads[количество потоков (0 = все ядра)]' \
                        '--resume[продолжить обучение с последнего чекпоинта]' \
                        '--fast[экспериментальный флаг ускорения forward pass]'
                    ;;
                gen)
                    _arguments \
                        '--prompt[стартовый текст]' \
                        '--n[количество генерируемых токенов]' \
                        '--temp[температура выборки (0 = жадный)]' \
                        '--top-k[top-k sampling]' \
                        '--top-p[top-p nucleus sampling]' \
                        '--seed[зерно генератора]' \
                        '--no-stream[не печатать токены по мере генерации]' \
                        '--batch[несколько промптов через ;]' \
                        '--show-tokens[показывать ID токенов]'
                    ;;
                info|gradcheck|bench|update|tui)
                    _arguments \
                        '--fast[экспериментальный флаг ускорения]'
                    ;;
            esac
            ;;
    esac
}
if [[ "$funcstack[1]" == "_xenith" ]]; then
    _xenith "$@"
fi