#!/usr/bin/env bash
set -e

# @describe Perform a web search to get up-to-date information or additional context.
# Use this when you need current information or feel a search could provide a better answer.

# @option --query! The search query.

# @meta require-tools loki

# @env WEB_SEARCH_MODEL=gemini:gemini-2.5-flash The model for web-searching.
#
# supported loki models:
#   - gemini:gemini-2.0-*
#   - vertexai:gemini-*
#   - perplexity:*
#   - ernie:*
# @env LLM_OUTPUT=/dev/stdout The output path

# shellcheck disable=SC2154
main() {
    client="${WEB_SEARCH_MODEL%%:*}"

    if [[ "$client" == "gemini" ]]; then
        export LOKI_PATCH_GEMINI_CHAT_COMPLETIONS='{".*":{"body":{"tools":[{"google_search":{}}]}}}'
    elif [[ "$client" == "vertexai" ]]; then
        export LOKI_PATCH_VERTEXAI_CHAT_COMPLETIONS='{
    "gemini-1.5-.*":{"body":{"tools":[{"googleSearchRetrieval":{}}]}},
    "gemini-2.0-.*":{"body":{"tools":[{"google_search":{}}]}}
}'
    elif [[ "$client" == "ernie" ]]; then
        export LOKI_PATCH_ERNIE_CHAT_COMPLETIONS='{".*":{"body":{"web_search":{"enable":true}}}}'
    fi

    loki -m "$WEB_SEARCH_MODEL" "$argc_query" >> "$LLM_OUTPUT"
}


# ARGC-BUILD {
# This block was generated by argc (https://github.com/sigoden/argc).
# Modifying it manually is not recommended

_argc_run() {
    if [[ "${1:-}" == "___internal___" ]]; then
        _argc_die "error: unsupported ___internal___ command"
    fi
    if [[ "${OS:-}" == "Windows_NT" ]] && [[ -n "${MSYSTEM:-}" ]]; then
        set -o igncr
    fi
    argc__args=("$(basename "$0" .sh)" "$@")
    argc__positionals=()
    _argc_index=1
    _argc_len="${#argc__args[@]}"
    _argc_tools=()
    _argc_parse
    _argc_require_tools "${_argc_tools[@]}"
    if [ -n "${argc__fn:-}" ]; then
        $argc__fn "${argc__positionals[@]}"
    fi
}

_argc_usage() {
    cat <<-'EOF'
Perform a web search to get up-to-date information or additional context.
Use this when you need current information or feel a search could provide a better answer.

USAGE:  --query <QUERY>

OPTIONS:
      --query <QUERY>  The query to search for.
  -h, --help           Print help
  -V, --version        Print version

ENVIRONMENTS:
  WEB_SEARCH_MODEL
          The model for web-searching.
          
          supported loki models:
            - gemini:gemini-2.0-*
            - vertexai:gemini-*
            - perplexity:*
            - ernie:*
          [default: gemini:gemini-2.5-flash]

  LLM_OUTPUT
          The output path [default: /dev/stdout]
EOF
    exit
}

_argc_version() {
    echo  0.0.0
    exit
}

_argc_parse() {
    local _argc_key _argc_action
    local _argc_subcmds=""
    while [[ $_argc_index -lt $_argc_len ]]; do
        _argc_item="${argc__args[_argc_index]}"
        _argc_key="${_argc_item%%=*}"
        case "$_argc_key" in
        --help | -help | -h)
            _argc_usage
            ;;
        --version | -version | -V)
            _argc_version
            ;;
        --)
            _argc_dash="${#argc__positionals[@]}"
            argc__positionals+=("${argc__args[@]:$((_argc_index + 1))}")
            _argc_index=$_argc_len
            break
            ;;
        --query)
            _argc_take_args "--query <QUERY>" 1 1 "-" ""
            _argc_index=$((_argc_index + _argc_take_args_len + 1))
            if [[ -z "${argc_query:-}" ]]; then
                argc_query="${_argc_take_args_values[0]:-}"
            else
                _argc_die "error: the argument \`--query\` cannot be used multiple times"
            fi
            ;;
        *)
            if _argc_maybe_flag_option "-" "$_argc_item"; then
                _argc_die "error: unexpected argument \`$_argc_key\` found"
            fi
            argc__positionals+=("$_argc_item")
            _argc_index=$((_argc_index + 1))
            ;;
        esac
    done
    _argc_require_params "error: the following required arguments were not provided:" \
        'argc_query:--query <QUERY>'
    _argc_tools=(loki)
    if [[ -n "${_argc_action:-}" ]]; then
        $_argc_action
    else
        argc__fn=main
        if [[ "${argc__positionals[0]:-}" == "help" ]] && [[ "${#argc__positionals[@]}" -eq 1 ]]; then
            _argc_usage
        fi
        if [[ -z "${WEB_SEARCH_MODEL:-}" ]]; then
            export WEB_SEARCH_MODEL=gemini:gemini-2.5-flash
        fi
        if [[ -z "${LLM_OUTPUT:-}" ]]; then
            export LLM_OUTPUT=/dev/stdout
        fi
    fi
}

