@echo off
chcp 936 >nul
setlocal enabledelayedexpansion

:: 设置路径
set "llama_path=E:\ai\llama"
set "model_path=E:\ai\llama\models"
set "cli_exe=%llama_path%\llama-cli.exe"
set "server_exe=%llama_path%\llama-server.exe"
set "log_dir=%llama_path%\logs"
set "config_file=%llama_path%\model_configs.txt"

:: 创建必要目录
if not exist "%log_dir%" mkdir "%log_dir%"

:: 全局变量
set "selected_model="
set "full_model_path="

:main_menu
cls
echo ========================================
echo         LLaMA 模型管理菜单
echo ========================================
echo.
echo [1] 选择并运行模型
echo [2] 查看模型详细信息
echo [3] 批量测试模式
echo [4] 模型管理
echo [5] 参数配置模板
echo [6] 模型对比模式
echo [7] 查看对话日志
echo [8] 性能测试
echo [9] 加载保存的配置
echo [10] 打开模型文件夹
echo [11] 下载模型
echo [0] 退出
echo.
echo ========================================
set /p main_choice=请选择 (0-10):

if "%main_choice%"=="0" exit /b
if "%main_choice%"=="1" goto select_model
if "%main_choice%"=="2" goto model_info_menu
if "%main_choice%"=="3" goto batch_test
if "%main_choice%"=="4" goto model_management
if "%main_choice%"=="5" goto preset_config
if "%main_choice%"=="6" goto compare_models
if "%main_choice%"=="7" goto view_logs
if "%main_choice%"=="8" goto performance_test
if "%main_choice%"=="9" goto load_config
if "%main_choice%"=="10" goto open_folder
if "%main_choice%"=="11" goto download_model
echo 无效选项！
timeout /t 2 >nul
goto main_menu

:: ==========================================
:: 选择并运行模型
:: ==========================================
:select_model
cls
call :list_models
if %count% equ 0 (
    echo 没有找到任何模型文件！
    pause
    goto main_menu
)

echo ----------------------------------------
echo 0. 返回主菜单
echo.
set /p model_choice=请选择模型 (0-%count%)：

if "%model_choice%"=="0" goto main_menu

:: 验证输入
call :validate_input "%model_choice%" %count%
if errorlevel 1 goto select_model

:: 获取选中的模型
set "selected_model=!model%model_choice%!"
set "full_model_path=%model_path%\%selected_model%"

:launch_mode
cls
echo ========================================
echo 已选择模型: %selected_model%
echo ========================================
echo.
echo 请选择启动方式:
echo.
echo [1] 交互式对话模式 (llama-cli.exe)
echo [2] Web服务器模式 (llama-server.exe)
echo [3] 使用自定义参数运行
echo [4] 返回模型选择
echo [5] 返回主菜单
echo [0] 退出
echo.
echo ========================================
set /p mode_choice=请选择 (0-5):

if "%mode_choice%"=="0" exit /b
if "%mode_choice%"=="4" goto select_model
if "%mode_choice%"=="5" goto main_menu

:: 交互式对话
if "%mode_choice%"=="1" (
    call :run_interactive
    goto launch_mode
)

:: Web服务器
if "%mode_choice%"=="2" (
    call :run_server
    goto launch_mode
)

:: 自定义参数
if "%mode_choice%"=="3" (
    call :custom_params
    goto launch_mode
)

echo 无效选项！
timeout /t 2 >nul
goto launch_mode

:: ==========================================
:: 查看模型详细信息
:: ==========================================
:model_info_menu
cls
call :list_models
if %count% equ 0 (
    echo 没有找到任何模型文件！
    pause
    goto main_menu
)

echo ----------------------------------------
echo 0. 返回主菜单
echo.
set /p info_choice=请选择要查看的模型 (0-%count%)：

if "%info_choice%"=="0" goto main_menu
call :validate_input "%info_choice%" %count%
if errorlevel 1 goto model_info_menu

set "info_model=!model%info_choice%!"
set "info_path=%model_path%\%info_model%"

