#!/usr/bin/env bash

set -euo pipefail

URL="http://127.0.0.1:8000/mcp"
TOKEN=""
PRETTY=true
MODE="mcp"   # mcp | openapi

usage() {
cat <<EOF
Usage:

  MCP — Discovery
    $0 --token TOKEN list
  MCP — Discovery (custom URL)
    $0 --url URL --token TOKEN list
  MCP — Tool call
    $0 --token TOKEN call TOOL_NAME key=value ...

  OpenAPI — Discovery
    $0 --openapi --url URL --token TOKEN list
  OpenAPI — Tool call
    $0 --openapi --url URL --token TOKEN call OPERATION_ID key=value ...

Options:

  --url URL
      Endpoint base URL
      Default (MCP): http://127.0.0.1:8000/mcp
  --token TOKEN
      Bearer token (required)
  --openapi
      Switch to OpenAPI mode (fetches /openapi.json, calls operations directly)
  --brut
      No pretty-print — raw JSON output
  -h, --help
      Show help

Examples:

  MCP — Discover tools
    $0 --token xxx list
  MCP — Call a tool
    $0 --token xxx call get_weather city=Paris

  OpenAPI — Discover operations
    $0 --openapi --url https://api.example.com --token xxx list
  OpenAPI — Call an operation
    $0 --openapi --url https://api.example.com --token xxx call get_weather city=Paris
  Brut output
    $0 --openapi --url https://api.example.com --token xxx --brut call get_weather city=Paris

EOF
}

pretty_print() {
    if [[ "$PRETTY" != true ]]; then
        cat
        return
    fi

    if ! command -v jq >/dev/null 2>&1; then
        echo "ERROR: jq is required for pretty-print" >&2
        exit 1
    fi

    local RESPONSE
    RESPONSE=$(cat)

    # MCP SSE responses carry a "data: " prefix
    local JSON
    JSON=$(echo "$RESPONSE" | sed -n 's/^data: //p')

    if [[ -z "$JSON" ]]; then
        # Plain JSON (OpenAPI responses) — just pretty-print as-is
        echo "$RESPONSE" | jq '.'
        return
    fi

    echo "$JSON" | jq '
        if .result.structuredContent then
            .result.structuredContent
        elif .result then
            .result
        else
            .
        end
    '
}

# ---------------------------------------------------------------------------
# MCP helpers
# ---------------------------------------------------------------------------

create_session() {
    local headers
    headers=$(mktemp)

    curl -s \
      -X POST "$URL" \
      -H "Content-Type: application/json" \
      -H "Accept: application/json, text/event-stream" \
      -H "Authorization: Bearer $TOKEN" \
      -D "$headers" \
      -d '{
            "jsonrpc":"2.0",
            "method":"initialize",
            "id":0,
            "params":{
              "protocolVersion":"2024-11-05",
              "capabilities":{},
              "clientInfo":{
                "name":"mcp.sh",
                "version":"1.0"
              }
            }
          }' >/dev/null

    SESSION=$(
        grep -i '^mcp-session-id:' "$headers" |
        awk '{print $2}' |
        tr -d '\r'
    )

    rm -f "$headers"

    if [[ -z "${SESSION:-}" ]]; then
        echo "ERROR: unable to retrieve MCP session id" >&2
        exit 1
    fi
}

mcp_list() {
    create_session

    curl -s \
      -X POST "$URL" \
      -H "Content-Type: application/json" \
      -H "Accept: application/json, text/event-stream" \
      -H "Authorization: Bearer $TOKEN" \
      -H "mcp-session-id: $SESSION" \
      -d '{
            "jsonrpc":"2.0",
            "method":"tools/list",
            "id":1
          }' \
    | pretty_print

    echo
}

