#!/bin/bash
set -uo pipefail # -e: exit if any command fails

SCRIPT_NAME=$(basename "$0")
USAGE="Usage: $SCRIPT_NAME [-h] [-i <input_file>] [-a <auth_type>] [-t <api_key>] [-p <aws_profile>] [-H <openapi_host>]"

# Declare global variables
LAST_RESPONSE=""
EXIT_CODE=0

# Debug function
debug() {
  local log_dir="/tmp/llm-scenario-generator/logs"
  mkdir -p "$log_dir"
  local log_file="$log_dir/do-action-$(date +%Y%m%d).log"
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] DEBUG: $*" >> "$log_file"
}

# Function to check if a command exists
function command_exists {
  local cmd="$1"
  if command -v "$cmd" &>/dev/null; then
    return 0 # Success
  else
    return 1 # Failure
  fi
}

function py_module_exists {
  local module="$1"
  if python -m "$module" -h &>/dev/null; then
    return 0 # Success
  else
    return 1 # Failure
  fi
}

function is_http_response_expected {
  local expect=$1
  local resp_body=$2
  local status_code=$3

  local expected_status_code=""
  local expected_cel_expr=""
  if [[ -n "$expect" && "$expect" != "null" ]]; then
    # read expected_status_code expected_cel_expr <<<$(echo "$expect" | yq -r ".status_code, .cel")
    local expected_status_code=$(echo "$expect" | yq -r ".status_code")
    local expected_cel_expr=$(echo "$expect" | yq -r ".cel")
    # expected_cel_expr="${expected_cel_expr#body}"
    expected_cel_expr="${expected_cel_expr//body\./.}"
  fi

  if [[ -n "$expected_status_code" && "$expected_status_code" != "0" && "$expected_status_code" != "null" ]]; then
    if [[ "$expected_status_code" != $status_code ]]; then return 1; fi
  fi

  if [[ -n "$expected_cel_expr" && "$expected_cel_expr" != "null" ]]; then
    eval_result=$(python -m celpy -sb "$expected_cel_expr" <<<"$resp_body")
    if [[ "$eval_result" != "true" ]]; then return 1; fi
  fi
  return 0
}

function is_cmd_output_expected {
  local expect=$1
  local output=$2

  local expected_cel_expr=""
  if [[ -n "$expect" && "$expect" != "null" ]]; then
    expected_cel_expr=$(echo "$expect" | yq -r ".cel")
    expected_cel_expr="${expected_cel_expr#output}"
  fi

  if [[ -n "$expected_cel_expr" && "$expected_cel_expr" != "null" ]]; then
    eval_result=$(python -m celpy -sb "$expected_cel_expr" <<<"$output")
    if [[ "$eval_result" != "true" ]]; then return 1; fi
  fi
  return 0
}