cls
echo ========================================
echo 模型详细信息
echo ========================================
echo 文件名: %info_model%
echo 完整路径: %info_path%
echo.
echo 文件大小:
dir "%info_path%" /-c | findstr "%info_model%"
echo.
echo 文件哈希 (MD5):
certutil -hashfile "%info_path%" MD5 2>nul | findstr /v "MD5 certutil"
if errorlevel 1 echo 无法计算哈希值
echo.
pause
goto main_menu

:: ==========================================
:: 批量测试模式
:: ==========================================
:batch_test
cls
echo ========================================
echo 批量测试模式
echo ========================================
echo 将对所有模型执行简单测试
echo 测试提示词: "你好，请介绍一下你自己"
echo.
echo 按任意键开始测试...
pause >nul

call :list_models
if %count% equ 0 (
    echo 没有找到任何模型文件！
    pause
    goto main_menu
)

echo.
echo 开始批量测试...
echo ========================================

for /l %%i in (1,1,%count%) do (
    set "test_model=!model%%i!"
    set "test_path=%model_path%\!test_model!"
    echo.
    echo [%%i/%count%] 测试模型: !test_model!
    echo ----------------------------------------
    
    if exist "!test_path!" (
        "%cli_exe%" -m "!test_path!" -p "你好，请介绍一下你自己" -n 100 --no-display-prompt 2>nul
    ) else (
        echo 错误: 模型文件不存在！
    )
    
    echo.
    echo ----------------------------------------
    if %%i lss %count% (
        echo 按任意键测试下一个模型...
        pause >nul
    )
)

echo.
echo 批量测试完成！
pause
goto main_menu

:: ==========================================
:: 模型管理
:: ==========================================
:model_management
cls
echo ========================================
echo 模型管理
echo ========================================
echo [1] 刷新模型列表
echo [2] 重命名模型
echo [3] 删除模型（谨慎使用）
echo [4] 查看磁盘空间
echo [5] 备份模型
echo [6] 返回主菜单
echo.
set /p mgmt_choice=请选择 (1-6):

if "%mgmt_choice%"=="1" goto model_management
if "%mgmt_choice%"=="2" goto rename_model
if "%mgmt_choice%"=="3" goto delete_model
if "%mgmt_choice%"=="4" goto check_space
if "%mgmt_choice%"=="5" goto backup_model
if "%mgmt_choice%"=="6" goto main_menu

goto model_management

:rename_model
cls
call :list_models
echo ----------------------------------------
set /p rename_choice=请选择要重命名的模型 (1-%count%)：
call :validate_input "%rename_choice%" %count%
if errorlevel 1 goto rename_model

set "old_name=!model%rename_choice%!"
set "old_path=%model_path%\%old_name%"

echo.
echo 当前名称: %old_name%
set /p new_name=请输入新名称（含扩展名）:

if exist "%model_path%\%new_name%" (
    echo 错误: 文件 %new_name% 已存在！
    pause
    goto model_management
)

ren "%old_path%" "%new_name%"
if errorlevel 1 (
    echo 重命名失败！
) else (
    echo 重命名成功！
)
pause
goto model_management

:delete_model
cls
call :list_models
echo ----------------------------------------
set /p del_choice=请选择要删除的模型 (1-%count%)：
call :validate_input "%del_choice%" %count%
if errorlevel 1 goto delete_model

set "del_name=!model%del_choice%!"
set "del_path=%model_path%\%del_name%"

echo.
echo 警告: 即将删除模型: %del_name%
echo 文件大小: 
dir "%del_path%" /-c | findstr "%del_name%"
echo.
set /p confirm=确认删除？(输入 YES 确认):

if /i "%confirm%"=="YES" (
    del "%del_path%"
    if errorlevel 1 (
        echo 删除失败！
    ) else (
        echo 删除成功！
    )
) else (
    echo 操作已取消
)
pause
goto model_management

:check_space
cls
echo ========================================
echo 磁盘空间信息
echo ========================================
echo 模型文件夹: %model_path%
echo.
wmic logicaldisk where "caption='%model_path:~0,2%'" get deviceid,size,freespace /format:table
echo.
echo 模型文件夹大小:
powershell -command "Get-ChildItem -Path '%model_path%' -Recurse | Measure-Object -Property Length -Sum | ForEach-Object { '{0:N2} GB' -f ($_.Sum / 1GB) }" 2>nul
if errorlevel 1 echo 无法计算大小（需要PowerShell支持）
pause
goto model_management

