#!/usr/bin/env bash
# tidb-branch.sh — Create, query, and delete TiDB Cloud serverless branches.
#
# Environment variables (required):
#   TIDB_CLOUD_PUBLIC_KEY   — API public key
#   TIDB_CLOUD_PRIVATE_KEY  — API private key
#   TIDB_CLOUD_CLUSTER_ID   — Target cluster ID
#
# Usage:
#   ./scripts/tidb-branch.sh create <display-name> [--parent <parent-branch-id>]
#     → creates branch (optional parent defaults to cluster root)
#   ./scripts/tidb-branch.sh delete <branch-id>       → deletes branch
#   ./scripts/tidb-branch.sh list                      → lists all branches
#   ./scripts/tidb-branch.sh query <branch-id> <sql>   → runs SQL query (requires Go)
#   ./scripts/tidb-branch.sh describe <branch-id>      → shows branch details
#   ./scripts/tidb-branch.sh reset <branch-id>         → syncs branch with parent
#   ./scripts/tidb-branch.sh set-password <branch-id> <password>  → changes root password

set -euo pipefail

BASE_URL="https://serverless.tidbapi.com/v1beta1"
QUERY_TOOL_DIR="/tmp/tidb_branch_query_tool"

die() { echo "ERROR: $*" >&2; exit 1; }

[ -z "${TIDB_CLOUD_PUBLIC_KEY:-}" ] && die "TIDB_CLOUD_PUBLIC_KEY not set"
[ -z "${TIDB_CLOUD_PRIVATE_KEY:-}" ] && die "TIDB_CLOUD_PRIVATE_KEY not set"
[ -z "${TIDB_CLOUD_CLUSTER_ID:-}" ] && die "TIDB_CLOUD_CLUSTER_ID not set"

DIGEST_AUTH="${TIDB_CLOUD_PUBLIC_KEY}:${TIDB_CLOUD_PRIVATE_KEY}"
CLUSTER_URL="${BASE_URL}/clusters/${TIDB_CLOUD_CLUSTER_ID}"

api() {
  local method="$1" path="$2"
  shift 2
  curl -s --digest -u "$DIGEST_AUTH" -X "$method" "${CLUSTER_URL}${path}" \
    -H "Content-Type: application/json" "$@"
}

generate_password() {
  # 24 chars, alphanumeric + special, guaranteed mixed case
  openssl rand -base64 32 | tr -dc 'a-zA-Z0-9!@#$%' | head -c 24
}

