#!/usr/bin/env python3
import argparse
import os
from rkllm.api import RKLLM

def parse_args():
    parser = argparse.ArgumentParser(description='Компиляция LLM для Rockchip NPU')
    
    # Обязательные параметры
    parser.add_argument('--model', type=str, required=True,
                        help='Путь к модели (Hugging Face или GGUF)')
    parser.add_argument('--output', type=str, required=True,
                        help='Путь для сохранения .rkllm файла')
    
    # Опциональные параметры с умными значениями по умолчанию
    parser.add_argument('--format', type=str, choices=['huggingface', 'gguf'], default='huggingface',
                        help='Формат входной модели (по умолчанию: huggingface)')
    parser.add_argument('--enable-q', type=str, choices=['True', 'False'], default='True',
                        help='Включить квантование (по умолчаниюЖ  True)')                    
    parser.add_argument('--quant', type=str, default='w8a8_g256',
                        choices=['w4a16', 'w4a16_g32', 'w4a16_g64', 'w4a16_g128', 
                                 'w8a8', 'w8a8_g128', 'w8a8_g256', 'w8a8_g512'],
                        help='Тип квантования (по умолчанию: w8a8_g256)')
    parser.add_argument('--platform', type=str, default='rk3588',
                        choices=['rk3588', 'rk3576', 'rk3562', 'rv1126b'],
                        help='Целевая платформа (по умолчанию: rk3588)')
    parser.add_argument('--npu-cores', type=int, default=3, choices=[1, 2, 3],
                        help='Количество ядер NPU (по умолчанию: 3)')
    parser.add_argument('--max-context', type=int, default=8192,
                        help='Максимальная длина контекста (по умолчанию: 8192)')
    
    return parser.parse_args()

def main():
    args = parse_args()
    os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True)
    
    print(f"🔧 Компиляция модели:")
    print(f"   Модель: {args.model}")
    print(f"   Формат: {args.format}")
    print(f"   Включено квантование: {args.enable_q}")
    print(f"   Квантование: {args.quant}")
    print(f"   Платформа: {args.platform}")
    print(f"   Ядра NPU: {args.npu_cores}")
    print(f"   Длина контекста: {args.max_context}")
    
    llm = RKLLM()
    
    # Загрузка в зависимости от формата
    if args.format == 'huggingface':
        ret = llm.load_huggingface(
            model=args.model,
            device='cpu',
            trust_remote_code=True,
            quantized_model=False
        )
    else:  # gguf
        ret = llm.load_gguf(model=args.model)
    
    if ret != 0:
        print(f"❌ Ошибка загрузки модели (код {ret})")
        exit(ret)
    
    # Определение алгоритма квантования
    algo = 'grq' if 'w4a16' in args.quant else 'normal'
    
    ret = llm.build(
        do_quantization=args.enable_q,
        optimization_level=1,
        quantized_dtype=args.quant,
        quantized_algorithm=algo,
        target_platform=args.platform,
        num_npu_core=args.npu_cores,
        max_context=args.max_context
    )
    
    if ret != 0:
        print(f"❌ Ошибка компиляции (код {ret})")
        exit(ret)
    
    ret = llm.export_rkllm(args.output)
    if ret != 0:
        print(f"❌ Ошибка экспорта (код {ret})")
        exit(ret)
    
    size_gb = os.path.getsize(args.output) / (1024**3)
    print(f"\n✅ Успешно! Сохранено: {args.output}")
    print(f"📦 Размер: {size_gb:.2f} GB")

if __name__ == '__main__':
    main()