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

PROGRAM="ai:sbx"
CONFIG_ROOT="${XDG_CONFIG_HOME:-$HOME/.config}/ai-sbx"
DEFAULT_AGENT="${AI_SBX_AGENT:-codex}"
DEFAULT_MODE="${AI_SBX_MODE:-clone}"
DEFAULT_BRANCH="${AI_SBX_BRANCH:-ai-sbx}"
APP_CONFIG_FILE="$CONFIG_ROOT/github-app"

# Keys and levels are validated against GitHub's app-permissions schema.
# "workflows" has no read level; write is required to push any commit that
# touches .github/workflows.
GITHUB_APP_PERMISSIONS='{
  "metadata": "read",
  "contents": "write",
  "pull_requests": "write",
  "issues": "write",
  "workflows": "write",
  "actions": "write",
  "checks": "read",
  "statuses": "read",
  "security_events": "write",
  "secret_scanning_alerts": "read",
  "vulnerability_alerts": "read"
}'

usage() {
    cat <<'EOF'
Usage:
  mise run ai:sbx -- app --app-id ID --key PATH
  mise run ai:sbx -- setup [options]
  mise run ai:sbx -- refresh
  mise run ai:sbx -- run [-- agent arguments...]
  mise run ai:sbx -- status
  mise run ai:sbx -- remove

App options (configured once, for every repository):
  --app-id ID            Numeric GitHub App ID, not the client ID.
  --key PATH             The App's RSA private key (.pem).

Setup options:
  --aws-profile NAME     Host AWS profile to expose inside the sandbox.
                         May be supplied more than once. A trailing
                         -readonly is stripped from the profile name
                         written into the sandbox.
  --agent NAME           Sandbox agent. Default: codex
  --direct               Mount the host working tree read-write.
  --clone                Give the agent a Git worktree on its own
                         branch. This is the default.
  --branch NAME          Branch used by --clone. Default: ai-sbx
  --replace              Replace the existing sandbox.

Examples:
  mise run ai:sbx -- setup \
    --aws-profile api-portal-readonly \
    --aws-profile prod-readonly

  mise run ai:sbx -- run

  mise run ai:sbx -- run -- --dangerously-bypass-approvals-and-sandbox \
    "Review the Terraform plan"
EOF
}

die() {
    printf '%s: %s\n' "$PROGRAM" "$*" >&2
    exit 1
}

require_command() {
    command -v "$1" >/dev/null 2>&1 ||
        die "Required command not found: $1"
}

github_app_configured() {
    [[ -f "$APP_CONFIG_FILE" ]]
}

load_app_config() {
    github_app_configured ||
        die "No GitHub App configured. Run: mise run ai:sbx -- app --app-id ID --key PATH"

    # shellcheck disable=SC1090
    source "$APP_CONFIG_FILE"

    [[ -n "${APP_ID:-}" ]] ||
        die "APP_ID is missing from $APP_CONFIG_FILE"

    [[ -r "${APP_PRIVATE_KEY_FILE:-}" ]] ||
        die "GitHub App private key is not readable: ${APP_PRIVATE_KEY_FILE:-unset}"
}

base64url() {
    openssl base64 -A | tr '+/' '-_' | tr -d '='
}

# GitHub caps App JWT lifetime at 10 minutes and rejects future iat values, so
# backdate slightly to tolerate clock skew and stay well inside the cap.
github_app_jwt() {
    local now header payload signing_input signature

    now="$(date +%s)"
    header='{"alg":"RS256","typ":"JWT"}'
    payload="$(printf '{"iat":%d,"exp":%d,"iss":"%s"}' \
        "$((now - 60))" "$((now + 540))" "$APP_ID")"

    signing_input="$(printf '%s' "$header" | base64url).$(printf '%s' "$payload" | base64url)"

    signature="$(
        printf '%s' "$signing_input" |
            openssl dgst -sha256 -sign "$APP_PRIVATE_KEY_FILE" -binary |
            base64url
    )"

    printf '%s.%s' "$signing_input" "$signature"
}

github_api() {
    local method="$1" path="$2" token="$3"
    shift 3

    curl --silent --show-error \
        --request "$method" \
        --header "Authorization: Bearer $token" \
        --header "Accept: application/vnd.github+json" \
        --header "X-GitHub-Api-Version: 2022-11-28" \
        "https://api.github.com$path" \
        "$@"
}