function do_curl {
  local basepath=$1
  local auth_type=$2
  local apikey=$3
  local request=$4
  local expect=$5

  debug "do_curl: basepath=$basepath"
  debug "do_curl: auth_type=$auth_type"
  debug "do_curl: request=$request"
  debug "do_curl: expect=$expect"

  # read method path <<<$(echo "$request" | yq -r ".method, .path")
  local method=$(echo "$request" | yq -r ".method")
  local path=$(echo "$request" | yq -r ".path")
  local body=$(echo "$request" | yq -r ".body")
  if [[ -z "$body" ]]; then
    body='{}'
  fi

  debug "do_curl: method=$method"
  debug "do_curl: path=$path"
  debug "do_curl: body=$body"

  case "${auth_type}" in
  "digest")
    debug "do_curl: using digest auth"
    local response=$(curl \
      -s \
      --connect-timeout 10 \
      --max-time 30 \
      --retry 2 \
      --retry-delay 1 \
      --digest \
      -u "${apikey}" \
      -H 'content-type: application/json' \
      -X "$method" \
      -w "\n%{http_code}" \
      -d "${body}" \
      "${basepath}${path}" 2>&1)
    ;;
  "bearer")
    debug "do_curl: using bearer auth"
    local response=$(curl \
      -s \
      --connect-timeout 10 \
      --max-time 30 \
      --retry 2 \
      --retry-delay 1 \
      -H "Authorization: Bearer ${apikey}" \
      -H 'content-type: application/json' \
      -X "$method" \
      -w "\n%{http_code}" \
      -d "${body}" \
      "${basepath}${path}" 2>&1)
    ;;
  esac

  local curl_exit_code=$?
  local status_code=$(echo "$response" | tail -n 1)
  local resp_body=$(echo "$response" | sed '$d')

  debug "do_curl: status_code=$status_code"
  debug "do_curl: resp_body=$resp_body"

  if [ $curl_exit_code -ne 0 ] || [ "$status_code" == "000" ]; then
    LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"status_code\": $status_code, \"error\": \"Connection failed: $resp_body\"}}"
    EXIT_CODE=$curl_exit_code
    debug "do_curl: connection failed, LAST_RESPONSE=$LAST_RESPONSE"
    return $EXIT_CODE
  fi

  if [[ -z "$resp_body" ]]; then
    resp_body="{}"
  fi

  # Store response in global variable
  LAST_RESPONSE="{\"status_code\": $status_code, \"headers\": {}, \"body\": $resp_body}"
  debug "do_curl: LAST_RESPONSE=$LAST_RESPONSE"

  # Check expectations and return status
  is_http_response_expected "$expect" "$resp_body" "$status_code"
  local ret=$?
  debug "do_curl: is_http_response_expected returned $ret"

  return $ret
}

function do_cmd_aws {
  local request=$1
  local profile=$2
  local expect=$3

  debug "do_cmd_aws: request=$request"
  debug "do_cmd_aws: profile=$profile"
  debug "do_cmd_aws: expect=$expect"

  local args=($(echo "$request" | yq -r ".args[]"))
  if [[ -n "$profile" ]]; then
    args+=("--profile" "$profile")
  fi
  local output=$(aws "${args[@]}" 2>&1)
  local status=$?

  debug "do_cmd_aws: status=$status"
  debug "do_cmd_aws: output=$output"

  # If command failed, create an error response
  if [ $status -ne 0 ]; then
    LAST_RESPONSE="{\"status_code\": $status, \"json\": {\"output\": \"\"}, \"headers\": {}, \"error\": \"$(echo "$output" | jq -Rsa . | sed 's/^\"\(.*\)\"$/\1/')\"}"
    EXIT_CODE=$status
    debug "do_cmd_aws: command failed, LAST_RESPONSE=$LAST_RESPONSE"
    return $EXIT_CODE
  fi

  # Store response in global variable
  LAST_RESPONSE="$output"
  debug "do_cmd_aws: LAST_RESPONSE=$LAST_RESPONSE"

  # Check expectations and return status
  is_cmd_output_expected "$expect" "$output"
  local ret=$?
  debug "do_cmd_aws: is_cmd_output_expected returned $ret"

  return $ret
}

function do_cmd_gcloud {
  local request=$1
  local expect=$2

  debug "do_cmd_gcloud: request=$request"
  debug "do_cmd_gcloud: expect=$expect"

  local args=($(echo "$request" | yq -r ".args[]"))
  if [[ "${args[@]}" =~ "--format json" ]]; then
    args+=("--format" "json")
  fi
  local output=$(gcloud "${args[@]}" 2>&1)
  local status=$?

  debug "do_cmd_gcloud: output=$output"

  # If command failed, create an error response
  if [ $status -ne 0 ]; then
    LAST_RESPONSE="{\"status_code\": $status, \"json\": {\"output\": \"\"}, \"headers\": {}, \"error\": \"$(echo "$output" | jq -Rsa . | sed 's/^\"\(.*\)\"$/\1/')\"}"
    EXIT_CODE=$status
    debug "do_cmd_gcloud: command failed, LAST_RESPONSE=$LAST_RESPONSE"
    return $EXIT_CODE
  fi

  # Store response in global variable
  LAST_RESPONSE="$output"
  debug "do_cmd_gcloud: LAST_RESPONSE=$LAST_RESPONSE"

  # Check expectations and return status
  is_cmd_output_expected "$expect" "$output"
  local ret=$?
  debug "do_cmd_gcloud: is_cmd_output_expected returned $ret"

  return $ret
}