mcp_call() {
    local TOOL_NAME="$1"; shift

    create_session

    local ARGS_JSON
    ARGS_JSON=$(build_args_json "$@")

    local PAYLOAD
    PAYLOAD=$(cat <<EOF
{
  "jsonrpc":"2.0",
  "method":"tools/call",
  "id":2,
  "params":{
    "name":"$TOOL_NAME",
    "arguments":$ARGS_JSON
  }
}
EOF
)

    curl -s \
      -X POST "$URL" \
      -H "Content-Type: application/json" \
      -H "Accept: application/json, text/event-stream" \
      -H "Authorization: Bearer $TOKEN" \
      -H "mcp-session-id: $SESSION" \
      -d "$PAYLOAD" \
    | pretty_print

    echo
}

# ---------------------------------------------------------------------------
# OpenAPI helpers
# ---------------------------------------------------------------------------

# Fetch and cache the OpenAPI spec (stored in a temp file for the session)
OPENAPI_SPEC_FILE=""

fetch_openapi_spec() {
    if [[ -n "$OPENAPI_SPEC_FILE" && -f "$OPENAPI_SPEC_FILE" ]]; then
        return
    fi

    if ! command -v jq >/dev/null 2>&1; then
        echo "ERROR: jq is required for OpenAPI mode" >&2
        exit 1
    fi

    OPENAPI_SPEC_FILE=$(mktemp)
    # Try /openapi.json first, then /openapi.yaml (json only — yaml requires extra tooling)
    local BASE_URL="${URL%/}"   # strip trailing slash

    local HTTP_CODE
    HTTP_CODE=$(curl -s -o "$OPENAPI_SPEC_FILE" -w "%{http_code}" \
      -H "Authorization: Bearer $TOKEN" \
      "${BASE_URL}/openapi.json")

    if [[ "$HTTP_CODE" != "200" ]]; then
        echo "ERROR: could not fetch OpenAPI spec from ${BASE_URL}/openapi.json (HTTP $HTTP_CODE)" >&2
        rm -f "$OPENAPI_SPEC_FILE"
        exit 1
    fi

    # Validate it looks like JSON
    if ! jq empty "$OPENAPI_SPEC_FILE" 2>/dev/null; then
        echo "ERROR: ${BASE_URL}/openapi.json did not return valid JSON" >&2
        rm -f "$OPENAPI_SPEC_FILE"
        exit 1
    fi
}