# GitHub answers errors with HTTP 4xx and a .message body, which curl alone
# treats as success, so every response is inspected before it is used.
github_api_field() {
    local response="$1" field="$2" context="$3" value

    if value="$(jq -er "$field" <<<"$response" 2>/dev/null)"; then
        printf '%s' "$value"
        return
    fi

    local message
    message="$(jq -r '.message // "unrecognized response"' <<<"$response" 2>/dev/null)" ||
        message="unparseable response"

    die "$context: $message"
}

resolve_installation_id() {
    local jwt response

    jwt="$(github_app_jwt)"
    response="$(github_api GET "/repos/$REPOSITORY/installation" "$jwt")"

    github_api_field "$response" '.id' \
        "GitHub App is not installed on $REPOSITORY"
}

mint_github_token() {
    local jwt response body

    jwt="$(github_app_jwt)"
    body="$(
        jq -nc \
            --arg repo "${REPOSITORY#*/}" \
            --argjson permissions "$GITHUB_APP_PERMISSIONS" \
            '{repositories: [$repo], permissions: $permissions}'
    )"

    response="$(
        github_api POST \
            "/app/installations/$CONFIG_INSTALLATION_ID/access_tokens" \
            "$jwt" --data "$body"
    )"

    github_api_field "$response" '.token' \
        "Could not mint an installation token for $REPOSITORY"
}

install_github_token() {
    require_command openssl
    require_command curl
    require_command jq

    load_app_config

    [[ -n "${CONFIG_INSTALLATION_ID:-}" ]] ||
        die "Installation ID is missing. Re-run: mise run ai:sbx -- setup"

    # --force is mandatory: without it a second write prompts for confirmation,
    # reads the prompt from the already-consumed stdin, cancels, and still
    # exits 0 — leaving the previous, expired token in place.
    mint_github_token |
        sbx secret set --force "$SANDBOX_NAME" github >/dev/null

    printf 'Installed a fresh GitHub App token for %s (expires in 1 hour).\n' \
        "$REPOSITORY"
}

# Terraform and provider blocks reference the account profile name, while the
# host distinguishes the read-only grant with a -readonly suffix. The suffix is
# a host-side naming convention, so it is stripped on the way into the sandbox.
sandbox_profile_name() {
    local profile="$1"
    local mapped="${profile%-readonly}"

    [[ -n "$mapped" ]] ||
        die "AWS profile name is empty after stripping -readonly: $profile"

    printf '%s' "$mapped"
}

# A task included from the global mise config runs with the config root as its
# working directory ($HOME), not the directory the user invoked it from, so the
# repository would otherwise be undiscoverable from anywhere.
enter_invocation_directory() {
    local invoked_from="${MISE_ORIGINAL_CWD:-$PWD}"

    cd "$invoked_from" ||
        die "Could not enter the invoking directory: $invoked_from"
}

repository_context() {
    enter_invocation_directory

    REPO_ROOT="$(git rev-parse --show-toplevel 2>/dev/null)" ||
        die "This command must be run inside a Git repository."

    local remote
    remote="$(git -C "$REPO_ROOT" remote get-url origin 2>/dev/null)" ||
        die "The repository has no origin remote."

    case "$remote" in
        git@github.com:*)
            REPOSITORY="${remote#git@github.com:}"
            ;;
        ssh://git@github.com/*)
            REPOSITORY="${remote#ssh://git@github.com/}"
            ;;
        https://github.com/*)
            REPOSITORY="${remote#https://github.com/}"
            ;;
        http://github.com/*)
            REPOSITORY="${remote#http://github.com/}"
            ;;
        *)
            die "Unsupported GitHub origin: $remote"
            ;;
    esac

    REPOSITORY="${REPOSITORY%.git}"
    REPOSITORY="${REPOSITORY%/}"

    [[ "$REPOSITORY" =~ ^[^/]+/[^/]+$ ]] ||
        die "Could not derive owner/repository from origin: $remote"

    local slug
    slug="$(
        printf '%s' "$REPOSITORY" |
            tr '[:upper:]' '[:lower:]' |
            tr '/_' '--' |
            tr -cd 'a-z0-9.-'
    )"

    # Include a short digest to avoid collisions caused by normalization.
    local digest
    digest="$(
        printf '%s' "$REPOSITORY" |
            sha256sum |
            cut -c1-10
    )"

    SANDBOX_NAME="ai-${slug}-${digest}"
    REPO_CONFIG_DIR="$CONFIG_ROOT/repos/$digest"
    REPO_CONFIG_FILE="$REPO_CONFIG_DIR/config"
}