function do_cmd_az {
  local request=$1
  local expect=$2

  debug "do_cmd_az: request=$request"
  debug "do_cmd_az: expect=$expect"

  local args=($(echo "$request" | yq -r ".args[]"))

  # Execute command and capture both stdout and stderr
  local output
  local status
  output=$(az "${args[@]}" 2>&1)
  status=$?

  debug "do_cmd_az: status=$status"
  debug "do_cmd_az: output=$output"

  # If command failed, create an error response
  if [ $status -ne 0 ]; then
    LAST_RESPONSE="{\"status_code\": $status, \"json\": {\"output\": \"\"}, \"headers\": {}, \"error\": \"$(echo "$output" | jq -Rsa . | sed 's/^"\(.*\)"$/\1/')\"}"
    EXIT_CODE=$status
    debug "do_cmd_az: command failed, LAST_RESPONSE=$LAST_RESPONSE"
    return $EXIT_CODE
  fi

  # Store successful response in global variable
  LAST_RESPONSE="$output"
  debug "do_cmd_az: LAST_RESPONSE=$LAST_RESPONSE"

  # Check expectations and return status
  is_cmd_output_expected "$expect" "$output"
  local ret=$?
  debug "do_cmd_az: is_cmd_output_expected returned $ret"

  return $ret
}

function do_cmd_mysqlsh {
  local request=$1
  local expect=$2

  debug "do_cmd_mysqlsh: request=$request"
  debug "do_cmd_mysqlsh: expect=$expect"

  # 使用 -r 保持原始字符串格式
  local cmd_str
  cmd_str=$(echo "$request" | yq -r '.args | join(" ")')

  debug "do_cmd_mysqlsh: mysqlsh $cmd_str"

  local output
  # 使用 eval 执行完整命令
  output=$(eval "mysqlsh $cmd_str" 2>&1)
  local status=$?

  debug "do_cmd_mysqlsh: output=$output"

  # If command failed, create an error response
  if [ $status -ne 0 ]; then
    LAST_RESPONSE="{\"status_code\": $status, \"json\": {\"output\": \"\"}, \"headers\": {}, \"error\": \"$(echo "$output" | jq -Rsa . | sed 's/^\"\(.*\)\"$/\1/')\"}"
    EXIT_CODE=$status
    debug "do_cmd_mysqlsh: command failed, LAST_RESPONSE=$LAST_RESPONSE"
    return $EXIT_CODE
  fi

  LAST_RESPONSE="$output"
  debug "do_cmd_mysqlsh: LAST_RESPONSE=$LAST_RESPONSE"

  is_cmd_output_expected "$expect" "$output"
  local ret=$?
  debug "do_cmd_mysqlsh: is_cmd_output_expected returned $ret"

  return $ret
}

AUTH_TYPE="digest"
API_KEY=""
INPUT_FILE=""
PROFILE_AWS=""
OPENAPI_HOST="" #https://dedicated.staging.tidbapi.com/v1beta1

