#!/usr/bin/env bash

set -euo pipefail

MAX_ROWS_PER_QUERY=10000
SESSION_DIR=".session"

decode_binary_rows() {
  local file="$1"
  python3 - "$file" <<'PY'
import json
import string
import sys

path = sys.argv[1]
with open(path, "r", encoding="utf-8") as f:
    payload = json.load(f)

types = payload.get("types")
rows = payload.get("rows")

if not isinstance(types, list) or not isinstance(rows, list):
    json.dump(payload, sys.stdout, ensure_ascii=False)
    sys.exit(0)

binary_types = {
    "BLOB",
    "TINYBLOB",
    "MEDIUMBLOB",
    "LONGBLOB",
    "BINARY",
    "VARBINARY",
    "BIT",
}
hexdigits = set(string.hexdigits)


def to_bytes(value):
    if isinstance(value, list):
        if all(isinstance(x, int) and 0 <= x <= 255 for x in value):
            return bytes(value)
        return None
    if isinstance(value, str):
        text = value.strip()
        if text.startswith(("0x", "0X")):
            text = text[2:]
        if len(text) % 2 != 0:
            return None
        if not text:
            return b""
        if all(ch in hexdigits for ch in text):
            try:
                return bytes.fromhex(text)
            except ValueError:
                return None
    return None


def to_printable_utf8_or_hex(blob):
    if blob == b"":
        return ""
    try:
        text = blob.decode("utf-8")
        if all((ch >= " " or ch in "\t\r\n") for ch in text):
            return text
    except UnicodeDecodeError:
        pass
    return "0x" + blob.hex()


def decode_value(value, type_name):
    normalized = (type_name or "").upper()
    if normalized not in binary_types:
        return value

    blob = to_bytes(value)
    if blob is None:
        return value

    if normalized == "BIT":
        if len(blob) == 0:
            return "0b0"
        return "0b" + "".join(f"{b:08b}" for b in blob)

    return to_printable_utf8_or_hex(blob)


decoded_rows = []
for row in rows:
    if not isinstance(row, list):
        decoded_rows.append(row)
        continue

    out_row = []
    for i, value in enumerate(row):
        col_type = ""
        if i < len(types) and isinstance(types[i], dict):
            col_type = types[i].get("type", "")
        out_row.append(decode_value(value, col_type))
    decoded_rows.append(out_row)

payload["rows"] = decoded_rows
json.dump(payload, sys.stdout, ensure_ascii=False)
PY
}

load_database_url() {
  if [[ -n "${DATABASE_URL:-}" ]]; then
    return 0
  fi

  if [[ -f ".env" ]]; then
    while IFS= read -r line; do
      line="${line%%$'\r'}"
      case "$line" in
        ''|'#'*)
          ;;
        DATABASE_URL=*)
          DATABASE_URL="${line#DATABASE_URL=}"
          DATABASE_URL="${DATABASE_URL%\"}"
          DATABASE_URL="${DATABASE_URL#\"}"
          DATABASE_URL="${DATABASE_URL%\'}"
          DATABASE_URL="${DATABASE_URL#\'}"
          export DATABASE_URL
          return 0
          ;;
      esac
    done < ".env"
  fi

  return 1
}

setup_message() {
  cat <<'EOF'
DATABASE_URL not found. Provide it in environment or .env.
Expected format:
  mysql://<username>:<password>@<host>/<database>
Example:
  mysql://4FriSs61wM5o1zu.root:ZwuR6cdSNdjpxEY8@gateway01.us-east-1.dev.shared.aws.tidbcloud.com/test
Setup options:
  1) export DATABASE_URL="mysql://<username>:<password>@<host>/<database>"
  2) create .env with: DATABASE_URL="mysql://<username>:<password>@<host>/<database>"
EOF
}

