#!/bin/bash
# Usage: ./test-workflows.sh [create-sharedpool|create-cluster|delete-cluster|delete-sharedpool|all]

set -e

CLUSTER_SERVICE_GRPC="localhost:10000"
MOCK_DPM_ADMIN="http://localhost:8081"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

SHARED_POOL_ID="testsp$(date +%s)"
CLUSTER_ID="testcluster$(date +%s)"
TENANT_ID="testtenant001"
DATAPLANE_NAME="testdataplane001"

log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

check_prerequisites() {
    log_info "Checking prerequisites..."
    
    if ! command -v grpcurl &> /dev/null; then
        log_error "grpcurl not installed. Install with: brew install grpcurl"
        exit 1
    fi
    
    if ! grpcurl -plaintext $CLUSTER_SERVICE_GRPC list &> /dev/null; then
        log_error "cluster-service gRPC is not responding at $CLUSTER_SERVICE_GRPC"
        log_error "Make sure to start with --mock flag for reflection support"
        exit 1
    fi
    
    if ! curl -s "$MOCK_DPM_ADMIN/healthz" &> /dev/null; then
        log_error "mock-dpm is not running at $MOCK_DPM_ADMIN"
        exit 1
    fi
    
    log_info "All prerequisites met!"
}

list_services() {
    log_info "Listing gRPC services..."
    grpcurl -plaintext $CLUSTER_SERVICE_GRPC list
}

create_sharedpool() {
    log_info "Creating SharedPool: $SHARED_POOL_ID"
    
    grpcurl -plaintext -d "{
        \"shared_pool\": {
            \"id\": \"$SHARED_POOL_ID\"
        },
        \"request_unit\": {\"min\": 1000, \"max\": 10000},
        \"init_shared_pool_config\": {
            \"tikv_data_storage\": {\"type\": \"gp3\", \"size_gi\": 100, \"iops\": 3000, \"throughput_mbps\": 125},
            \"tikv_raftlog_storage\": {\"type\": \"gp3\", \"size_gi\": 50, \"iops\": 3000, \"throughput_mbps\": 125},
            \"tiflash_cn_storage\": {\"type\": \"gp3\", \"size_gi\": 100, \"iops\": 3000, \"throughput_mbps\": 125},
            \"tiflash_wn_storage\": {\"type\": \"gp3\", \"size_gi\": 100, \"iops\": 3000, \"throughput_mbps\": 125},
            \"network_cidr\": \"10.0.0.0/16\",
            \"service_cidr\": \"10.1.0.0/16\"
        },
        \"tenant_id\": \"$TENANT_ID\"
    }" $CLUSTER_SERVICE_GRPC cluster_service_ng.v1beta.SharedPoolService/CreateSharedPool
    
    echo "$SHARED_POOL_ID" > /tmp/test_sharedpool_id
    log_info "SharedPool creation initiated!"
}

get_sharedpool() {
    local sp_id="${1:-$SHARED_POOL_ID}"
    [ -f /tmp/test_sharedpool_id ] && sp_id=$(cat /tmp/test_sharedpool_id)
    
    log_info "Getting SharedPool: $sp_id"
    grpcurl -plaintext -d "{\"shared_pool_id\": \"$sp_id\"}" \
        $CLUSTER_SERVICE_GRPC cluster_service_ng.v1beta.SharedPoolService/GetSharedPool
}

list_sharedpools() {
    log_info "Listing SharedPools..."
    grpcurl -plaintext -d '{}' $CLUSTER_SERVICE_GRPC cluster_service_ng.v1beta.SharedPoolService/ListSharedPools
}

delete_sharedpool() {
    local sp_id="${1:-$SHARED_POOL_ID}"
    [ -f /tmp/test_sharedpool_id ] && sp_id=$(cat /tmp/test_sharedpool_id)
    
    log_info "Deleting SharedPool: $sp_id"
    grpcurl -plaintext -d "{\"shared_pool_id\": \"$sp_id\"}" \
        $CLUSTER_SERVICE_GRPC cluster_service_ng.v1beta.SharedPoolService/DeleteSharedPool
    log_info "SharedPool deletion initiated!"
}