if [[ $# == 0 ]]; then
  echo $USAGE
  exit 1
fi

# check args number and required commands.
while getopts ":hi:a:t:p:H:" opt; do
  debug "Parsing option: $opt = $OPTARG"
  case "${opt}" in
  h)
    echo "$USAGE"
    exit 0
    ;;
  i)
    INPUT_FILE="${OPTARG}"
    ;;
  p)
    PROFILE_AWS="${OPTARG}"
    ;;
  t)
    API_KEY="${OPTARG}"
    ;;
  H)
    OPENAPI_HOST="${OPTARG}"
    ;;
  a)
    AUTH_TYPE="${OPTARG}"
    ;;
  *)
    echo "$USAGE" >&2
    exit 1
    ;;
  esac
done

if [[ -z "$INPUT_FILE" ]]; then
  echo "ERROR: -i <input_file> is required." >&2
  LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"error\": \"input file required\"}}"
  EXIT_CODE=1
fi

if ! command_exists "yq"; then
  echo "ERROR: 'yq' is required but not found." >&2
  LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"error\": \"yq not found\"}}"
  EXIT_CODE=1
fi

if ! py_module_exists "celpy"; then
  echo "ERROR: python module 'cel-python' is required but not found." >&2
  LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"error\": \"cel-python not found\"}}"
  EXIT_CODE=1
fi

debug "Reading request data from $INPUT_FILE"
REQUEST_TYPE=$(yq ".request_type" ${INPUT_FILE})
REQUEST=$(yq ".request" ${INPUT_FILE})
MAX_RETRIES=$(yq ".max_retries" ${INPUT_FILE})
DELAY_AFTER=$(yq ".delay_after" ${INPUT_FILE})
EXPECT=$(yq ".expect" ${INPUT_FILE})
RETRY_COUNT=0

debug "REQUEST_TYPE=$REQUEST_TYPE"
debug "REQUEST=$REQUEST"
debug "MAX_RETRIES=$MAX_RETRIES"
debug "DELAY_AFTER=$DELAY_AFTER"
debug "EXPECT=$EXPECT"

[[ ${MAX_RETRIES} -eq 0 ]] && MAX_RETRIES=1

case "$REQUEST_TYPE" in
"http")
  if [[ -z "$API_KEY" ]]; then
    echo "ERROR: -t <api_key> is required for http." >&2
    LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"error\": \"API key required\"}}"
    EXIT_CODE=1
  elif ! command_exists "curl"; then
    echo "ERROR: 'curl' is required but not found." >&2
    LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"error\": \"curl not found\"}}"
    EXIT_CODE=1
  else
    debug "Starting HTTP request loop"
    while [[ ${RETRY_COUNT} -lt ${MAX_RETRIES} ]]; do
      debug "Attempt $((RETRY_COUNT + 1))/$MAX_RETRIES"
      if do_curl "$OPENAPI_HOST" "$AUTH_TYPE" "$API_KEY" "$REQUEST" "$EXPECT"; then
        debug "Request succeeded"
        EXIT_CODE=0
        break
      fi

      debug "Request failed"
      ((RETRY_COUNT += 1))
      debug "Retry count is now $RETRY_COUNT"

      if [[ ${RETRY_COUNT} -eq ${MAX_RETRIES} ]]; then
        debug "All retries failed"
        EXIT_CODE=1
        break
      fi

      debug "Sleeping for $DELAY_AFTER seconds"
      sleep ${DELAY_AFTER}
    done
  fi
  ;;
"cmd_aws")
  if ! command_exists "aws"; then
    echo "ERROR: 'aws' is required but not found." >&2
    LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"error\": \"aws not found\"}}"
    EXIT_CODE=1
  else
    debug "Starting AWS command loop"
    while [[ ${RETRY_COUNT} -lt ${MAX_RETRIES} ]]; do
      debug "Attempt $((RETRY_COUNT + 1))/$MAX_RETRIES"
      if do_cmd_aws "$REQUEST" "$PROFILE_AWS" "$EXPECT"; then
        debug "Command succeeded"
        EXIT_CODE=0
        break
      fi

      debug "Command failed"
      ((RETRY_COUNT += 1))
      debug "Retry count is now $RETRY_COUNT"

      if [[ ${RETRY_COUNT} -eq ${MAX_RETRIES} ]]; then
        debug "All retries failed"
        EXIT_CODE=1
        break
      fi

      debug "Sleeping for $DELAY_AFTER seconds"
      sleep ${DELAY_AFTER}
    done
  fi
  ;;