:backup_model
cls
call :list_models
echo ----------------------------------------
set /p backup_choice=请选择要备份的模型 (1-%count%)：
call :validate_input "%backup_choice%" %count%
if errorlevel 1 goto backup_model

set "backup_name=!model%backup_choice%!"
set "backup_path=%model_path%\%backup_name%"
set "backup_dest=%llama_path%\backups"

if not exist "%backup_dest%" mkdir "%backup_dest%"

echo 正在备份 %backup_name% ...
copy "%backup_path%" "%backup_dest%\%backup_name%" >nul
if errorlevel 1 (
    echo 备份失败！
) else (
    echo 备份成功！已保存到 %backup_dest%
)
pause
goto model_management

:: ==========================================
:: 参数配置模板
:: ==========================================
:preset_config
cls
echo ========================================
echo 参数配置模板
echo ========================================
echo [1] 快速响应 - 适合实时对话
echo [2] 高质量输出 - 适合创作内容
echo [3] 平衡模式 - 默认推荐
echo [4] 低资源模式 - 适合低配置电脑
echo [5] 返回主菜单
echo.
set /p preset=请选择 (1-5):

if "%preset%"=="5" goto main_menu

if "%preset%"=="1" (
    set "cli_preset=--temp 0.8 --top-p 0.8 --ctx-size 2048 --threads 4"
    set "preset_name=快速响应"
)
if "%preset%"=="2" (
    set "cli_preset=--temp 0.1 --top-p 0.95 --ctx-size 4096 --repeat-penalty 1.1 --threads 8"
    set "preset_name=高质量输出"
)
if "%preset%"=="3" (
    set "cli_preset=--temp 0.7 --top-p 0.9 --ctx-size 3072 --threads 6"
    set "preset_name=平衡模式"
)
if "%preset%"=="4" (
    set "cli_preset=--temp 0.7 --top-p 0.9 --ctx-size 1024 --threads 2"
    set "preset_name=低资源模式"
)

echo.
echo 已选择: %preset_name%
echo 参数: %cli_preset%
echo.
echo 是否保存此配置？
set /p save_preset=保存为默认配置？(Y/N):

if /i "%save_preset%"=="Y" (
    echo default_preset=%preset_name% > "%llama_path%\user_config.txt"
    echo default_args=%cli_preset% >> "%llama_path%\user_config.txt"
    echo 配置已保存！
)

echo.
echo 现在可以选择模型运行...
pause
goto select_model

:: ==========================================
:: 模型对比模式
:: ==========================================
:compare_models
cls
echo ========================================
echo 模型对比模式
echo ========================================
call :list_models
echo ----------------------------------------
echo 请选择第一个模型:
set /p model1=选择模型 (1-%count%)：
call :validate_input "%model1%" %count%
if errorlevel 1 goto compare_models

set "model1_name=!model%model1%!"
set "model1_path=%model_path%\!model1_name!"

echo.
echo 请选择第二个模型:
set /p model2=选择模型 (1-%count%)：
call :validate_input "%model2%" %count%
if errorlevel 1 goto compare_models

set "model2_name=!model%model2%!"
set "model2_path=%model_path%\!model2_name!"

echo.
set /p compare_prompt=请输入测试提示词 (直接回车使用默认):
if "%compare_prompt%"=="" set "compare_prompt=请用一句话介绍人工智能"

cls
echo ========================================
echo 模型对比测试
echo ========================================
echo 提示词: %compare_prompt%
echo.
echo [模型1] %model1_name%
echo ----------------------------------------
"%cli_exe%" -m "%model1_path%" -p "%compare_prompt%" -n 150 --no-display-prompt 2>nul
echo.
echo ========================================
echo [模型2] %model2_name%
echo ----------------------------------------
"%cli_exe%" -m "%model2_path%" -p "%compare_prompt%" -n 150 --no-display-prompt 2>nul
echo.
echo ========================================
pause
goto main_menu

:: ==========================================
:: 查看对话日志
:: ==========================================
:view_logs
cls
echo ========================================
echo 对话日志
echo ========================================
if not exist "%log_dir%\*.log" (
    echo 暂无日志文件
    pause
    goto main_menu
)