_argc_take_args() {
    _argc_take_args_values=()
    _argc_take_args_len=0
    local param="$1" min="$2" max="$3" signs="$4" delimiter="$5"
    if [[ "$min" -eq 0 ]] && [[ "$max" -eq 0 ]]; then
        return
    fi
    local _argc_take_index=$((_argc_index + 1)) _argc_take_value
    if [[ "$_argc_item" == *=* ]]; then
        _argc_take_args_values=("${_argc_item##*=}")
    else
        while [[ $_argc_take_index -lt $_argc_len ]]; do
            _argc_take_value="${argc__args[_argc_take_index]}"
            if _argc_maybe_flag_option "$signs" "$_argc_take_value"; then
                if [[ "${#_argc_take_value}" -gt 1 ]]; then
                    break
                fi
            fi
            _argc_take_args_values+=("$_argc_take_value")
            _argc_take_args_len=$((_argc_take_args_len + 1))
            if [[ "$_argc_take_args_len" -ge "$max" ]]; then
                break
            fi
            _argc_take_index=$((_argc_take_index + 1))
        done
    fi
    if [[ "${#_argc_take_args_values[@]}" -lt "$min" ]]; then
        _argc_die "error: incorrect number of values for \`$param\`"
    fi
    if [[ -n "$delimiter" ]] && [[ "${#_argc_take_args_values[@]}" -gt 0 ]]; then
        local item values arr=()
        for item in "${_argc_take_args_values[@]}"; do
            IFS="$delimiter" read -r -a values <<<"$item"
            arr+=("${values[@]}")
        done
        _argc_take_args_values=("${arr[@]}")
    fi
}

_argc_require_params() {
    local message="$1" missed_envs="" item name render_name
    for item in "${@:2}"; do
        name="${item%%:*}"
        render_name="${item##*:}"
        if [[ -z "${!name:-}" ]]; then
            missed_envs="$missed_envs"$'\n'"  $render_name"
        fi
    done
    if [[ -n "${missed_envs}" ]]; then
        _argc_die "$message$missed_envs"
    fi
}

_argc_maybe_flag_option() {
    local signs="$1" arg="$2"
    if [[ -z "$signs" ]]; then
        return 1
    fi
    local cond=false
    if [[ "$signs" == *"+"* ]]; then
        if [[ "$arg" =~ ^\+[^+].* ]]; then
            cond=true
        fi
    elif [[ "$arg" == -* ]]; then
        if (( ${#arg} < 3 )) || [[ ! "$arg" =~ ^---.* ]]; then
            cond=true
        fi
    fi
    if [[ "$cond" == "false" ]]; then
        return 1
    fi
    local value="${arg%%=*}"
    if [[ "$value" =~ [[:space:]] ]]; then
        return 1
    fi
    return 0
}

_argc_require_tools() {
    local tool missing_tools=()
    for tool in "$@"; do
        if ! command -v "$tool" >/dev/null 2>&1; then
            missing_tools+=("$tool")
        fi
    done
    if [[ "${#missing_tools[@]}" -gt 0 ]]; then
        echo "error: missing tools: ${missing_tools[*]}" >&2
        exit 1
    fi
}

_argc_die() {
    if [[ $# -eq 0 ]]; then
        cat
    else
        echo "$*" >&2
    fi
    exit 1
}

_argc_run "$@"

# ARGC-BUILD }