cmd_create() {
  local name=""
  local parent_id=""
  while [ $# -gt 0 ]; do
    case "$1" in
      --parent)
        [ -n "${2:-}" ] || die "create --parent requires a branch ID (e.g. bran-...)"
        parent_id="$2"
        shift 2
        ;;
      *)
        [ -z "$name" ] || die "Unexpected argument: $1 (usage: create <display-name> [--parent <parent-branch-id>])"
        name="$1"
        shift
        ;;
    esac
  done
  [ -n "$name" ] || die "Usage: create <display-name> [--parent <parent-branch-id>]"

  local password
  password=$(generate_password)

  local body
  if [ -n "$parent_id" ]; then
    echo "Creating branch '${name}' from parent '${parent_id}'..."
    body=$(printf '{"displayName": %s, "rootPassword": %s, "parentId": %s}' \
      "$(python3 -c "import json,sys; print(json.dumps(sys.argv[1]))" "$name")" \
      "$(python3 -c "import json,sys; print(json.dumps(sys.argv[1]))" "$password")" \
      "$(python3 -c "import json,sys; print(json.dumps(sys.argv[1]))" "$parent_id")")
  else
    echo "Creating branch '${name}'..."
    body=$(printf '{"displayName": %s, "rootPassword": %s}' \
      "$(python3 -c "import json,sys; print(json.dumps(sys.argv[1]))" "$name")" \
      "$(python3 -c "import json,sys; print(json.dumps(sys.argv[1]))" "$password")")
  fi

  local result
  result=$(api POST "/branches" -d "$body")

  local branch_id state
  branch_id=$(echo "$result" | python3 -c "import sys,json; print(json.load(sys.stdin).get('branchId',''))")
  state=$(echo "$result" | python3 -c "import sys,json; print(json.load(sys.stdin).get('state',''))")

  [ -z "$branch_id" ] && die "Branch creation failed: $result"
  echo "Branch ID: ${branch_id} (state: ${state})"
  echo "Polling until ACTIVE (up to 90s)..."

  local host="" port="" user_prefix=""
  for i in $(seq 1 30); do
    result=$(api GET "/branches/${branch_id}")
    state=$(echo "$result" | python3 -c "import sys,json; print(json.load(sys.stdin).get('state',''))")
    if [ "$state" = "ACTIVE" ]; then
      host=$(echo "$result" | python3 -c "import sys,json; print(json.load(sys.stdin)['endpoints']['public']['host'])")
      port=$(echo "$result" | python3 -c "import sys,json; print(json.load(sys.stdin)['endpoints']['public']['port'])")
      user_prefix=$(echo "$result" | python3 -c "import sys,json; print(json.load(sys.stdin)['userPrefix'])")
      break
    fi
    sleep 3
  done

  [ -z "$host" ] && die "Branch did not become ACTIVE within 90s (last state: ${state})"

  echo ""
  echo "=== Branch Ready ==="
  echo "BRANCH_ID=${branch_id}"
  echo "TIDB_HOST=${host}"
  echo "TIDB_PORT=${port}"
  echo "TIDB_USER=${user_prefix}.root"
  echo "TIDB_PASSWORD=${password}"
  echo "TIDB_DATABASE=nexuscrm"
  echo ""
  echo "# Export for use with query command:"
  echo "export BRANCH_ID=${branch_id}"
  echo "export TIDB_HOST=${host}"
  echo "export TIDB_PORT=${port}"
  echo "export TIDB_USER=${user_prefix}.root"
  echo "export TIDB_PASSWORD=${password}"
  echo "export TIDB_DATABASE=nexuscrm"
}

cmd_delete() {
  local branch_id="${1:?Usage: delete <branch-id>}"
  echo "Deleting branch ${branch_id}..."
  local result
  result=$(api DELETE "/branches/${branch_id}")
  local state
  state=$(echo "$result" | python3 -c "import sys,json; print(json.load(sys.stdin).get('state',''))" 2>/dev/null || echo "unknown")
  echo "Branch ${branch_id}: ${state}"
}

