#!/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_TOKEN_DAYS="${AI_SBX_TOKEN_DAYS:-30}"
DEFAULT_TEMPLATE="${AI_SBX_TEMPLATE:-}"

# GitHub accepts these as query parameters on the token creation form. A write
# level implies read, so only the highest level is listed. "workflows" is
# required to push any commit touching .github/workflows and is separate from
# "actions".
TOKEN_URL_PERMISSIONS=(
    metadata=read
    contents=write
    pull_requests=write
    issues=write
    workflows=write
    actions=write
    statuses=read
    security_events=write
)

# GitHub omits these from the pre-fill parameters, so they can only be ticked
# on the form itself.
TOKEN_MANUAL_PERMISSIONS=(
    "Checks: Read"
    "Dependabot alerts: Read"
    "Secret scanning alerts: Read"
)

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

Setup opens a pre-filled GitHub token form in your browser. Use "token" on
its own to replace an expired or revoked token later.

Set AI_SBX_TEMPLATE in your mise config to reuse one custom image across
every repository without repeating --template.

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 private in-container clone
                         of the repository, mounted read-only.
                         Its commits reach the host through the
                         sandbox-<name> git remote. This is the default.
  --template REF         Custom sandbox image. Defaults to
                         AI_SBX_TEMPLATE when set.
  --stock-template       Ignore AI_SBX_TEMPLATE and use the agent's
                         stock image.
  --kit PATH             Mixin kit to apply. May be supplied more
                         than once.
  --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"
}

url_encode() {
    local string="$1" index character encoded=""

    for ((index = 0; index < ${#string}; index++)); do
        character="${string:index:1}"
        case "$character" in
            [a-zA-Z0-9.~_-])
                encoded+="$character"
                ;;
            *)
                printf -v character '%%%02X' "'$character"
                encoded+="$character"
                ;;
        esac
    done

    printf '%s' "$encoded"
}

token_url() {
    local owner="${REPOSITORY%%/*}"
    local name="${REPOSITORY#*/}"
    local url="https://github.com/settings/personal-access-tokens/new"

    url+="?name=$(url_encode "ai-sbx $name")"
    url+="&description=$(url_encode "AI agent sandbox for $REPOSITORY")"
    url+="&target_name=$(url_encode "$owner")"
    url+="&expires_in=$(url_encode "$DEFAULT_TOKEN_DAYS")"

    local permission
    for permission in "${TOKEN_URL_PERMISSIONS[@]}"; do
        url+="&$permission"
    done

    printf '%s' "$url"
}

open_browser() {
    local url="$1" opener

    for opener in "${BROWSER:-}" xdg-open open; do
        [[ -n "$opener" ]] || continue

        if command -v "$opener" >/dev/null 2>&1; then
            "$opener" "$url" >/dev/null 2>&1 &
            return 0
        fi
    done

    return 1
}

read_token() {
    local token

    # -s keeps the token off the terminal; it never reaches shell history
    # because it is read into a variable rather than typed as an argument.
    IFS= read -rsp 'Paste token: ' token </dev/tty
    printf '\n' >&2

    [[ -n "$token" ]] ||
        die "No token entered."

    case "$token" in
        github_pat_*) ;;
        ghp_*)
            die "That is a classic token. Generate a fine-grained token from the link above."
            ;;
        *)
            die "That does not look like a fine-grained token (expected a github_pat_ prefix)."
            ;;
    esac

    printf '%s' "$token"
}

install_github_token() {
    local url
    url="$(token_url)"

    cat >&2 <<EOF

Create a fine-grained token for $REPOSITORY.

Everything except the repository is pre-filled. On the page:

  1. Repository access -> Only select repositories -> ${REPOSITORY#*/}
  2. Tick the permissions the form cannot pre-fill:
EOF

    local permission
    for permission in "${TOKEN_MANUAL_PERMISSIONS[@]}"; do
        printf '       %s\n' "$permission" >&2
    done

    cat >&2 <<EOF
  3. Generate token, then paste it below.

EOF

    if open_browser "$url"; then
        printf 'Opened your browser.\n\n' >&2
    else
        printf 'Open this link:\n\n%s\n\n' "$url" >&2
    fi

    # --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.
    read_token |
        sbx secret set --force "$SANDBOX_NAME" github >/dev/null

    printf 'Stored the token for %s in sandbox %s.\n' "$REPOSITORY" "$SANDBOX_NAME" >&2
}
# 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"


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

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

    CONFIG_TEMPLATE="${CONFIG_TEMPLATE:-}"
}

save_config() {
    local agent="$1"
    local mode="$2"
    local template="$3"
    local kit_count="$4"
    shift 4

    local -a kits=("${@:1:kit_count}")
    local -a profiles=("${@:kit_count + 1}")

    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_TEMPLATE=%q\n' "$template"

        printf 'CONFIG_KITS=('
        local kit
        for kit in ${kits[@]+"${kits[@]}"}; do
            printf ' %q' "$kit"
        done
        printf ' )\n'

        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 its own in-container clone, so its commits
    # never land on whatever the host has checked out.
    if [[ "$CONFIG_MODE" == "clone" ]]; then
        create_args+=(--clone)
    fi

    if [[ -n "$CONFIG_TEMPLATE" ]]; then
        create_args+=(--template "$CONFIG_TEMPLATE")
    fi

    local kit
    for kit in ${CONFIG_KITS[@]+"${CONFIG_KITS[@]}"}; do
        create_args+=(--kit "$kit")
    done

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

    sbx "${create_args[@]}"
}

setup_command() {
    local agent="$DEFAULT_AGENT"
    local mode="$DEFAULT_MODE"
    local template="$DEFAULT_TEMPLATE"
    local replace=false
    local -a aws_profiles=()
    local -a kits=()

    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
                ;;
            --direct)
                mode="direct"
                shift
                ;;
            --template)
                (($# >= 2)) || die "--template requires a value"
                template="$2"
                shift 2
                ;;
            --stock-template)
                template=""
                shift
                ;;
            --kit)
                (($# >= 2)) || die "--kit requires a value"
                kits+=("$2")
                shift 2
                ;;
            --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 kit
    for kit in ${kits[@]+"${kits[@]}"}; do
        [[ -e "$kit" ]] ||
            die "Kit does not exist: $kit"
    done

    save_config "$agent" "$mode" "$template" "${#kits[@]}" \
        ${kits[@]+"${kits[@]}"} ${aws_profiles[@]+"${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

    install_github_token

    install_sandbox_aws_files

    cat <<EOF

Setup complete.

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

Run it with:

  mise run ai:sbx -- run
EOF
}

token_command() {
    load_config

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

    install_github_token
}

refresh_command() {
    load_config

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

    install_sandbox_aws_files
}

run_command() {
    load_config

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

    # The GitHub token is long-lived and stays in the sbx secret store; only
    # the AWS credentials expire between sessions.
    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"

    printf 'Template:   %s\n' "${CONFIG_TEMPLATE:-stock}"

    if ((${#CONFIG_KITS[@]})); then
        printf 'Kits:\n'
        printf '  %s\n' "${CONFIG_KITS[@]}"
    fi

    printf 'Token days: %s\n' "$DEFAULT_TOKEN_DAYS"

    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
            ;;
    esac

    require_command git
    require_command sbx
    require_command sha256sum

    repository_context

    case "$command" in
        setup)
            setup_command "$@"
            ;;
        token)
            token_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