openapi_list() {
    fetch_openapi_spec

    # Résout un $ref "#/components/schemas/Foo" → l'objet schema correspondant
    # Extrait les propriétés du body en résolvant les $ref si nécessaire
    jq -r '
        . as $root |
        # Fonction de résolution de $ref
        def resolve_ref(schema):
            if schema | has("$ref") then
                schema["$ref"] |
                ltrimstr("#/") |
                split("/") |
                reduce .[] as $part ($root; .[$part])
            else
                schema
            end;

        .paths // {} |
        to_entries[] |
        .key as $path |
        .value | to_entries[] |
        select(.key | test("^(get|post|put|patch|delete|options|head)$"; "i")) |
        .key as $method |
        .value |
        . as $op |
        {
            operationId: (.operationId // ($method + " " + $path)),
            method: ($method | ascii_upcase),
            path: $path,
            summary: (.summary // "-"),
            parameters: (
                # Paramètres de chemin/query classiques
                [
                    ($op.parameters // [])[] |
                    .name + "=" + (.schema.type // "string") +
                    (if .required then " (required)" else "" end)
                ]
                +
                # Paramètres du requestBody (résolution du $ref si besoin)
                (
                    if $op.requestBody then
                        ($op.requestBody.content | to_entries[0].value.schema) |
                        resolve_ref(.) |
                        . as $resolved_schema |
                        ($resolved_schema.properties // {}) |
                        to_entries[] |
                        .key as $pname |
                        [
                            $pname + "=" + (.value.type // "?") +
                            (if (.value.default != null) then " [default: " + (.value.default | tostring) + "]" else "" end) +
                            (if ($resolved_schema.required // [] | contains([$pname])) then " (required)" else "" end)
                        ]
                    else
                        []
                    end
                )
            )
        }
    ' "$OPENAPI_SPEC_FILE" \
    | if [[ "$PRETTY" == true ]]; then
          jq -r '"─────────────────────────────────────────────\n" +
                 "operationId : " + .operationId + "\n" +
                 "method      : " + .method + "  " + .path + "\n" +
                 "summary     : " + .summary + "\n" +
                 "parameters  : " + (if (.parameters | length) > 0 then (.parameters | join(", ")) else "-" end)'
      else
          cat
      fi

    echo
}

openapi_call() {
    local OPERATION_ID="$1"; shift

    fetch_openapi_spec

    # Resolve the operation: find path + method + parameter locations
    # Résout les $ref pour les bodyParams
    local OP_INFO
    OP_INFO=$(jq -r --arg op "$OPERATION_ID" '
        . as $root |
        def resolve_ref(schema):
            if schema | has("$ref") then
                schema["$ref"] |
                ltrimstr("#/") |
                split("/") |
                reduce .[] as $part ($root; .[$part])
            else
                schema
            end;

        .paths // {} |
        to_entries[] |
        .key as $path |
        .value | to_entries[] |
        select(.key | test("^(get|post|put|patch|delete|options|head)$"; "i")) |
        .key as $method |
        .value |
        select((.operationId // "") == $op) |
        {
            method: ($method | ascii_upcase),
            path: $path,
            queryParams: [(.parameters // [])[] | select(.in == "query") | .name],
            pathParams:  [(.parameters // [])[] | select(.in == "path")  | .name],
            bodyParams:  (
                if .requestBody then
                    (.requestBody.content | to_entries[0].value.schema) |
                    resolve_ref(.) |
                    .properties // {} | keys
                else []
                end
            ),
            contentType: (
                if .requestBody then
                    (.requestBody.content | to_entries[0].key)
                else "application/json"
                end
            )
        }
    ' "$OPENAPI_SPEC_FILE")

    if [[ -z "$OP_INFO" ]]; then
        echo "ERROR: operationId '$OPERATION_ID' not found in OpenAPI spec" >&2
        exit 1
    fi

    local METHOD PATH_TMPL
    METHOD=$(echo "$OP_INFO"    | jq -r '.method')
    PATH_TMPL=$(echo "$OP_INFO" | jq -r '.path')
    local CONTENT_TYPE
    CONTENT_TYPE=$(echo "$OP_INFO" | jq -r '.contentType')

    # Parse key=value arguments
    declare -A KV
    for ARG in "$@"; do
        local KEY="${ARG%%=*}"
        local VALUE="${ARG#*=}"
        KV["$KEY"]="$VALUE"
    done

    # 1. Substitute path parameters
    local FINAL_PATH="$PATH_TMPL"
    while IFS= read -r PNAME; do
        [[ -z "$PNAME" ]] && continue
        local PVAL="${KV[$PNAME]:-}"
        if [[ -z "$PVAL" ]]; then
            echo "ERROR: missing required path parameter '$PNAME'" >&2
            exit 1
        fi
        FINAL_PATH="${FINAL_PATH//\{$PNAME\}/$PVAL}"
        unset "KV[$PNAME]"
    done < <(echo "$OP_INFO" | jq -r '.pathParams[]')

    # 2. Build query string
    local QUERY_STRING=""
    while IFS= read -r QNAME; do
        [[ -z "$QNAME" ]] && continue
        local QVAL="${KV[$QNAME]:-}"
        if [[ -n "$QVAL" ]]; then
            QUERY_STRING+="${QUERY_STRING:+&}${QNAME}=$(python3 -c "import urllib.parse,sys; print(urllib.parse.quote(sys.argv[1]))" "$QVAL" 2>/dev/null || echo "$QVAL")"
            unset "KV[$QNAME]"
        fi
    done < <(echo "$OP_INFO" | jq -r '.queryParams[]')

    # 3. Remaining KV pairs → request body
    local BODY_JSON="{}"
    if [[ ${#KV[@]} -gt 0 ]]; then
        BODY_JSON=$(build_args_json_from_assoc)
    fi

    # 4. Assemble final URL
    local BASE_URL="${URL%/}"
    local FULL_URL="${BASE_URL}${FINAL_PATH}"
    [[ -n "$QUERY_STRING" ]] && FULL_URL+="?${QUERY_STRING}"

    # 5. Fire the request
    local CURL_ARGS=(
        -s
        -X "$METHOD"
        -H "Authorization: Bearer $TOKEN"
        -H "Accept: application/json"
    )

    if [[ "$METHOD" != "GET" && "$METHOD" != "HEAD" ]]; then
        CURL_ARGS+=(
            -H "Content-Type: ${CONTENT_TYPE}"
            -d "$BODY_JSON"
        )
    fi

    curl "${CURL_ARGS[@]}" "$FULL_URL" | pretty_print

    echo
}

# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

build_args_json() {
    local ARGS_JSON="{"
    local FIRST=true

    for ARG in "$@"; do
        local KEY="${ARG%%=*}"
        local VALUE="${ARG#*=}"

        if [[ "$FIRST" == true ]]; then
            FIRST=false
        else
            ARGS_JSON+=","
        fi

        local ESCAPED_VALUE
        ESCAPED_VALUE=$(printf '%s' "$VALUE" | sed 's/"/\\"/g')
        ARGS_JSON+="\"$KEY\":\"$ESCAPED_VALUE\""
    done

    ARGS_JSON+="}"
    echo "$ARGS_JSON"
}

# Build JSON body from the associative array KV (bash 4+)
build_args_json_from_assoc() {
    local JSON="{"
    local FIRST=true

    for KEY in "${!KV[@]}"; do
        local VALUE="${KV[$KEY]}"
        if [[ "$FIRST" == true ]]; then
            FIRST=false
        else
            JSON+=","
        fi
        local ESCAPED_VALUE
        ESCAPED_VALUE=$(printf '%s' "$VALUE" | sed 's/"/\\"/g')
        JSON+="\"$KEY\":\"$ESCAPED_VALUE\""
    done

    JSON+="}"
    echo "$JSON"
}

# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

COMMAND=""

while [[ $# -gt 0 ]]; do
    case "$1" in

        --url)
            URL="$2"
            shift 2
            ;;

        --token)
            TOKEN="$2"
            shift 2
            ;;

        --openapi)
            MODE="openapi"
            shift
            ;;

        --brut)
            PRETTY=false
            shift
            ;;

        list|call)
            COMMAND="$1"
            shift
            break
            ;;

        -h|--help)
            usage
            exit 0
            ;;

        *)
            echo "Unknown argument: $1" >&2
            usage
            exit 1
            ;;
    esac
done

if [[ -z "$TOKEN" ]]; then
    echo "ERROR: --token is required" >&2
    exit 1
fi

if [[ -z "$COMMAND" ]]; then
    usage
    exit 1
fi

# ---------------------------------------------------------------------------
# Dispatch
# ---------------------------------------------------------------------------

if [[ "$MODE" == "openapi" ]]; then

    if [[ "$COMMAND" == "list" ]]; then
        openapi_list
        exit 0
    fi

    if [[ "$COMMAND" == "call" ]]; then
        TOOL_NAME="${1:-}"
        if [[ -z "$TOOL_NAME" ]]; then
            echo "ERROR: missing operation id" >&2
            exit 1
        fi
        shift
        openapi_call "$TOOL_NAME" "$@"
        exit 0
    fi

else  # MCP mode

    if [[ "$COMMAND" == "list" ]]; then
        mcp_list
        exit 0
    fi

    if [[ "$COMMAND" == "call" ]]; then
        TOOL_NAME="${1:-}"
        if [[ -z "$TOOL_NAME" ]]; then
            echo "ERROR: missing tool name" >&2
            exit 1
        fi
        shift
        mcp_call "$TOOL_NAME" "$@"
        exit 0
    fi

fi

echo "ERROR: unsupported command" >&2
exit 1