sandbox_exists() {
    sbx ls --quiet 2>/dev/null |
        grep -Fxq "$SANDBOX_NAME"
}

load_config() {
    [[ -f "$REPO_CONFIG_FILE" ]] ||
        die "Repository is not configured. Run: mise run ai:sbx -- setup"

    # This file is user-owned, mode 600, and contains no credentials.
    # shellcheck disable=SC1090
    source "$REPO_CONFIG_FILE"

    [[ "${CONFIG_REPOSITORY:-}" == "$REPOSITORY" ]] ||
        die "Repository configuration does not match the current origin."

    [[ -n "${CONFIG_AGENT:-}" ]] ||
        die "Agent is missing from $REPO_CONFIG_FILE"

    CONFIG_BRANCH="${CONFIG_BRANCH:-$DEFAULT_BRANCH}"
    CONFIG_INSTALLATION_ID="${CONFIG_INSTALLATION_ID:-}"

    declare -p CONFIG_AWS_PROFILES >/dev/null 2>&1 ||
        CONFIG_AWS_PROFILES=()
}

save_config() {
    local agent="$1"
    local mode="$2"
    local branch="$3"
    local installation_id="$4"
    shift 4
    local -a profiles=("$@")

    mkdir -p "$REPO_CONFIG_DIR"
    chmod 700 "$CONFIG_ROOT" "$CONFIG_ROOT/repos" "$REPO_CONFIG_DIR" 2>/dev/null || true

    {
        printf 'CONFIG_REPOSITORY=%q\n' "$REPOSITORY"
        printf 'CONFIG_SANDBOX=%q\n' "$SANDBOX_NAME"
        printf 'CONFIG_AGENT=%q\n' "$agent"
        printf 'CONFIG_MODE=%q\n' "$mode"
        printf 'CONFIG_BRANCH=%q\n' "$branch"
        printf 'CONFIG_INSTALLATION_ID=%q\n' "$installation_id"

        printf 'CONFIG_AWS_PROFILES=('
        local profile
        for profile in "${profiles[@]}"; do
            printf ' %q' "$profile"
        done
        printf ' )\n'
    } >"$REPO_CONFIG_FILE"

    chmod 600 "$REPO_CONFIG_FILE"
}

validate_aws_profile() {
    local profile="$1"

    aws configure list-profiles | grep -Fxq "$profile" ||
        die "AWS profile does not exist on the host: $profile"

    printf 'Validating AWS profile %s...\n' "$profile" >&2

    if ! aws sts get-caller-identity \
        --profile "$profile" \
        --output json \
        >/dev/null; then
        printf '\nAWS authentication failed for profile %s.\n' "$profile" >&2
        printf 'Run:\n\n  aws sso login --profile %q\n\n' "$profile" >&2
        exit 1
    fi
}

# Two host profiles mapping to the same sandbox name would silently write two
# sections with one identity, so reject it before any credentials are exported.
validate_profile_mapping() {
    local -A claimed_by=()
    local profile mapped

    for profile in "$@"; do
        mapped="$(sandbox_profile_name "$profile")"

        if [[ -n "${claimed_by[$mapped]:-}" ]]; then
            die "AWS profiles ${claimed_by[$mapped]} and $profile both map to sandbox profile $mapped"
        fi

        claimed_by["$mapped"]="$profile"
    done
}

