#!/usr/bin/env python3
"""
Tavily AI Search API Client
Usage: python tavily_search.py "your search query" [--max-results 5] [--depth basic|advanced]
"""

import os
import sys
import json
import argparse
from typing import List, Dict, Any, Optional


def get_api_key(api_key: Optional[str] = None) -> str:
    """Get Tavily API key from environment variable or parameter."""
    # Priority: parameter > environment variable
    if api_key:
        return api_key
    
    env_key = os.environ.get("TAVILY_API_KEY")
    if env_key:
        return env_key
    
    raise ValueError(
        "Tavily API key required. Set via:\n"
        "  1. Environment variable: export TAVILY_API_KEY=your_key\n"
        "     (Add to ~/.bashrc or ~/.zshrc for persistence)\n"
        "  2. Direct parameter: pass api_key when calling the function\n"
        "\nGet your API key at: https://tavily.com"
    )


def tavily_search(
    query: str,
    max_results: int = 5,
    search_depth: str = "basic",
    include_answer: bool = False,
    include_images: bool = False,
    api_key: Optional[str] = None
) -> Dict[str, Any]:
    """
    Search using Tavily AI Search API.
    
    Args:
        query: Search query string
        max_results: Number of results to return (1-20)
        search_depth: "basic" or "advanced"
        include_answer: Include AI-generated answer
        include_images: Include image URLs
        api_key: Tavily API key (optional, auto-detected from env/config)
    
    Returns:
        Dictionary containing search results
    """
    api_key = get_api_key(api_key)
    
    try:
        import requests
    except ImportError:
        raise ImportError("requests package required. Install with: pip install requests")
    
    url = "https://api.tavily.com/search"
    
    payload = {
        "api_key": api_key,
        "query": query,
        "max_results": min(max(max_results, 1), 20),
        "search_depth": search_depth,
        "include_answer": include_answer,
        "include_images": include_images,
    }
    
    response = requests.post(url, json=payload, timeout=30)
    response.raise_for_status()
    
    return response.json()


def format_results(results: Dict[str, Any]) -> str:
    """Format search results for display."""
    output = []
    
    if "answer" in results and results["answer"]:
        output.append("=" * 60)
        output.append("AI ANSWER")
        output.append("=" * 60)
        output.append(results["answer"])
        output.append("")
    
    output.append("=" * 60)
    output.append("SEARCH RESULTS")
    output.append("=" * 60)
    output.append(f"Query: {results.get('query', 'N/A')}")
    output.append("")
    
    for i, result in enumerate(results.get("results", []), 1):
        output.append(f"{i}. {result.get('title', 'No title')}")
        output.append(f"   URL: {result.get('url', 'N/A')}")
        
        if result.get('published_date'):
            output.append(f"   Published: {result['published_date']}")
        
        if result.get('score'):
            output.append(f"   Relevance: {result['score']:.2f}")
        
        content = result.get('content', '')
        if content:
            # Truncate long content
            if len(content) > 300:
                content = content[:297] + "..."
            output.append(f"   {content}")
        
        output.append("")
    
    return "\n".join(output)


def main():
    parser = argparse.ArgumentParser(description="Tavily AI Search")
    parser.add_argument("query", help="Search query")
    parser.add_argument("--max-results", type=int, default=5, help="Number of results (1-20)")
    parser.add_argument("--depth", choices=["basic", "advanced"], default="basic", help="Search depth")
    parser.add_argument("--answer", action="store_true", help="Include AI-generated answer")
    parser.add_argument("--images", action="store_true", help="Include images")
    parser.add_argument("--json", action="store_true", help="Output raw JSON")
    
    args = parser.parse_args()
    
    try:
        results = tavily_search(
            query=args.query,
            max_results=args.max_results,
            search_depth=args.depth,
            include_answer=args.answer,
            include_images=args.images
        )
        
        if args.json:
            print(json.dumps(results, indent=2, ensure_ascii=False))
        else:
            print(format_results(results))
    
    except ValueError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Search failed: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()