#!/usr/bin/env bash
set -euo pipefail

# ─────────────────────────────────────────────────────────────────────────────
# 21-llm-security.sh — SecureGit LLM Red Teaming & Hardening
#
# Workflows to demonstrate armyknife-redteam scanning of local and remote
# language models via SecureGit. Provides interactive options to configure 
# scans, benchmarks, and apply pipeline remediations.
#
# Usage:
#   ./21-llm-security.sh                     # Interactive redteam menu
#   ./21-llm-security.sh scan                # Run standard local scan
#   ./21-llm-security.sh gpubox              # Run remote scan against gpubox
#   ./21-llm-security.sh compare             # Compare two endpoint models
# ─────────────────────────────────────────────────────────────────────────────

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
source "$SCRIPT_DIR/lib/securegit-common.sh"

export SECUREGIT_REDTEAM_BIN="$HOME/Projects/armyknife-redteam/target/release/armyknife-llm-redteam-mcp"
export ARMYKNIFE_BIN="$HOME/Projects/armyknife-redteam/target/release/armyknife-llm-redteam"

_print_header() {
    sg_header "ArmyKnife AI Red Team Interface"
    echo "============================================================"
    echo "  This workflow executes automated penetration tests and     "
    echo "  vulnerability scans against LLMs using the ArmyKnife    "
    echo "  security engine and the SecureGit integration bridge.     "
    echo "============================================================"
    echo
}

op_scan() {
    sg_info "Executing local model scan via SecureGit MCP Bridge..."
    echo
    if ! command -v securegit &>/dev/null; then
        sg_die "securegit binary not found in PATH."
    fi
    
    local model
    model="$(sg_prompt "[Local] Model name to scan" "glm-4.7-flash:latest")"
    
    sg_info "Running: securegit hf scan \"$model\""
    securegit hf scan "$model"
}

op_gpubox() {
    sg_info "Executing remote GPU Box model scan via SecureGit MCP Bridge..."
    echo
    
    local endpoint
    endpoint="$(sg_prompt "Remote Address" "100.127.20.90:11434")"
    
    local model
    model="$(sg_prompt "[Remote] Model name to scan" "gemma3:4b")"
    
    local full_uri="ollama://$endpoint/$model"
    
    sg_info "Running: securegit hf scan \"$full_uri\""
    securegit hf scan "$full_uri"
}

op_compare() {
    sg_info "Executing A/B Red Team Comparison via ArmyKnife..."
    echo
    
    local model_a
    model_a="$(sg_prompt "Model A (Local)" "glm-4.7-flash:latest")"
    
    local model_b
    model_b="$(sg_prompt "Model B (Remote)" "ollama://100.127.20.90:11434/gemma3:4b")"
    
    local probes
    probes="$(sg_prompt "Probe Category (e.g. SystemPromptLeak, PromptInjection)" "SystemPromptLeak")"
    
    sg_info "Running comparison: $model_a vs $model_b"
    "$ARMYKNIFE_BIN" compare --model-a "$model_a" --model-b "$model_b" --probes "$probes"
}

op_benchmark() {
    sg_info "Executing Latency and Security Benchmark across LLMs..."
    echo
    
    local model
    model="$(sg_prompt "Model to benchmark" "glm-4.7-flash:latest")"
    
    local delay
    delay="$(sg_prompt "Throttling Delay between requests (ms)" "500")"
    
    sg_info "Running benchmark with ${delay}ms delay..."
    "$ARMYKNIFE_BIN" benchmark run --model "$model" --delay-ms "$delay" --concurrency 1
}

op_harden() {
    sg_info "Executing Remediation Generation for Vulnerable Model..."
    echo
    
    local model
    model="$(sg_prompt "Vulnerable Model" "glm-4.7-flash:latest")"
    
    sg_info "Generating Modelfiles and API Firewalls for $model..."
    "$ARMYKNIFE_BIN" remediate --model "$model"
}

op_pipelines() {
    sg_info "Executing Defensive Pipeline Checks..."
    echo
    
    local model
    model="$(sg_prompt "Model to evaluate" "glm-4.7-flash:latest")"
    
    sg_info "Running gate scan..."
    "$ARMYKNIFE_BIN" pipeline scan --model "$model" --output /tmp/pipeline-results
    cat /tmp/pipeline-results
}

main() {
    _print_header

    if [[ "${1:-}" == "scan" ]]; then
        op_scan
    elif [[ "${1:-}" == "gpubox" ]]; then
        op_gpubox
    elif [[ "${1:-}" == "compare" ]]; then
        op_compare
    elif [[ "${1:-}" == "benchmark" ]]; then
        op_benchmark
    elif [[ "${1:-}" == "harden" ]]; then
        op_harden    
    elif [[ "${1:-}" == "pipelines" ]]; then
        op_pipelines
    else
        local options=(
            "Standard Security Scan (Localhost)"
            "Remote Security Scan (GPU Box)"
            "A/B Model Vulnerability Comparison"
            "Throttled Rate-Limited Benchmark"
            "Generate Remediation Firewalls"
            "CI/CD Pipeline Security Gate Evaluation"
            "Exit"
        )
        
        while true; do
            local choice
            choice="$(sg_menu "Select an AI Security Operation" "${options[@]}")"
            
            case "$choice" in
                1) op_scan ;;
                2) op_gpubox ;;
                3) op_compare ;;
                4) op_benchmark ;;
                5) op_harden ;;
                6) op_pipelines ;;
                7) exit 0 ;;
            esac
            
            echo
            sg_prompt_any_key "Press any key to return to menu..."
            clear
            _print_header
        done
    fi
}

main "$@"