write_aws_files() {
    load_config

    local output_dir="$1"
    local config_file="$output_dir/config"
    local credentials_file="$output_dir/credentials"

    validate_profile_mapping "${CONFIG_AWS_PROFILES[@]}"

    mkdir -p "$output_dir"
    chmod 700 "$output_dir"

    : >"$config_file"
    : >"$credentials_file"

    local profile
    for profile in "${CONFIG_AWS_PROFILES[@]}"; do
        validate_aws_profile "$profile"

        local sandbox_profile
        sandbox_profile="$(sandbox_profile_name "$profile")"

        local credential_json
        credential_json="$(
            aws configure export-credentials \
                --profile "$profile" \
                --format process
        )"

        local access_key secret_key session_token expiration region output
        access_key="$(jq -er '.AccessKeyId' <<<"$credential_json")"
        secret_key="$(jq -er '.SecretAccessKey' <<<"$credential_json")"
        session_token="$(jq -er '.SessionToken' <<<"$credential_json")"
        expiration="$(jq -er '.Expiration // empty' <<<"$credential_json" || true)"

        region="$(
            aws configure get region --profile "$profile" 2>/dev/null ||
                true
        )"
        output="$(
            aws configure get output --profile "$profile" 2>/dev/null ||
                true
        )"

        region="${region:-us-east-1}"
        output="${output:-json}"

        cat >>"$config_file" <<EOF
[profile $sandbox_profile]
region = $region
output = $output

EOF

        cat >>"$credentials_file" <<EOF
[$sandbox_profile]
aws_access_key_id = $access_key
aws_secret_access_key = $secret_key
aws_session_token = $session_token

EOF

        printf 'Exported %-30s as %-30s expires %s\n' \
            "$profile" \
            "$sandbox_profile" \
            "${expiration:-unknown}" >&2

        unset credential_json access_key secret_key session_token
    done

    chmod 600 "$config_file" "$credentials_file"
}

