#!/usr/bin/env bash
# update-model.sh — Update model configuration in openclaw state file
# Usage:
#   update-model.sh --primary <model-id>
#   update-model.sh --fallbacks <model1,model2,model3>
#   update-model.sh --primary <model-id> --fallbacks <model1,model2>
#
# All model IDs are validated against /v1/models before writing.
# A backup of the current config is created before any changes.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

PRIMARY=""
FALLBACKS=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --primary)
      PRIMARY="$2"
      shift 2
      ;;
    --fallbacks)
      FALLBACKS="$2"
      shift 2
      ;;
    --help|-h)
      echo "Usage: update-model.sh [--primary <model-id>] [--fallbacks <model1,model2,...>]"
      echo ""
      echo "Options:"
      echo "  --primary   Set the primary model (will be prefixed with llm-proxy/)"
      echo "  --fallbacks Comma-separated list of fallback models (min 2 required)"
      echo ""
      echo "All model IDs are validated against /v1/models before writing."
      exit 0
      ;;
    *)
      echo "ERROR: Unknown argument: $1" >&2
      exit 1
      ;;
  esac
done

if [[ -z "$PRIMARY" && -z "$FALLBACKS" ]]; then
  echo "ERROR: Must specify --primary and/or --fallbacks" >&2
  exit 1
fi

# Find the state file
STATE_FILE="${OPENCLAW_STATE_DIR:-$HOME/.openclaw/state}/openclaw.json"
if [[ ! -f "$STATE_FILE" ]]; then
  for alt in \
    "/opt/openclaw/state/openclaw.json" \
    "${OPENCLAW_CONFIG_DIR:-$HOME/.openclaw/config}/openclaw.json"; do
    if [[ -f "$alt" ]]; then
      STATE_FILE="$alt"
      break
    fi
  done
fi

if [[ ! -f "$STATE_FILE" ]]; then
  echo "ERROR: Cannot find openclaw.json state file" >&2
  exit 1
fi

echo "📁 Config file: $STATE_FILE"

# Validate all model IDs first
echo ""
echo "🔍 Validating model IDs against /v1/models..."
VALIDATION_FAILED=0

if [[ -n "$PRIMARY" ]]; then
  # Strip llm-proxy/ prefix for validation
  PRIMARY_CLEAN="${PRIMARY#llm-proxy/}"
  if ! "$SCRIPT_DIR/validate-model.sh" "$PRIMARY_CLEAN" 2>&1; then
    VALIDATION_FAILED=1
  fi
fi

if [[ -n "$FALLBACKS" ]]; then
  IFS=',' read -ra FB_ARRAY <<< "$FALLBACKS"
  
  if [[ ${#FB_ARRAY[@]} -lt 2 ]]; then
    echo "❌ ERROR: Minimum 2 fallback models required (got ${#FB_ARRAY[@]})" >&2
    VALIDATION_FAILED=1
  fi
  
  for fb in "${FB_ARRAY[@]}"; do
    fb_clean="${fb#llm-proxy/}"
    fb_clean="$(echo "$fb_clean" | xargs)"  # trim whitespace
    if ! "$SCRIPT_DIR/validate-model.sh" "$fb_clean" 2>&1; then
      VALIDATION_FAILED=1
    fi
  done
fi

if [[ $VALIDATION_FAILED -ne 0 ]]; then
  echo ""
  echo "❌ Validation failed. No changes made." >&2
  exit 1
fi

# Create backup
BACKUP="${STATE_FILE}.backup.$(date +%Y%m%d-%H%M%S)"
cp "$STATE_FILE" "$BACKUP"
echo ""
echo "💾 Backup saved: $BACKUP"

# Apply changes using Python for safe JSON manipulation
python3 -c "
import json, sys, re

state_file = '$STATE_FILE'
primary = '${PRIMARY}'.strip() or None
fallbacks_raw = '${FALLBACKS}'.strip() or None

# Read and parse (handle JSON5 comments)
with open(state_file, 'r') as f:
    content = f.read()

# Strip comments for parsing
lines = content.split('\n')
cleaned = []
for line in lines:
    s = line.lstrip()
    if s.startswith('//'):
        continue
    in_string = False
    result = []
    i = 0
    while i < len(line):
        c = line[i]
        if c == '\"' and (i == 0 or line[i-1] != '\\\\'):
            in_string = not in_string
        elif c == '/' and i + 1 < len(line) and line[i+1] == '/' and not in_string:
            break
        result.append(c)
        i += 1
    cleaned.append(''.join(result))

# Remove trailing commas before } or ] (JSON5 feature)
json_str = '\n'.join(cleaned)
json_str = re.sub(r',\s*([}\]])', r'\1', json_str)

try:
    cfg = json.loads(json_str)
except json.JSONDecodeError:
    try:
        cfg = json.loads(content)
    except json.JSONDecodeError as e:
        print(f'ERROR: Failed to parse config: {e}', file=sys.stderr)
        sys.exit(1)

# Ensure path exists
if 'agents' not in cfg:
    cfg['agents'] = {}
if 'defaults' not in cfg['agents']:
    cfg['agents']['defaults'] = {}
if 'model' not in cfg['agents']['defaults']:
    cfg['agents']['defaults']['model'] = {}

model = cfg['agents']['defaults']['model']
if isinstance(model, str):
    model = {'primary': model}
    cfg['agents']['defaults']['model'] = model

old_primary = model.get('primary', '(none)')
old_fallbacks = model.get('fallbacks', [])

# Apply primary
if primary:
    # Ensure llm-proxy/ prefix
    if not primary.startswith('llm-proxy/'):
        primary = f'llm-proxy/{primary}'
    model['primary'] = primary

# Apply fallbacks
if fallbacks_raw:
    fbs = [fb.strip() for fb in fallbacks_raw.split(',') if fb.strip()]
    fbs = [f'llm-proxy/{fb}' if not fb.startswith('llm-proxy/') else fb for fb in fbs]
    model['fallbacks'] = fbs

# Remove per-agent model overrides that match the old primary
# (they were likely set by the same drift that caused the issue)
agent_list = cfg.get('agents', {}).get('list', [])
removed_overrides = []
for agent in agent_list:
    if 'model' in agent:
        removed_overrides.append((agent.get('id', '?'), agent['model']))
        del agent['model']

# Write back
with open(state_file, 'w') as f:
    json.dump(cfg, f, indent=2)
    f.write('\n')

# Print summary
print()
print('✅ Configuration updated:')
print()
print(f'  Primary: {old_primary} → {model.get(\"primary\", \"(none)\")}')
print(f'  Fallbacks:')
for i, fb in enumerate(model.get('fallbacks', []), 1):
    old_marker = '' if fb in old_fallbacks else ' (new)'
    print(f'    {i}. {fb}{old_marker}')
if removed_overrides:
    print()
    print('  🧹 Cleared per-agent model overrides:')
    for aid, amodel in removed_overrides:
        print(f'    {aid}: {amodel} → (uses default)')
" 2>&1

echo ""
echo "Done. Restart OpenClaw for changes to take effect."