parse_database_url() {
  if [[ "${DATABASE_URL:-}" != mysql://* ]]; then
    printf 'invalid DATABASE_URL: scheme must be mysql\n' >&2
    return 1
  fi

  local rest creds hostdb
  rest="${DATABASE_URL#mysql://}"
  hostdb="${rest##*@}"
  creds="${rest%@$hostdb}"

  DB_USERNAME="${creds%%:*}"
  DB_PASSWORD="${creds#*:}"
  DB_HOST="${hostdb%%/*}"
  if [[ "$hostdb" == */* ]]; then
    DB_DATABASE="${hostdb#*/}"
  else
    DB_DATABASE=""
  fi

  if [[ "$DB_USERNAME" == "$creds" || -z "$DB_USERNAME" ]]; then
    printf 'invalid DATABASE_URL: username is required\n' >&2
    return 1
  fi
  if [[ "$DB_PASSWORD" == "$creds" ]]; then
    printf 'invalid DATABASE_URL: password is required\n' >&2
    return 1
  fi
  if [[ -z "$DB_HOST" ]]; then
    printf 'invalid DATABASE_URL: host is required\n' >&2
    return 1
  fi
}

json_error() {
  local msg="$1"
  jq -n --arg m "$msg" '{ok:false,status:0,session:"",error:{message:$m}}'
}

read_session() {
  local file="$1"
  if [[ -f "$file" ]]; then
    cat "$file"
  else
    printf ''
  fi
}

write_session() {
  local file="$1" session="$2"
  mkdir -p "$(dirname "$file")"
  printf '%s' "$session" > "$file"
}

normalize_session_file() {
  local file="$1"
  local base

  base="${file##*/}"
  if [[ -z "$base" ]]; then
    base="session.token"
  fi

  mkdir -p "$SESSION_DIR"
  printf './%s/%s' "$SESSION_DIR" "$base"
}

read_response_header() {
  local file="$1" header_name="$2"
  local line name value target

  target="$(printf '%s' "$header_name" | tr '[:upper:]' '[:lower:]')"

  while IFS= read -r line; do
    line="${line%%$'\r'}"
    name="${line%%:*}"
    name="$(printf '%s' "$name" | tr '[:upper:]' '[:lower:]')"
    if [[ "$name" == "$target" ]]; then
      value="${line#*:}"
      value="${value#"${value%%[![:space:]]*}"}"
      value="${value%"${value##*[![:space:]]}"}"
      printf '%s' "$value"
      return 0
    fi
  done < "$file"

  printf ''
}

request_id() {
  printf '%s%s' "$(date -u +%Y%m%d%H%M%S)" "$(LC_ALL=C tr -dc 'a-z0-9' < /dev/urandom | head -c 20)"
}

post_sql() {
  local sql="$1" isolation="$2" stateful_action="$3" session_file="$4" debug="$5" insecure="$6" cleanup_on_success="$7" retry_on_invalid_session="$8"
  local session auth url rid body

  session="$(read_session "$session_file")"
  rid="$(request_id)"
  auth="$(printf '%s' "${DB_USERNAME}:${DB_PASSWORD}" | base64)"
  url="https://http-${DB_HOST}/v1beta/sql"
  body="$(jq -n --arg q "$sql" '{query:$q}')"

  local headers_file body_file err_file
  headers_file="$(mktemp)"
  body_file="$(mktemp)"
  err_file="$(mktemp)"

  local -a curl_args
  curl_args=(-sS --compressed -X POST "$url" -D "$headers_file" -o "$body_file" -w "%{http_code}" \
    -H "Content-Type: application/json" \
    -H "User-Agent: tidbcloud-http/0.3" \
    -H "Authorization: Basic ${auth}" \
    -H "TiDB-Database: ${DB_DATABASE_OVERRIDE:-$DB_DATABASE}" \
    -H "TiDB-Session: ${session}" \
    -H "X-Debug-Trace-Id: ${rid}" \
    -H "Accept-Encoding: gzip" \
    --data "$body")

  if [[ -n "$isolation" ]]; then
    curl_args+=(-H "TiDB-Isolation-Level: ${isolation}")
  fi
  if [[ -n "$stateful_action" ]]; then
    curl_args+=(-H "TiDB-Stateful-Action: ${stateful_action}")
  fi
  if [[ "$insecure" == "1" ]]; then
    curl_args+=(-k)
  fi

  if [[ "$debug" == "1" ]]; then
    jq -n --arg u "$url" --arg r "$rid" --arg s "$stateful_action" --arg i "$isolation" '{debug:{url:$u,request_id:$r,stateful_action:$s,isolation:$i}}' >&2
  fi

  local http_code
  if ! http_code="$(curl "${curl_args[@]}" 2>"$err_file")"; then
    local emsg
    emsg="$(tr '\n' ' ' < "$err_file")"
    if [[ "$emsg" == *"Could not resolve host"* || "$emsg" == *"Name or service not known"* ]]; then
      emsg="$emsg; host not found. This direct HTTP flow only supports TiDB Cloud Starter and Essential plans. Dedicated/Premium are not supported by this skill."
    fi
    json_error "$emsg"
    rm -f "$headers_file" "$body_file" "$err_file"
    return 1
  fi

  local next_session result_session
  next_session="$(read_response_header "$headers_file" "TiDB-Session")"
  result_session="$next_session"

  if [[ "$stateful_action" == "close" ]]; then
    rm -f "$session_file"
    result_session=""
  elif [[ -n "$next_session" ]]; then
    write_session "$session_file" "$next_session"
  elif [[ -n "$session" ]]; then
    result_session="$session"
  fi

  if [[ "$http_code" -ge 200 && "$http_code" -lt 300 ]]; then
    if jq -e . < "$body_file" >/dev/null 2>&1; then
      local row_count
      row_count="$(jq -r 'if (type=="object" and has("rows") and (.rows|type)=="array") then (.rows|length) else 0 end' "$body_file" 2>/dev/null || printf '0')"
      if [[ "$row_count" -gt "$MAX_ROWS_PER_QUERY" ]]; then
        jq -n --argjson c "$http_code" --arg s "$result_session" --argjson r "$row_count" --argjson m "$MAX_ROWS_PER_QUERY" '{ok:false,status:$c,session:$s,error:{message:("query returned too many rows: " + ($r|tostring) + ", max allowed is " + ($m|tostring) + ". add LIMIT to your SQL")}}'
        rm -f "$headers_file" "$body_file" "$err_file"
        return 1
      fi
      local payload_json should_decode
      payload_json="$(cat "$body_file")"
      should_decode="0"
      if jq -e '((.types | type) == "array") and any(.types[]?; ((.type // "") | ascii_upcase | test("(^| )((TINY|MEDIUM|LONG)?BLOB|BINARY|VARBINARY|BIT)( |$)")))' "$body_file" >/dev/null 2>&1; then
        should_decode="1"
      fi
      if [[ "$should_decode" == "1" ]]; then
        local decoded_json
        if decoded_json="$(decode_binary_rows "$body_file" 2>/dev/null)"; then
          payload_json="$decoded_json"
        fi
      fi
      jq -n --argjson d "$payload_json" --arg s "$result_session" --argjson c "$http_code" '{ok:true,status:$c,session:$s,data:$d}'
    else
      jq -n --arg t "$(cat "$body_file")" --arg s "$result_session" --argjson c "$http_code" '{ok:true,status:$c,session:$s,text:$t}'
    fi
    if [[ "$cleanup_on_success" == "1" ]]; then
      rm -f "$session_file"
    fi
    rm -f "$headers_file" "$body_file" "$err_file"
  else
    local invalid_session="0"
    if jq -e '((.code? == 49900001) or (((.message? // "") | tostring | ascii_downcase) | test("invalid session")))' < "$body_file" >/dev/null 2>&1; then
      invalid_session="1"
    fi

    if [[ "$invalid_session" == "1" && "$retry_on_invalid_session" == "1" && "$cleanup_on_success" != "1" ]]; then
      rm -f "$session_file"
      if [[ "$debug" == "1" ]]; then
        jq -n --arg sf "$session_file" '{debug:{action:"invalid-session-retry",session_file:$sf}}' >&2
      fi
      rm -f "$headers_file" "$body_file" "$err_file"
      post_sql "$sql" "$isolation" "$stateful_action" "$session_file" "$debug" "$insecure" "$cleanup_on_success" "0"
      return $?
    fi

    if jq -e . < "$body_file" >/dev/null 2>&1; then
      jq -n --argjson e "$(cat "$body_file")" --argjson c "$http_code" --arg s "$result_session" '{ok:false,status:$c,session:$s,error:$e}'
    else
      jq -n --arg t "$(cat "$body_file")" --argjson c "$http_code" --arg s "$result_session" '{ok:false,status:$c,session:$s,error:{message:$t}}'
    fi
    rm -f "$headers_file" "$body_file" "$err_file"
    return 1
  fi
}

usage() {
  cat <<'EOF'
Usage:
  tidb_serverless_http.sh stateless --sql "select 1"
  tidb_serverless_http.sh stateful open
  tidb_serverless_http.sh stateful query --sql "show tables"
  tidb_serverless_http.sh stateful close
  tidb_serverless_http.sh tx begin [--isolation "READ COMMITTED|REPEATABLE READ"]
  tidb_serverless_http.sh tx query --sql "select 1"
  tidb_serverless_http.sh tx commit
  tidb_serverless_http.sh tx rollback

Common options:
  --database <db>
  --session-file <name>   (stored in .session/<name>)
  --debug
  --secure   (enable TLS verify; default is -k)

Default session files (.session directory):
  stateless: .session/stateless.token
  stateful:  .session/stateful.token  (removed on stateful close)
  tx:        .session/tx.token        (removed on tx commit/rollback)
EOF
}

main() {
  if ! command -v jq >/dev/null 2>&1; then
    printf 'jq is required\n' >&2
    exit 2
  fi

  if ! load_database_url; then
    json_error "$(setup_message)"
    exit 1
  fi

  if ! parse_database_url; then
    json_error "invalid DATABASE_URL. Expected mysql://<username>:<password>@<host>/<database>"
    exit 1
  fi

  if [[ $# -lt 1 ]]; then
    usage
    exit 2
  fi

  local mode action sql isolation
  local session_file=""
  local debug="0"
  local insecure="1"
  local cleanup_on_success="0"
  DB_DATABASE_OVERRIDE=""

  mode="$1"
  shift

  case "$mode" in
    stateless)
      session_file="stateless.token"
      sql=""
      while [[ $# -gt 0 ]]; do
        case "$1" in
          --sql) sql="$2"; shift 2 ;;
          --database) DB_DATABASE_OVERRIDE="$2"; shift 2 ;;
          --session-file) session_file="$2"; shift 2 ;;
          --debug) debug="1"; shift ;;
          --secure) insecure="0"; shift ;;
          *) printf 'unknown arg: %s\n' "$1" >&2; exit 2 ;;
        esac
      done
      session_file="$(normalize_session_file "$session_file")"
      if [[ -z "$sql" ]]; then
        printf 'missing --sql\n' >&2
        exit 2
      fi
      post_sql "$sql" "" "" "$session_file" "$debug" "$insecure" "$cleanup_on_success" "1"
      ;;
    stateful)
      session_file="stateful.token"
      if [[ $# -lt 1 ]]; then
        usage
        exit 2
      fi
      action="$1"
      shift
      sql=""
      while [[ $# -gt 0 ]]; do
        case "$1" in
          --sql) sql="$2"; shift 2 ;;
          --database) DB_DATABASE_OVERRIDE="$2"; shift 2 ;;
          --session-file) session_file="$2"; shift 2 ;;
          --debug) debug="1"; shift ;;
          --secure) insecure="0"; shift ;;
          *) printf 'unknown arg: %s\n' "$1" >&2; exit 2 ;;
        esac
      done
      session_file="$(normalize_session_file "$session_file")"
      case "$action" in
        open)
          cleanup_on_success="0"
          post_sql "" "" "open" "$session_file" "$debug" "$insecure" "$cleanup_on_success" "1"
          ;;
        query)
          if [[ -z "$sql" ]]; then
            printf 'missing --sql\n' >&2
            exit 2
          fi
          cleanup_on_success="0"
          post_sql "$sql" "" "" "$session_file" "$debug" "$insecure" "$cleanup_on_success" "1"
          ;;
        close)
          cleanup_on_success="1"
          post_sql "" "" "close" "$session_file" "$debug" "$insecure" "$cleanup_on_success" "0"
          ;;
        *) printf 'invalid stateful action: %s\n' "$action" >&2; exit 2 ;;
      esac
      ;;
    tx)
      session_file="tx.token"
      if [[ $# -lt 1 ]]; then
        usage
        exit 2
      fi
      action="$1"
      shift
      sql=""
      isolation=""
      while [[ $# -gt 0 ]]; do
        case "$1" in
          --sql) sql="$2"; shift 2 ;;
          --isolation) isolation="$2"; shift 2 ;;
          --database) DB_DATABASE_OVERRIDE="$2"; shift 2 ;;
          --session-file) session_file="$2"; shift 2 ;;
          --debug) debug="1"; shift ;;
          --secure) insecure="0"; shift ;;
          *) printf 'unknown arg: %s\n' "$1" >&2; exit 2 ;;
        esac
      done
      session_file="$(normalize_session_file "$session_file")"
      case "$action" in
        begin)
          cleanup_on_success="0"
          post_sql "BEGIN" "$isolation" "" "$session_file" "$debug" "$insecure" "$cleanup_on_success" "1"
          ;;
        query)
          if [[ -z "$sql" ]]; then
            printf 'missing --sql\n' >&2
            exit 2
          fi
          cleanup_on_success="0"
          post_sql "$sql" "" "" "$session_file" "$debug" "$insecure" "$cleanup_on_success" "1"
          ;;
        commit)
          cleanup_on_success="1"
          post_sql "COMMIT" "" "" "$session_file" "$debug" "$insecure" "$cleanup_on_success" "0"
          ;;
        rollback)
          cleanup_on_success="1"
          post_sql "ROLLBACK" "" "" "$session_file" "$debug" "$insecure" "$cleanup_on_success" "0"
          ;;
        *) printf 'invalid tx action: %s\n' "$action" >&2; exit 2 ;;
      esac
      ;;
    *)
      usage
      exit 2
      ;;
  esac
}

main "$@"