echo 最近的日志文件:
dir "%log_dir%\*.log" /od /b
echo.
set /p view_log=请输入要查看的日志文件名 (直接回车返回):

if "%view_log%"=="" goto main_menu
if exist "%log_dir%\%view_log%" (
    type "%log_dir%\%view_log%"
) else (
    echo 文件不存在！
)
pause
goto main_menu

:: ==========================================
:: 性能测试
:: ==========================================
:performance_test
cls
call :list_models
echo ----------------------------------------
set /p perf_choice=请选择要测试的模型 (1-%count%)：
call :validate_input "%perf_choice%" %count%
if errorlevel 1 goto performance_test

set "perf_model=!model%perf_choice%!"
set "perf_path=%model_path%\%perf_model%"

cls
echo ========================================
echo 性能测试 - %perf_model%
echo ========================================
echo 测试内容: 生成100个token
echo.
echo 开始时间: %time%
set "start_time=%time%"

"%cli_exe%" -m "%perf_path%" -p "介绍一下人工智能的发展历史" -n 100 --no-display-prompt >nul 2>&1

echo 结束时间: %time%
echo.
echo 按任意键查看详细性能报告...
pause >nul

"%cli_exe%" -m "%perf_path%" -p "测试" -n 10 --metrics 2>&1 | findstr /i "time speed"
pause
goto main_menu

:: ==========================================
:: 加载保存的配置
:: ==========================================
:load_config
cls
if not exist "%config_file%" (
    echo 暂无保存的配置！
    echo.
    echo 提示: 在运行模型时可以选择保存配置
    pause
    goto main_menu
)

echo ========================================
echo 保存的配置
echo ========================================
type "%config_file%"
echo.
echo 0. 返回主菜单
set /p config_choice=请选择配置 (输入配置名称):

if "%config_choice%"=="0" goto main_menu

:: 查找配置
for /f "tokens=1,2 delims=|" %%a in ('findstr /i "^%config_choice%|" "%config_file%" 2^>nul') do (
    set "saved_model=%%a"
    set "saved_args=%%b"
)

if defined saved_model (
    set "full_model_path=%saved_model%"
    echo 加载配置: %config_choice%
    echo 模型: %saved_model%
    echo 参数: %saved_args%
    echo.
    "%cli_exe%" -m "%saved_model%" %saved_args%
) else (
    echo 未找到配置！
)
pause
goto main_menu

:open_folder
start explorer "%model_path%"
goto main_menu

:: ==========================================
:: 子函数
:: ==========================================

:list_models
set count=0
echo 可用的模型列表:
echo ----------------------------------------
for /f "delims=" %%i in ('dir /b /a-d "%model_path%\*" 2^>nul') do (
    set /a count+=1
    set "model!count!=%%i"
    echo !count!. %%i
)
exit /b

:validate_input
set "input=%~1"
set "max=%2"
echo %input%|findstr /r "^[0-9]*$" >nul
if errorlevel 1 (
    echo 无效输入，请输入数字！
    timeout /t 2 >nul
    exit /b 1
)
if %input% lss 1 (
    echo 输入不能小于1！
    timeout /t 2 >nul
    exit /b 1
)
if %input% gtr %max% (
    echo 输入超出范围！
    timeout /t 2 >nul
    exit /b 1
)
exit /b 0

:run_interactive
set "timestamp=%date:~0,4%%date:~5,2%%date:~8,2%_%time:~0,2%%time:~3,2%%time:~6,2%"
set "timestamp=%timestamp: =0%"
set "log_file=%log_dir%\conversation_%timestamp%.log"

echo.
echo ========================================
echo 启动交互式对话模式
echo ========================================
echo 模型: %selected_model%
echo.
echo 使用说明:
echo - 直接输入问题，按回车发送
echo - 按 Ctrl+C 退出对话
echo ========================================
echo.

:: 直接运行模型，不使用 tee（Windows不支持）
if defined cli_preset (
    echo 使用预设参数: %preset_name%
    "%cli_exe%" -m "%full_model_path%" --color auto %cli_preset%
) else (
    "%cli_exe%" -m "%full_model_path%" --color auto
)

echo.
echo 对话已结束
pause
exit /b