install_sandbox_aws_files() {
    load_config

    if ((${#CONFIG_AWS_PROFILES[@]} == 0)); then
        printf 'No AWS profiles configured; skipping AWS credential refresh.\n'
        return
    fi

    require_command aws
    require_command jq

    local temporary_directory
    temporary_directory="$(mktemp -d)"
    trap 'rm -rf "$temporary_directory"' RETURN

    write_aws_files "$temporary_directory"

    # The agent user differs between sandbox images, so ask rather than assume.
    local sandbox_home
    # shellcheck disable=SC2016
    sandbox_home="$(sbx exec "$SANDBOX_NAME" bash -c 'printf %s "$HOME"')"

    [[ -n "$sandbox_home" ]] ||
        die "Could not determine the sandbox home directory."

    # shellcheck disable=SC2016
    sbx exec "$SANDBOX_NAME" \
        bash -c 'mkdir -p "$HOME/.aws" && chmod 700 "$HOME/.aws"'

    sbx cp \
        "$temporary_directory/config" \
        "$SANDBOX_NAME:$sandbox_home/.aws/config"

    sbx cp \
        "$temporary_directory/credentials" \
        "$SANDBOX_NAME:$sandbox_home/.aws/credentials"

    # Every expansion below belongs to the sandbox shell, not the host.
    # shellcheck disable=SC2016
    sbx exec "$SANDBOX_NAME" bash -c '
        chmod 700 "$HOME/.aws"
        chmod 600 "$HOME/.aws/config" "$HOME/.aws/credentials"

        persistent=/etc/sandbox-persistent.sh
        marker="# BEGIN ai-sbx AWS configuration"

        if grep -Fq "$marker" "$persistent" 2>/dev/null; then
            exit 0
        fi

        cat >>"$persistent" <<'"'"'EOF'"'"'
# BEGIN ai-sbx AWS configuration
export AWS_CONFIG_FILE="$HOME/.aws/config"
export AWS_SHARED_CREDENTIALS_FILE="$HOME/.aws/credentials"
export AWS_SDK_LOAD_CONFIG=1
export AWS_EC2_METADATA_DISABLED=true
unset AWS_ACCESS_KEY_ID
unset AWS_SECRET_ACCESS_KEY
unset AWS_SESSION_TOKEN
unset AWS_SECURITY_TOKEN
# END ai-sbx AWS configuration
EOF
    '

    rm -rf "$temporary_directory"
    trap - RETURN

    printf 'Installed isolated AWS profiles in sandbox %s:\n' "$SANDBOX_NAME"

    local profile
    for profile in "${CONFIG_AWS_PROFILES[@]}"; do
        printf '  %s -> %s\n' "$profile" "$(sandbox_profile_name "$profile")"
    done
}

create_sandbox() {
    load_config

    local -a create_args=(
        create
        --name "$SANDBOX_NAME"
    )

    # Clone mode gives the agent a Git worktree on its own branch, so its
    # commits never land on whatever the host has checked out.
    if [[ "$CONFIG_MODE" == "clone" ]]; then
        create_args+=(--branch "$CONFIG_BRANCH")
    fi

    create_args+=(
        "$CONFIG_AGENT"
        "$REPO_ROOT"
    )

    sbx "${create_args[@]}"
}

setup_command() {
    local agent="$DEFAULT_AGENT"
    local mode="$DEFAULT_MODE"
    local branch="$DEFAULT_BRANCH"
    local replace=false
    local -a aws_profiles=()

    while (($#)); do
        case "$1" in
            --aws-profile)
                (($# >= 2)) || die "--aws-profile requires a value"
                aws_profiles+=("$2")
                shift 2
                ;;
            --agent)
                (($# >= 2)) || die "--agent requires a value"
                agent="$2"
                shift 2
                ;;
            --clone)
                mode="clone"
                shift
                ;;
            --branch)
                (($# >= 2)) || die "--branch requires a value"
                branch="$2"
                shift 2
                ;;
            --direct)
                mode="direct"
                shift
                ;;
            --replace)
                replace=true
                shift
                ;;
            -h | --help)
                usage
                exit 0
                ;;
            *)
                die "Unknown setup option: $1"
                ;;
        esac
    done

    validate_profile_mapping "${aws_profiles[@]}"

    local profile
    for profile in "${aws_profiles[@]}"; do
        validate_aws_profile "$profile"
    done

    local installation_id=""
    if github_app_configured; then
        require_command openssl
        require_command curl
        require_command jq
        load_app_config
        installation_id="$(resolve_installation_id)"
        printf 'GitHub App installation for %s: %s\n' "$REPOSITORY" "$installation_id"
    fi

    save_config "$agent" "$mode" "$branch" "$installation_id" "${aws_profiles[@]}"

    if sandbox_exists; then
        if [[ "$replace" == true ]]; then
            printf 'Removing existing sandbox %s...\n' "$SANDBOX_NAME"
            sbx rm "$SANDBOX_NAME"
        else
            printf 'Using existing sandbox %s.\n' "$SANDBOX_NAME"
        fi
    fi

    if ! sandbox_exists; then
        printf 'Creating sandbox %s for %s...\n' \
            "$SANDBOX_NAME" "$REPOSITORY"
        create_sandbox
    fi

    if github_app_configured; then
        install_github_token
    else
        cat <<EOF

No GitHub App is configured, so this sandbox needs a fine-grained token.

Configure the App once instead, and every repository afterwards is automatic:

  mise run ai:sbx -- app --app-id ID --key PATH

Otherwise, create a token restricted to:

  Repository: $REPOSITORY
  Sandbox:    $SANDBOX_NAME

Permissions:
  Metadata:                Read
  Contents:                Read and write
  Pull requests:           Read and write
  Issues:                  Read and write
  Workflows:               Read and write
  Actions:                 Read and write
  Checks:                  Read
  Commit statuses:         Read
  Code scanning alerts:    Read and write
  Secret scanning alerts:  Read
  Dependabot alerts:       Read

EOF

        # Interactive prompt; the token is not placed in shell history.
        sbx secret set --force "$SANDBOX_NAME" github
    fi

    install_sandbox_aws_files

    cat <<EOF

Setup complete.

Repository: $REPOSITORY
Sandbox:   $SANDBOX_NAME
Agent:     $agent
Mode:      $mode
Branch:    $branch

Run it with:

  mise run ai:sbx -- run
EOF
}

app_command() {
    local app_id="" key=""

    while (($#)); do
        case "$1" in
            --app-id)
                (($# >= 2)) || die "--app-id requires a value"
                app_id="$2"
                shift 2
                ;;
            --key)
                (($# >= 2)) || die "--key requires a value"
                key="$2"
                shift 2
                ;;
            -h | --help)
                usage
                exit 0
                ;;
            *)
                die "Unknown app option: $1"
                ;;
        esac
    done

    [[ "$app_id" =~ ^[0-9]+$ ]] ||
        die "--app-id must be the numeric App ID, not the client ID"

    [[ -r "$key" ]] ||
        die "Private key is not readable: ${key:-unset}"

    key="$(cd "$(dirname "$key")" && printf '%s/%s' "$PWD" "$(basename "$key")")"

    openssl rsa -in "$key" -noout 2>/dev/null ||
        die "Not a usable RSA private key: $key"

    mkdir -p "$CONFIG_ROOT"
    chmod 700 "$CONFIG_ROOT"

    {
        printf 'APP_ID=%q\n' "$app_id"
        printf 'APP_PRIVATE_KEY_FILE=%q\n' "$key"
    } >"$APP_CONFIG_FILE"

    chmod 600 "$APP_CONFIG_FILE"

    printf 'Recorded GitHub App %s in %s\n' "$app_id" "$APP_CONFIG_FILE"
    printf 'Install it on each repository, then run setup there.\n'
}

refresh_command() {
    load_config

    sandbox_exists ||
        die "Sandbox does not exist. Run: mise run ai:sbx -- setup"

    if github_app_configured; then
        install_github_token
    fi

    install_sandbox_aws_files
}

run_command() {
    load_config

    sandbox_exists ||
        die "Sandbox does not exist. Run: mise run ai:sbx -- setup"

    # Both credentials are short-lived, so re-mint before every session.
    if github_app_configured; then
        install_github_token
    fi

    install_sandbox_aws_files

    if (($#)) && [[ "$1" == "--" ]]; then
        shift
    fi

    if (($#)); then
        exec sbx run "$SANDBOX_NAME" -- "$@"
    else
        exec sbx run "$SANDBOX_NAME"
    fi
}

status_command() {
    load_config

    printf 'Repository: %s\n' "$REPOSITORY"
    printf 'Root:       %s\n' "$REPO_ROOT"
    printf 'Sandbox:    %s\n' "$SANDBOX_NAME"
    printf 'Agent:      %s\n' "$CONFIG_AGENT"
    printf 'Mode:       %s\n' "$CONFIG_MODE"

    if [[ "$CONFIG_MODE" == "clone" ]]; then
        printf 'Branch:     %s\n' "$CONFIG_BRANCH"
    fi

    if github_app_configured; then
        # shellcheck disable=SC1090
        source "$APP_CONFIG_FILE"
        printf 'GitHub:     App %s, installation %s\n' \
            "${APP_ID:-unset}" "${CONFIG_INSTALLATION_ID:-unresolved}"
    else
        printf 'GitHub:     manual fine-grained token\n'
    fi

    printf 'AWS profiles (host -> sandbox):\n'
    if ((${#CONFIG_AWS_PROFILES[@]})); then
        local profile
        for profile in "${CONFIG_AWS_PROFILES[@]}"; do
            printf '  %s -> %s\n' "$profile" "$(sandbox_profile_name "$profile")"
        done
    else
        printf '  none\n'
    fi

    printf 'Sandbox exists: '
    if sandbox_exists; then
        printf 'yes\n'
    else
        printf 'no\n'
    fi

    printf '\nConfigured sandbox secrets:\n'
    sbx secret ls
}

remove_command() {
    load_config

    if sandbox_exists; then
        sbx rm "$SANDBOX_NAME"
    fi

    rm -rf "$REPO_CONFIG_DIR"

    printf 'Removed sandbox and local configuration for %s.\n' "$REPOSITORY"
}

main() {
    local command="${1:-}"
    if (($#)); then
        shift
    fi

    # These work outside a repository and without the sandbox toolchain.
    case "$command" in
        -h | --help | help | "")
            usage
            return
            ;;
        app)
            app_command "$@"
            return
            ;;
    esac

    require_command git
    require_command sbx
    require_command sha256sum

    repository_context

    case "$command" in
        setup)
            setup_command "$@"
            ;;
        refresh)
            refresh_command "$@"
            ;;
        run)
            run_command "$@"
            ;;
        status)
            status_command "$@"
            ;;
        remove)
            remove_command "$@"
            ;;
        *)
            die "Unknown command: $command"
            ;;
    esac
}

# Sourcing the task exposes its functions for tests without running a command.
if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then
    main "$@"
fi