#!/usr/bin/env python3
"""
Tavily 搜索示例
需要 API Key: https://tavily.com/
"""

import os
from tavily import TavilyClient

# 配置
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY", "")

def tavily_search(query, include_raw_content=False):
    """使用 Tavily 搜索"""
    if not TAVILY_API_KEY:
        print("❌ 请设置 TAVILY_API_KEY 环境变量")
        print("   export TAVILY_API_KEY='你的API Key'")
        return None
    
    client = TavilyClient(api_key=TAVILY_API_KEY)
    
    try:
        # 搜索
        result = client.search(query, include_raw_content=include_raw_content)
        return result
    except Exception as e:
        print(f"❌ 搜索失败: {e}")
        return None

def tavily_extract(urls):
    """提取网页内容"""
    if not TAVILY_API_KEY:
        print("❌ 请设置 TAVILY_API_KEY")
        return None
    
    client = TavilyClient(api_key=TAVILY_API_KEY)
    
    try:
        result = client.extract(urls)
        return result
    except Exception as e:
        print(f"❌ 提取失败: {e}")
        return None

def print_results(result):
    """打印搜索结果"""
    if not result:
        return
    
    print(f"\n📊 搜索结果")
    print(f"查询: {result.get('query', '')}")
    print(f"耗时: {result.get('response_time', 'N/A')}秒")
    print(f"结果数: {len(result.get('results', []))}")
    print("-" * 60)
    
    for i, r in enumerate(result.get('results', [])[:5], 1):
        print(f"\n{i}. {r.get('title', 'No title')}")
        print(f"   链接: {r.get('url', '')}")
        print(f"   摘要: {r.get('content', '')[:150]}...")

def main():
    import sys
    
    print(f"\n{'='*60}")
    print("🔍 Tavily AI 搜索")
    print(f"{'='*60}\n")
    
    if not TAVILY_API_KEY:
        print("⚠️ 未设置 API Key")
        print("请设置: export TAVILY_API_KEY='你的API Key'\n")
    
    if len(sys.argv) < 2:
        print("用法:")
        print("  python tavily_example.py \"搜索关键词\"")
        print("  python tavily_example.py \"关键词\" --extract")
        print("\n选项:")
        print("  --extract, -e  提取网页完整内容")
        return
    
    query = " ".join(sys.argv[1:])
    extract = "--extract" in sys.argv or "-e" in sys.argv
    
    if extract:
        print(f"🔍 搜索并提取: {query}")
    else:
        print(f"🔍 搜索: {query}")
    
    if extract:
        # 先搜索获取链接
        result = tavily_search(query)
        if result:
            urls = [r['url'] for r in result.get('results', [])[:3]]
            if urls:
                print(f"\n📄 提取内容...")
                extract_result = tavily_extract(urls)
                if extract_result:
                    for content in extract_result.get('content', []):
                        print(f"\n{'='*60}")
                        print(f"来源: {content.get('url', '')}")
                        print(f"{'='*60}")
                        print(content.get('content', '')[:1000])
    else:
        result = tavily_search(query)
        print_results(result)

if __name__ == "__main__":
    main()