:run_server
echo.
echo ========================================
echo 启动 Web 服务器模式
echo ========================================
echo 模型: %selected_model%
echo 地址: http://127.0.0.1:8080
echo.
echo 正在启动服务器...
echo 提示: 服务器将在后台运行，关闭窗口停止服务
echo.
start "LLaMA Server" "%server_exe%" -m "%full_model_path%" --host 127.0.0.1 --port 8080
timeout /t 3 >nul
start http://127.0.0.1:8080
echo 服务器已启动！
echo.
echo 按任意键返回菜单（服务器仍在后台运行）
pause >nul
exit /b

:custom_params
echo.
echo 请输入自定义参数 (例如: --temp 0.5 --top-p 0.95)
echo 直接回车使用默认参数
set /p custom_args=参数:
echo.
echo 启动模型...
if "%custom_args%"=="" (
    "%cli_exe%" -m "%full_model_path%" --color auto
) else (
    "%cli_exe%" -m "%full_model_path%" --color auto %custom_args%
)
echo.
echo 对话已结束
pause
:: ==========================================
:: 下载模型 ★新增功能★
:: ==========================================
:download_model
cls
echo ========================================
echo           下载模型
echo ========================================
echo.
echo 支持从 Hugging Face 下载模型
echo 下载的模型会自动保存到: %model_path%
echo.
echo ========================================
echo.
echo [1] 从模型ID下载 (推荐)
echo [2] 搜索模型 (需要python环境)
echo [3] 查看已安装模型
echo [4] 返回主菜单
echo.
set /p dl_choice=请选择 (1-4):

if "%dl_choice%"=="4" goto main_menu
if "%dl_choice%"=="1" goto download_by_id
if "%dl_choice%"=="2" goto search_model
if "%dl_choice%"=="3" goto list_downloaded

goto download_model

:download_by_id
cls
echo ========================================
echo 从Hugging Face下载模型
echo ========================================
echo.
echo 示例模型ID:
echo   - Qwen/Qwen2-7B-Instruct-GGUF
echo   - TheBloke/Llama-2-7B-Chat-GGUF
echo   - mistralai/Mistral-7B-Instruct-v0.3
echo.
echo 国内镜像加速: 将自动使用 hf-mirror.com
echo.
set /p model_id=请输入模型ID (例如: TheBloke/Llama-2-7B-Chat-GGUF):

if "%model_id%"=="" goto download_model

echo.
echo ========================================
echo 选择下载类型:
echo [1] 下载完整模型 (HuggingFace格式)
echo [2] 下载GGUF格式文件 (推荐用于llama.cpp)
echo.
set /p format_choice=请选择 (1-2):

if "%format_choice%"=="1" goto download_hf_full
if "%format_choice%"=="2" goto download_gguf

goto download_model

:download_hf_full
cls
echo ========================================
echo 下载完整模型: %model_id%
echo ========================================
echo.

:: 检查Python环境
python --version >nul 2>&1
if errorlevel 1 (
    echo 错误: 未找到Python！
    echo.
    echo 请先安装Python: https://www.python.org/downloads/
    echo 安装时请勾选 "Add Python to PATH"
    pause
    goto download_model
)

:: 安装huggingface_hub（如果需要）
pip show huggingface_hub >nul 2>&1
if errorlevel 1 (
    echo 正在安装 huggingface_hub 工具...
    pip install -U huggingface_hub -i https://pypi.tuna.tsinghua.edu.cn/simple
)

:: 使用国内镜像下载
echo 正在下载模型...
echo 提示: 大模型下载时间较长，请耐心等待...
echo.

set "target_dir=%model_path%\%model_id:\=%"
set "target_dir=%target_dir:/=%"

:: 设置国内镜像加速
set "HF_ENDPOINT=https://hf-mirror.com"

:: 执行下载
huggingface-cli download %model_id% --local-dir "%target_dir%" --local-dir-use-symlinks False --resume-download

if errorlevel 1 (
    echo.
    echo 下载失败！
    echo 可能原因: 模型ID错误或需要授权
    echo.
    echo 如果是需要授权的模型，请先登录:
    echo huggingface-cli login
) else (
    echo.
    echo ========================================
    echo 下载成功！
    echo 保存位置: %target_dir%
    echo ========================================
)

echo.
pause
goto download_model