create_cluster() {
    local sp_id="${1:-$SHARED_POOL_ID}"
    [ -f /tmp/test_sharedpool_id ] && sp_id=$(cat /tmp/test_sharedpool_id)
    
    log_info "Creating TiDBCluster: $CLUSTER_ID in SharedPool: $sp_id"
    
    grpcurl -plaintext -d "{
        \"id\": \"$CLUSTER_ID\",
        \"tenant_id\": \"$TENANT_ID\",
        \"shared_pool_id\": \"$sp_id\",
        \"request_unit\": {\"min\": 1000, \"max\": 5000},
        \"gateway_identity\": \"gateway-$CLUSTER_ID\"
    }" $CLUSTER_SERVICE_GRPC cluster_service_ng.v1beta.TiDBClusterService/CreateTiDBCluster
    
    echo "$CLUSTER_ID" > /tmp/test_cluster_id
    log_info "TiDBCluster creation initiated!"
}

get_cluster() {
    local cluster_id="${1:-$CLUSTER_ID}"
    [ -f /tmp/test_cluster_id ] && cluster_id=$(cat /tmp/test_cluster_id)
    
    log_info "Getting TiDBCluster: $cluster_id"
    grpcurl -plaintext -d "{\"cluster_id\": \"$cluster_id\"}" \
        $CLUSTER_SERVICE_GRPC cluster_service_ng.v1beta.TiDBClusterService/GetTiDBCluster
}

list_clusters() {
    log_info "Listing TiDBClusters..."
    grpcurl -plaintext -d '{}' $CLUSTER_SERVICE_GRPC cluster_service_ng.v1beta.TiDBClusterService/ListTiDBClusters
}

delete_cluster() {
    local cluster_id="${1:-$CLUSTER_ID}"
    [ -f /tmp/test_cluster_id ] && cluster_id=$(cat /tmp/test_cluster_id)
    
    log_info "Deleting TiDBCluster: $cluster_id"
    grpcurl -plaintext -d "{\"cluster_id\": \"$cluster_id\"}" \
        $CLUSTER_SERVICE_GRPC cluster_service_ng.v1beta.TiDBClusterService/DeleteTiDBCluster
    log_info "TiDBCluster deletion initiated!"
}

dump_mock_state() {
    log_info "Dumping mock-dpm state..."
    curl -s "$MOCK_DPM_ADMIN/dump" | jq .
}

reset_mock_state() {
    log_info "Resetting mock-dpm state..."
    curl -s -X POST "$MOCK_DPM_ADMIN/reset" | jq .
}

run_all_tests() {
    log_info "Running all tests..."
    
    check_prerequisites
    
    log_info "=== Test 1: Create SharedPool ==="
    create_sharedpool
    sleep 2
    
    log_info "=== Test 2: Get SharedPool ==="
    get_sharedpool
    
    log_info "=== Test 3: List SharedPools ==="
    list_sharedpools
    
    log_info "=== Test 4: Create TiDBCluster ==="
    create_cluster
    sleep 2
    
    log_info "=== Test 5: Get TiDBCluster ==="
    get_cluster
    
    log_info "=== Test 6: List TiDBClusters ==="
    list_clusters
    
    log_info "=== Test 7: Delete TiDBCluster ==="
    delete_cluster
    sleep 2
    
    log_info "=== Test 8: Delete SharedPool ==="
    delete_sharedpool
    
    log_info "=== Test 9: Check Mock-DPM State ==="
    dump_mock_state
    
    log_info "All tests completed!"
}

case "${1:-help}" in
    check)              check_prerequisites ;;
    list-services)      list_services ;;
    create-sharedpool)  check_prerequisites; create_sharedpool ;;
    get-sharedpool)     get_sharedpool "$2" ;;
    list-sharedpools)   list_sharedpools ;;
    delete-sharedpool)  delete_sharedpool "$2" ;;
    create-cluster)     check_prerequisites; create_cluster "$2" ;;
    get-cluster)        get_cluster "$2" ;;
    list-clusters)      list_clusters ;;
    delete-cluster)     delete_cluster "$2" ;;
    dump)               dump_mock_state ;;
    reset)              reset_mock_state ;;
    all)                run_all_tests ;;
    help|*)
        echo "Usage: $0 <command> [args]"
        echo ""
        echo "Commands:"
        echo "  check              Check prerequisites"
        echo "  list-services      List gRPC services"
        echo ""
        echo "  create-sharedpool  Create a test SharedPool"
        echo "  get-sharedpool     Get SharedPool status"
        echo "  list-sharedpools   List all SharedPools"
        echo "  delete-sharedpool  Delete SharedPool"
        echo ""
        echo "  create-cluster     Create a test TiDBCluster"
        echo "  get-cluster        Get TiDBCluster status"
        echo "  list-clusters      List all TiDBClusters"
        echo "  delete-cluster     Delete TiDBCluster"
        echo ""
        echo "  dump               Dump mock-dpm state"
        echo "  reset              Reset mock-dpm state"
        echo ""
        echo "  all                Run all tests"
        ;;
esac