"cmd_gcp")
  if ! command_exists "gcloud"; then
    echo "ERROR: 'gcloud' is required but not found." >&2
    LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"error\": \"gcloud not found\"}}"
    EXIT_CODE=1
  else
    debug "Starting GCP command loop"
    while [[ ${RETRY_COUNT} -lt ${MAX_RETRIES} ]]; do
      debug "Attempt $((RETRY_COUNT + 1))/$MAX_RETRIES"
      if do_cmd_gcloud "$REQUEST" "$EXPECT"; then
        debug "Command succeeded"
        EXIT_CODE=0
        break
      fi

      debug "Command failed"
      ((RETRY_COUNT += 1))
      debug "Retry count is now $RETRY_COUNT"

      if [[ ${RETRY_COUNT} -eq ${MAX_RETRIES} ]]; then
        debug "All retries failed"
        EXIT_CODE=1
        break
      fi

      debug "Sleeping for $DELAY_AFTER seconds"
      sleep ${DELAY_AFTER}
    done
  fi
  ;;
"cmd_azure")
  if ! command_exists "az"; then
    echo "ERROR: 'az' is required but not found." >&2
    LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"error\": \"az not found\"}}"
    EXIT_CODE=1
  else
    debug "Starting Azure command loop"
    while [[ ${RETRY_COUNT} -lt ${MAX_RETRIES} ]]; do
      debug "Attempt $((RETRY_COUNT + 1))/$MAX_RETRIES"
      if do_cmd_az "$REQUEST" "$EXPECT"; then
        debug "Command succeeded"
        EXIT_CODE=0
        break
      fi

      debug "Command failed"
      ((RETRY_COUNT += 1))
      debug "Retry count is now $RETRY_COUNT"

      if [[ ${RETRY_COUNT} -eq ${MAX_RETRIES} ]]; then
        debug "All retries failed"
        EXIT_CODE=1
        break
      fi

      debug "Sleeping for $DELAY_AFTER seconds"
      sleep ${DELAY_AFTER}
    done
  fi
  ;;
"cmd_mysqlsh")
  if ! command_exists "mysqlsh"; then
    echo "ERROR: 'mysqlsh' is required but not found." >&2
    LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"error\": \"mysqlsh not found\"}}"
    EXIT_CODE=1
  else
    debug "Starting MySQL Shell command loop"
    while [[ ${RETRY_COUNT} -lt ${MAX_RETRIES} ]]; do
      debug "Attempt $((RETRY_COUNT + 1))/$MAX_RETRIES"
      if do_cmd_mysqlsh "$REQUEST" "$EXPECT"; then
        debug "Command succeeded"
        EXIT_CODE=0
        break
      fi

      debug "Command failed"
      ((RETRY_COUNT += 1))
      debug "Retry count is now $RETRY_COUNT"

      if [[ ${RETRY_COUNT} -eq ${MAX_RETRIES} ]]; then
        debug "All retries failed"
        EXIT_CODE=1
        break
      fi

      debug "Sleeping for $DELAY_AFTER seconds"
      sleep ${DELAY_AFTER}
    done
  fi
  ;;
*)
  echo "ERROR: $REQUEST_TYPE is not supported." >&2
  LAST_RESPONSE="{\"status_code\": 500, \"headers\": {}, \"body\": {\"error\": \"unsupported request type\"}}"
  EXIT_CODE=1
  ;;
esac

# Always output response and exit with proper code
debug "Final response: $LAST_RESPONSE"
echo "$LAST_RESPONSE"
exit $EXIT_CODE