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

CLIENT_HOST="${CLIENT_HOST:-k_client}"
CA_FILE="${CA_FILE:-/home/user/chromecard/tls/phase2/ca.crt}"
PROXY_URL="${PROXY_URL:-https://127.0.0.1:9771}"
USERNAME="${USERNAME:-alice}"
REQUESTS="${REQUESTS:-20}"
PARALLELISM="${PARALLELISM:-8}"
CONNECT_TIMEOUT="${CONNECT_TIMEOUT:-8}"
LOGIN_TIMEOUT="${LOGIN_TIMEOUT:-90}"
INTERACTIVE_CARD="${INTERACTIVE_CARD:-0}"
EXPECT_AUTH_MODE="${EXPECT_AUTH_MODE:-}"
SSH_CONFIG="${SSH_CONFIG:-/home/user/.ssh/config}"

usage() {
  cat <<'EOF'
Usage: phase5_chain_regression.sh [options]

Runs the Phase 5 split-VM regression from the host by executing the client-side
flow inside k_client over SSH.

Options:
  --client-host HOST       SSH host alias for k_client (default: k_client)
  --ca-file PATH           CA bundle path inside k_client
  --proxy-url URL          Proxy URL visible from k_client
  --username NAME          Username for session login
  --requests N             Number of counter requests to issue
  --parallelism N          Number of concurrent workers
  --connect-timeout SEC    SSH connect timeout
  --login-timeout SEC      Timeout for the interactive login request (default: 90)
  --interactive-card       Print card-confirmation instructions before login
  --expect-auth-mode NAME  Require login response auth_mode to match
  --ssh-config PATH        SSH config file to use (default: /home/user/.ssh/config)
  -h, --help               Show this help text
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --client-host)
      CLIENT_HOST="$2"
      shift 2
      ;;
    --ca-file)
      CA_FILE="$2"
      shift 2
      ;;
    --proxy-url)
      PROXY_URL="$2"
      shift 2
      ;;
    --username)
      USERNAME="$2"
      shift 2
      ;;
    --requests)
      REQUESTS="$2"
      shift 2
      ;;
    --parallelism)
      PARALLELISM="$2"
      shift 2
      ;;
    --connect-timeout)
      CONNECT_TIMEOUT="$2"
      shift 2
      ;;
    --login-timeout)
      LOGIN_TIMEOUT="$2"
      shift 2
      ;;
    --interactive-card)
      INTERACTIVE_CARD=1
      shift
      ;;
    --expect-auth-mode)
      EXPECT_AUTH_MODE="$2"
      shift 2
      ;;
    --ssh-config)
      SSH_CONFIG="$2"
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "unknown argument: $1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

if [[ "${INTERACTIVE_CARD}" == "1" ]]; then
  cat <<EOF
Starting interactive login for ${USERNAME}.
When the card shows the authentication prompt, press yes to approve.
Press no only if you want to reject the login.
EOF
fi

ssh \
  -F "${SSH_CONFIG}" \
  -o BatchMode=yes \
  -o StrictHostKeyChecking=accept-new \
  -o ConnectTimeout="${CONNECT_TIMEOUT}" \
  "${CLIENT_HOST}" \
  env \
    CA_FILE="${CA_FILE}" \
    PROXY_URL="${PROXY_URL}" \
    USERNAME="${USERNAME}" \
    REQUESTS="${REQUESTS}" \
    PARALLELISM="${PARALLELISM}" \
    LOGIN_TIMEOUT="${LOGIN_TIMEOUT}" \
    EXPECT_AUTH_MODE="${EXPECT_AUTH_MODE}" \
    python3 - <<'PY'
import concurrent.futures
import json
import os
import ssl
import sys
import urllib.error
import urllib.request

ca_file = os.environ["CA_FILE"]
proxy_url = os.environ["PROXY_URL"].rstrip("/")
username = os.environ["USERNAME"]
requests = int(os.environ["REQUESTS"])
parallelism = int(os.environ["PARALLELISM"])
login_timeout = int(os.environ["LOGIN_TIMEOUT"])
expect_auth_mode = os.environ["EXPECT_AUTH_MODE"]

if requests < 1:
    raise SystemExit("REQUESTS must be >= 1")
if parallelism < 1:
    raise SystemExit("PARALLELISM must be >= 1")

ctx = ssl.create_default_context(cafile=ca_file)

def post_json(path: str, payload: dict | None = None, token: str | None = None, timeout: int = 10):
    data = None if payload is None else json.dumps(payload).encode("utf-8")
    headers = {}
    if payload is not None:
        headers["Content-Type"] = "application/json"
    if token:
        headers["Authorization"] = f"Bearer {token}"
    req = urllib.request.Request(
        f"{proxy_url}{path}",
        data=data,
        headers=headers,
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, context=ctx, timeout=timeout) as resp:
            return resp.status, json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8")
        try:
            payload = json.loads(body)
        except json.JSONDecodeError:
            payload = {"ok": False, "error": body}
        return exc.code, payload

status, login = post_json("/session/login", {"username": username}, timeout=login_timeout)
if status != 200 or "session_token" not in login:
    print(json.dumps({"ok": False, "stage": "login", "status": status, "response": login}))
    raise SystemExit(1)
if expect_auth_mode and login.get("auth_mode") != expect_auth_mode:
    print(
        json.dumps(
            {
                "ok": False,
                "stage": "login",
                "error": "unexpected auth_mode",
                "expected": expect_auth_mode,
                "response": login,
            }
        )
    )
    raise SystemExit(1)

token = login["session_token"]
values = []

def fetch_one(_: int) -> int:
    status, payload = post_json("/resource/counter", {}, token=token)
    if status != 200:
        raise RuntimeError(json.dumps({"status": status, "response": payload}))
    return int(payload["upstream"]["value"])

try:
    with concurrent.futures.ThreadPoolExecutor(max_workers=parallelism) as pool:
        for value in pool.map(fetch_one, range(requests)):
            values.append(value)

    status_resp, session = post_json("/session/status", {}, token=token)
    logout_status, logout = post_json("/session/logout", {}, token=token)
    invalid_status, invalid = post_json("/resource/counter", {}, token=token)
except Exception as exc:
    try:
        post_json("/session/logout", {}, token=token)
    finally:
        raise SystemExit(str(exc))

sorted_values = sorted(values)
expected = list(range(sorted_values[0], sorted_values[-1] + 1)) if sorted_values else []

summary = {
    "ok": True,
    "username": username,
    "proxy_url": proxy_url,
    "requests": requests,
    "parallelism": parallelism,
    "unique": len(set(values)) == len(values),
    "gap_free": sorted_values == expected,
    "min": min(sorted_values) if sorted_values else None,
    "max": max(sorted_values) if sorted_values else None,
    "values": sorted_values,
    "login": login,
    "session_status": {"status": status_resp, "response": session},
    "logout": {"status": logout_status, "response": logout},
    "post_logout": {"status": invalid_status, "response": invalid},
}
print(json.dumps(summary, indent=2, sort_keys=True))
if not summary["unique"] or not summary["gap_free"] or logout_status != 200 or invalid_status != 401:
    raise SystemExit(1)
PY