#!/bin/bash

# ask - Simple OpenRouter API CLI tool
# Usage: ask [OPTIONS] [PROMPT]

set -euo pipefail

# Check for API key
if [ -z "${OPENROUTER_API_KEY:-}" ]; then
    echo "Error: OPENROUTER_API_KEY environment variable is not set" >&2
    exit 1
fi

# Model shortcuts function
get_model() {
    case "$1" in
        c) echo "inception/mercury-coder:nitro" ;;
        g) echo "google/gemini-2.5-flash:nitro" ;;
        s) echo "anthropic/claude-sonnet-4:nitro" ;;
        k) echo "moonshotai/kimi-k2:nitro" ;;
        q) echo "qwen/qwen3-235b-a22b-2507:nitro" ;;
    esac
}

# Default values
MODEL="inception/mercury-coder:nitro"
SYSTEM_PROMPT=""
PROMPT=""
STREAMING=false
NO_SYSTEM=false
PROVIDER_ORDER=""

# Default system prompt (direct answers)
DEFAULT_PROMPT="You are a direct answer engine. Output ONLY the requested information.

For commands: Output executable syntax only. No explanations, no comments.
For questions: Output the answer only. No context, no elaboration.

Rules:
- If asked for a command, provide ONLY the command
- If asked a question, provide ONLY the answer
- Never include markdown formatting or code blocks
- Never add explanatory text before or after
- Assume output will be piped or executed directly
- For multi-step commands, use && or ; to chain them
- Make commands robust and handle edge cases silently"

# Function to show help
show_help() {
    cat << EOF
ask - Query AI models via OpenRouter API

Usage: ask [OPTIONS] [PROMPT]

Options:
  -c          Use inception/mercury-coder (default)
  -g          Use google/gemini-2.5-flash
  -s          Use anthropic/claude-sonnet-4
  -k          Use moonshotai/kimi-k2
  -q          Use qwen/qwen3-235b-a22b-2507
  -m MODEL    Use custom model
  -r          Disable system prompt (raw model behavior)
  --stream    Enable streaming output
  --system    Set system prompt for the conversation
  --provider  Comma-separated list of providers for routing
  -h, --help  Show this help message

Examples:
  ask "Write a hello world in Python"
  ask -g "Explain quantum computing"
  ask -m openai/gpt-4o "What is 2+2?"
  echo "Fix this code" | ask
  ask --system "You are a pirate" "Tell me about sailing"

EOF
    exit 0
}

# Parse command line arguments
while [ $# -gt 0 ]; do
    case "$1" in
        -h|--help) show_help ;;
        -[cgskq])
            MODEL="$(get_model "${1:1}")"
            shift ;;
        -m)
            MODEL="${2:?Error: -m requires a model name}"
            shift 2 ;;
        -r)
            NO_SYSTEM=true
            shift ;;
        --stream)
            STREAMING=true
            shift ;;
        --system)
            SYSTEM_PROMPT="${2:?Error: --system requires a prompt}"
            shift 2 ;;
        --provider)
            PROVIDER_ORDER="${2:?Error: --provider requires providers}"
            shift 2 ;;
        *)
            PROMPT="$*"
            break ;;
    esac
done

# If no prompt provided as argument, read from stdin
if [ -z "$PROMPT" ]; then
    if [ -t 0 ]; then
        echo "Error: No prompt provided. Use 'ask -h' for help." >&2
        exit 1
    fi
    # Read all stdin, preserving multi-line format
    PROMPT=$(cat)
fi

# Apply default system prompt unless disabled or custom prompt provided
if [ "$NO_SYSTEM" = false ] && [ -z "$SYSTEM_PROMPT" ]; then
    SYSTEM_PROMPT="$DEFAULT_PROMPT"
fi

# Build messages array with proper JSON escaping
if [ -n "$SYSTEM_PROMPT" ]; then
    MESSAGES='[{"role":"system","content":'"$(printf '%s' "$SYSTEM_PROMPT" | jq -Rs .)"'},{"role":"user","content":'"$(printf '%s' "$PROMPT" | jq -Rs .)"'}]'
else
    MESSAGES='[{"role":"user","content":'"$(printf '%s' "$PROMPT" | jq -Rs .)"'}]'
fi

# Record start time
START_TIME=$(date +%s.%N)

# Build JSON payload once
PROVIDER_JSON=""
if [ -n "$PROVIDER_ORDER" ]; then
    PROVIDER_JSON=',"provider":{"order":['$(echo "$PROVIDER_ORDER" | awk -F, '{for(i=1;i<=NF;i++) printf "\"%s\"%s", $i, (i<NF?",":"")}')']}'
fi

JSON_PAYLOAD='{
    "model": "'"$MODEL"'",
    "messages": '"$MESSAGES"',
    "stream": '$([ "$STREAMING" = true ] && echo true || echo false)"$PROVIDER_JSON"'
}'

API_URL="https://openrouter.ai/api/v1/chat/completions"

# Add newline before answer
echo

# Make API request
if [ "$STREAMING" = true ]; then
    # Streaming mode
    curl -sS "$API_URL" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer $OPENROUTER_API_KEY" \
      -d "$JSON_PAYLOAD" 2>&1 | while IFS= read -r line; do
        # Check for errors
        if echo "$line" | grep -q '"error"'; then
            echo "Error: $(echo "$line" | jq -r '.error.message // .error // "Unknown error"')" >&2
            exit 1
        fi

        # Process SSE data lines
        if [[ "$line" == data:* ]]; then
            json="${line#data: }"
            [ "$json" = "" ] || [ "$json" = "[DONE]" ] && continue

            content=$(echo "$json" | jq -r '.choices[0].delta.content // ""' 2>/dev/null)
            [ -n "$content" ] && printf '%s' "$content"
        fi
    done
    echo

    # Show metadata
    ELAPSED=$(printf "%.2f" $(echo "$(date +%s.%N) - $START_TIME" | bc))
    echo
    echo "[$MODEL - ${ELAPSED}s]" >&2
else
    # Non-streaming mode
    response=$(curl -sS "$API_URL" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer $OPENROUTER_API_KEY" \
      -d "$JSON_PAYLOAD" 2>&1)

    # Check for errors
    if echo "$response" | grep -q '"error"'; then
        echo "Error: $(echo "$response" | jq -r '.error.message // .error // "Unknown error"')" >&2
        exit 1
    fi

    # Extract and print content
    echo "$response" | jq -r '.choices[0].message.content // "No response received"'

    # Show metadata
    ELAPSED=$(printf "%.2f" $(echo "$(date +%s.%N) - $START_TIME" | bc))
    TOKENS=$(echo "$response" | jq -r '.usage.completion_tokens // 0')
    PROVIDER=$(echo "$response" | jq -r '.provider // "Unknown"')
    TPS=$(echo "scale=1; $TOKENS / $ELAPSED" | bc 2>/dev/null || echo "0.0")

    echo
    echo "[$MODEL via $PROVIDER - ${ELAPSED}s - ${TPS} tok/s]" >&2
fi