:download_gguf
cls
echo ========================================
echo 下载GGUF格式模型
echo ========================================
echo.
echo 提示: GGUF格式可直接用于llama.cpp
echo 推荐使用 TheBloke 等作者提供的GGUF版本
echo.
echo 示例: TheBloke/Llama-2-7B-Chat-GGUF
echo.
set /p gguf_id=请输入GGUF模型ID:

if "%gguf_id%"=="" goto download_model

echo.
echo 正在获取模型文件列表...
echo.

:: 使用Python脚本获取GGUF文件列表并下载
python -c "
import subprocess
import sys
import os

model_id = r'%gguf_id%'
target_dir = r'%model_path%\%gguf_id:\=%'
target_dir = target_dir.replace('/', '\\')

# 设置镜像
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

try:
    from huggingface_hub import list_repo_files, hf_hub_download
    
    print('正在获取文件列表...')
    files = list_repo_files(model_id)
    
    # 筛选GGUF文件
    gguf_files = [f for f in files if f.endswith('.gguf')]
    
    if not gguf_files:
        print('未找到GGUF文件！请检查模型ID是否正确')
        print('该模型可能不是GGUF格式，请尝试使用\"下载完整模型\"选项')
        sys.exit(1)
    
    print(f'找到 {len(gguf_files)} 个GGUF文件:')
    for i, f in enumerate(gguf_files):
        print(f'  {i+1}. {f}')
    
    print()
    if len(gguf_files) > 1:
        choice = input('请选择要下载的文件编号 (直接回车下载全部): ')
        if choice.strip():
            try:
                idx = int(choice) - 1
                gguf_files = [gguf_files[idx]]
            except:
                print('无效选择，将下载全部文件')
    
    # 创建目标目录
    os.makedirs(target_dir, exist_ok=True)
    
    # 下载文件
    for filename in gguf_files:
        print(f'正在下载: {filename}')
        local_path = hf_hub_download(
            repo_id=model_id,
            filename=filename,
            local_dir=target_dir,
            local_dir_use_symlinks=False,
            resume=True
        )
        print(f'完成: {local_path}')
    
    print('\n所有文件下载完成！')
    
except Exception as e:
    print(f'错误: {e}')
    print('\n提示: 如果下载失败，请尝试:')
    print('1. pip install -U huggingface_hub')
    print('2. 检查网络连接')
    sys.exit(1)
"

if errorlevel 1 (
    echo.
    echo 下载失败！
) else (
    echo.
    echo ========================================
    echo 下载成功！模型已保存到模型文件夹
    echo ========================================
)

pause
goto download_model

:search_model
cls
echo ========================================
echo 搜索Hugging Face模型
echo ========================================
echo.
set /p search_keyword=请输入搜索关键词:

if "%search_keyword%"=="" goto download_model

echo.
echo 正在搜索...
echo.

python -c "
import requests
import json

keyword = r'%search_keyword%'
url = f'https://huggingface.co/api/models?search={keyword}&limit=20'

try:
    response = requests.get(url, timeout=10)
    models = response.json()
    
    if not models:
        print('未找到相关模型')
    else:
        print(f'找到 {len(models)} 个相关模型:\n')
        for i, model in enumerate(models[:15]):
            model_id = model.get('modelId', 'unknown')
            downloads = model.get('downloads', 0)
            tags = model.get('tags', [])
            is_gguf = 'gguf' in str(tags).lower()
            marker = '[GGUF]' if is_gguf else '[HF]'
            print(f'{i+1}. {marker} {model_id}')
            print(f'   下载量: {downloads}')
            print()
            
except Exception as e:
    print(f'搜索失败: {e}')
    print('提示: 可能需要安装requests: pip install requests')
"

echo.
echo 提示: 找到喜欢的模型后，使用选项1输入模型ID即可下载
pause
goto download_model

:list_downloaded
cls
echo ========================================
echo 已下载的模型
echo ========================================
echo.
echo 模型文件夹: %model_path%
echo.
dir "%model_path%" /b /ad 2>nul
if errorlevel 1 echo 暂无下载的模型
echo.
echo ========================================
echo [1] 下载更多模型
echo [2] 返回主菜单
echo.
set /p list_choice=请选择:

if "%list_choice%"=="1" goto download_model
goto main_menu
exit /b