#!/bin/bash

# LLM生成测试脚本
# 提供简单的命令行界面来运行LLM生成测试

set -e  # 遇到错误时退出

# 颜色定义
RED='\\033[0;31m'
GREEN='\\033[0;32m'
YELLOW='\\033[1;33m'
BLUE='\\033[0;34m'
NC='\\033[0m' # No Color

# 脚本目录
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 显示帮助信息
show_help() {
    cat << EOF
LLM Generation Test Script

Usage: $0 [OPTIONS] [COMMAND]

Commands:
    basic           Run basic LLM generation tests
    comprehensive   Run comprehensive test suite
    unit           Run unit tests only
    integration    Run integration tests only
    performance    Run performance tests only
    cli            Run CLI tests only
    validation     Run validation tests only
    all            Run all tests (default)
    help           Show this help message

Options:
    -h, --help      Show this help message
    -v, --verbose   Enable verbose output
    -c, --config    Specify config file path
    -k, --api-key    Specify API key (overrides environment variable)

Environment Variables:
    SILICONFLOW_API_KEY    SiliconFlow API key (required for LLM tests)

Examples:
    $0                           # Run all tests
    $0 basic                     # Run basic tests only
    $0 -v --config custom.yaml   # Run with custom config and verbose output

EOF
}

# 检查依赖
check_dependencies() {
    log_info "Checking dependencies..."

    # 检查Python
    if ! command -v python3 &> /dev/null; then
        log_error "Python3 is required but not installed"
        exit 1
    fi

    # 检查pip
    if ! command -v pip3 &> /dev/null; then
        log_error "pip3 is required but not installed"
        exit 1
    fi

    # 检查必要的Python包
    local required_packages=("aiohttp" "pytest" "pyyaml" "click" "colorama")
    for package in "${required_packages[@]}"; do
        if ! python3 -c "import $package" 2>/dev/null; then
            log_warning "Package '$package' not found. Installing..."
            pip3 install "$package"
        fi
    done

    log_success "Dependencies check passed"
}

# 检查API密钥
check_api_key() {
    if [[ -z "$SILICONFLOW_API_KEY" && -z "$API_KEY" ]]; then
        log_warning "SILICONFLOW_API_KEY environment variable not set"
        log_warning "LLM tests will be skipped. Set the environment variable to run LLM tests."
        return 1
    fi

    # 如果提供了API密钥参数，设置环境变量
    if [[ -n "$API_KEY" ]]; then
        export SILICONFLOW_API_KEY="$API_KEY"
    fi

    return 0
}

# 设置环境
setup_environment() {
    log_info "Setting up environment..."

    # 切换到项目根目录
    cd "$PROJECT_ROOT"

    # 设置Python路径
    export PYTHONPATH="$PROJECT_ROOT:$PYTHONPATH"

    # 创建必要的目录
    mkdir -p test_results/{reports,logs,specs,performance,validation}

    log_success "Environment setup completed"
}

# 运行基本测试
run_basic_tests() {
    log_info "Running basic LLM generation tests..."

    if check_api_key; then
        python3 tests/run_llm_tests.py --test-type llm --llm-type basic ${VERBOSE_ARG} ${CONFIG_ARG}
    else
        log_warning "Skipping LLM tests due to missing API key"
        return 1
    fi
}

# 运行全面测试
run_comprehensive_tests() {
    log_info "Running comprehensive test suite..."

    if check_api_key; then
        python3 tests/run_llm_tests.py --test-type llm --llm-type comprehensive ${VERBOSE_ARG} ${CONFIG_ARG}
    else
        log_warning "Skipping LLM tests due to missing API key"
        return 1
    fi
}

# 运行单元测试
run_unit_tests() {
    log_info "Running unit tests..."

    python3 tests/run_llm_tests.py --test-type unit ${VERBOSE_ARG} ${CONFIG_ARG}
}

# 运行集成测试
run_integration_tests() {
    log_info "Running integration tests..."

    python3 tests/run_llm_tests.py --test-type integration ${VERBOSE_ARG} ${CONFIG_ARG}
}

# 运行性能测试
run_performance_tests() {
    log_info "Running performance tests..."

    if check_api_key; then
        python3 tests/run_llm_tests.py --test-type performance ${VERBOSE_ARG} ${CONFIG_ARG}
    else
        log_warning "Skipping performance tests due to missing API key"
        return 1
    fi
}

# 运行CLI测试
run_cli_tests() {
    log_info "Running CLI tests..."

    if check_api_key; then
        python3 tests/run_llm_tests.py --test-type cli ${VERBOSE_ARG} ${CONFIG_ARG}
    else
        log_warning "Skipping CLI tests due to missing API key"
        return 1
    fi
}

# 运行验证测试
run_validation_tests() {
    log_info "Running validation tests..."

    python3 tests/run_llm_tests.py --test-type validation ${VERBOSE_ARG} ${CONFIG_ARG}
}

# 运行所有测试
run_all_tests() {
    log_info "Running all tests..."

    python3 tests/run_llm_tests.py --test-type all ${VERBOSE_ARG} ${CONFIG_ARG}
}

# 显示结果
show_results() {
    local exit_code=$1
    local test_type=$2

    if [[ $exit_code -eq 0 ]]; then
        log_success "$test_type tests completed successfully"
    else
        log_error "$test_type tests failed with exit code $exit_code"
    fi

    # 显示结果目录
    if [[ -d "test_results" ]]; then
        log_info "Results saved in: test_results/"
        ls -la test_results/ 2>/dev/null || true
    fi
}

# 主函数
main() {
    local command="all"
    local verbose=false
    local config=""
    local api_key=""

    # 解析命令行参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -v|--verbose)
                verbose=true
                VERBOSE_ARG="--verbose"
                shift
                ;;
            -c|--config)
                config="$2"
                CONFIG_ARG="--config $config"
                shift 2
                ;;
            -k|--api-key)
                api_key="$2"
                API_KEY="$2"
                shift 2
                ;;
            basic|comprehensive|unit|integration|performance|cli|validation|all)
                command="$1"
                shift
                ;;
            *)
                log_error "Unknown command: $1"
                show_help
                exit 1
                ;;
        esac
    done

    # 显示欢迎信息
    echo "=========================================================="
    echo "LLM Generation Test Framework"
    echo "=========================================================="
    echo "Command: $command"
    echo "Project: $PROJECT_ROOT"
    if [[ -n "$config" ]]; then
        echo "Config: $config"
    fi
    echo "=========================================================="

    # 检查依赖
    check_dependencies

    # 设置环境
    setup_environment

    # 根据命令运行相应的测试
    local exit_code=0
    case $command in
        basic)
            run_basic_tests
            exit_code=$?
            ;;
        comprehensive)
            run_comprehensive_tests
            exit_code=$?
            ;;
        unit)
            run_unit_tests
            exit_code=$?
            ;;
        integration)
            run_integration_tests
            exit_code=$?
            ;;
        performance)
            run_performance_tests
            exit_code=$?
            ;;
        cli)
            run_cli_tests
            exit_code=$?
            ;;
        validation)
            run_validation_tests
            exit_code=$?
            ;;
        all)
            run_all_tests
            exit_code=$?
            ;;
        *)
            log_error "Unknown command: $command"
            show_help
            exit 1
            ;;
    esac

    # 显示结果
    show_results $exit_code "$command"

    # 退出
    exit $exit_code
}

# 运行主函数
main "$@"