#!/bin/bash
# Tavily 搜索

read -r input
QUERY=$(echo "$input" | python3 -c "
import sys, json
d = json.load(sys.stdin)
print(d.get('query', ''))
" 2>/dev/null)

COUNT=$(echo "$input" | python3 -c "
import sys, json
d = json.load(sys.stdin)
print(d.get('count', 5))
" 2>/dev/null)

if [ -z "$QUERY" ]; then
  echo "请提供 query 参数"
  exit 1
fi

API_KEY="${TAVILY_API_KEY}"
if [ -z "$API_KEY" ]; then
  echo "请设置 TAVILY_API_KEY 环境变量"
  exit 1
fi

RESP=$(curl -s -X POST "https://api.tavily.com/search" \
  -H "Content-Type: application/json" \
  -d "{\"api_key\":\"$API_KEY\",\"query\":\"$QUERY\",\"max_results\":$COUNT,\"include_answer\":true}")

# 解析并输出
ANSWER=$(echo "$RESP" | python3 -c "
import sys, json
d = json.load(sys.stdin)
answer = d.get('answer', '')
results = d.get('results', [])
lines = []
if answer:
    lines.append(f'总结: {answer}')
    lines.append('')
for i, r in enumerate(results[:$COUNT], 1):
    title = r.get('title', '无标题')
    url = r.get('url', '')
    content = r.get('content', '')[:200]
    lines.append(f'{i}. {title}')
    lines.append(f'   {url}')
    if content:
        lines.append(f'   {content}')
    lines.append('')
print('\\n'.join(lines))
" 2>/dev/null)

echo "$ANSWER"