cmd_list() {
  local result
  result=$(api GET "/branches?pageSize=100")
  echo "$result" | python3 -c "
import sys, json
data = json.load(sys.stdin)
branches = data.get('branches', [])
if not branches:
    print('No branches found.')
else:
    print(f'{'ID':<48} {'Name':<30} {'State':<10}')
    print('-' * 88)
    for b in branches:
        print(f\"{b['branchId']:<48} {b.get('displayName',''):<30} {b.get('state',''):<10}\")
    print(f'\nTotal: {len(branches)} branches')
"
}

cmd_describe() {
  local branch_id="${1:?Usage: describe <branch-id>}"
  api GET "/branches/${branch_id}" | python3 -m json.tool
}

cmd_reset() {
  local branch_id="${1:?Usage: reset <branch-id>}"
  echo "Resetting branch ${branch_id} to parent state..."
  api POST "/branches/${branch_id}:reset"
  echo ""
  echo "Reset initiated. Branch will sync with parent data."
  echo "Note: After reset, you must set a new password via: $0 set-password ${branch_id} <password>"
}

cmd_set_password() {
  local branch_id="${1:?Usage: set-password <branch-id> <password>}"
  local password="${2:?Usage: set-password <branch-id> <password>}"
  api PATCH "/branches/${branch_id}/sqlUsers/root" -d "{\"password\": \"${password}\"}"
  echo "Password updated for branch ${branch_id}"
}

cmd_query() {
  local branch_id="${1:?Usage: query <branch-id> <sql> (requires TIDB_HOST/USER/PASSWORD env vars)}"
  local sql="${2:?Usage: query <branch-id> <sql>}"

  [ -z "${TIDB_HOST:-}" ] && die "TIDB_HOST not set. Export the env vars from 'create' output first."
  [ -z "${TIDB_USER:-}" ] && die "TIDB_USER not set."
  [ -z "${TIDB_PASSWORD:-}" ] && die "TIDB_PASSWORD not set."

  # Build query tool if needed
  if [ ! -f "${QUERY_TOOL_DIR}/tidb_query" ]; then
    echo "Building query tool..." >&2
    mkdir -p "${QUERY_TOOL_DIR}"
    cat > "${QUERY_TOOL_DIR}/main.go" << 'GOEOF'
package main

import (
	"crypto/tls"
	"database/sql"
	"fmt"
	"os"

	"github.com/go-sql-driver/mysql"
)

func main() {
	host := os.Getenv("TIDB_HOST")
	port := os.Getenv("TIDB_PORT")
	user := os.Getenv("TIDB_USER")
	pass := os.Getenv("TIDB_PASSWORD")
	dbName := os.Getenv("TIDB_DATABASE")
	if port == "" { port = "4000" }
	if dbName == "" { dbName = "nexuscrm" }

	mysql.RegisterTLSConfig("tidb", &tls.Config{
		MinVersion: tls.VersionTLS12,
		ServerName: host,
	})

	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=UTC&tls=tidb",
		user, pass, host, port, dbName)

	db, err := sql.Open("mysql", dsn)
	if err != nil { fmt.Fprintf(os.Stderr, "open: %v\n", err); os.Exit(1) }
	defer db.Close()

	rows, err := db.Query(os.Args[1])
	if err != nil { fmt.Fprintf(os.Stderr, "query: %v\n", err); os.Exit(1) }
	defer rows.Close()

	cols, _ := rows.Columns()
	for i, c := range cols {
		if i > 0 { fmt.Print("\t") }
		fmt.Print(c)
	}
	fmt.Println()

	vals := make([]interface{}, len(cols))
	ptrs := make([]interface{}, len(cols))
	for i := range vals { ptrs[i] = &vals[i] }
	for rows.Next() {
		rows.Scan(ptrs...)
		for i, v := range vals {
			if i > 0 { fmt.Print("\t") }
			switch val := v.(type) {
			case []byte: fmt.Print(string(val))
			case nil:    fmt.Print("NULL")
			default:     fmt.Print(val)
			}
		}
		fmt.Println()
	}
}
GOEOF
    cd "${QUERY_TOOL_DIR}" && go mod init tidb_query 2>/dev/null && go get github.com/go-sql-driver/mysql 2>/dev/null && go build -o tidb_query . 2>&1 >&2
    cd - >/dev/null
  fi

  "${QUERY_TOOL_DIR}/tidb_query" "$sql"
}

case "${1:-help}" in
  create)       shift; cmd_create "$@" ;;
  delete)       shift; cmd_delete "$@" ;;
  list)         shift; cmd_list "$@" ;;
  describe)     shift; cmd_describe "$@" ;;
  reset)        shift; cmd_reset "$@" ;;
  set-password) shift; cmd_set_password "$@" ;;
  query)        shift; cmd_query "$@" ;;
  help|*)
    echo "Usage: $0 <command> [args]"
    echo ""
    echo "Commands:"
    echo "  create <name> [--parent <id>]     Create branch (optional parent branch ID)"
    echo "  delete <branch-id>                Delete branch"
    echo "  list                              List all branches"
    echo "  describe <branch-id>              Show branch details"
    echo "  reset <branch-id>                 Sync branch with parent"
    echo "  set-password <branch-id> <pass>   Change root password"
    echo "  query <branch-id> <sql>           Run SQL (needs TIDB_* env vars)"
    echo ""
    echo "Required env vars: TIDB_CLOUD_PUBLIC_KEY, TIDB_CLOUD_PRIVATE_KEY, TIDB_CLOUD_CLUSTER_ID"
    ;;
esac