#!/usr/bin/env python3
"""
Perplexity API Query Interface

Usage:
    python3 query.py "What is the capital of France?"
    python3 query.py "Latest AI news" --model sonar-pro --citations
"""

import json
import os
import sys
import urllib.request
from pathlib import Path

def load_config():
    """Load API configuration"""
    config_path = Path(__file__).parent.parent / "config.json"
    try:
        with open(config_path) as f:
            return json.load(f)
    except Exception as e:
        print(f"Error loading config: {e}", file=sys.stderr)
        return None

def query_perplexity(query, model=None, max_tokens=None, include_citations=False, search_context="low"):
    """
    Query Perplexity API
    
    Args:
        query: The question/prompt to send
        model: Model to use (sonar, sonar-pro, sonar-reasoning, sonar-deep-research)
        max_tokens: Maximum tokens in response
        include_citations: Whether to include source citations
        search_context: Search depth (low, medium, high)
    
    Returns:
        dict with response text, citations, and usage info
    """
    config = load_config()
    if not config:
        return {"error": "Failed to load configuration"}
    
    model = model or config.get("default_model", "sonar")
    max_tokens = max_tokens or config.get("default_max_tokens", 1000)
    api_key = config.get("api_key")
    base_url = config.get("base_url", "https://api.perplexity.ai")
    
    if not api_key:
        return {"error": "API key not configured"}
    
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": "Be precise and concise."},
            {"role": "user", "content": query}
        ],
        "max_tokens": max_tokens,
        "search_context_size": search_context
    }
    
    data = json.dumps(payload).encode()
    
    req = urllib.request.Request(
        f"{base_url}/chat/completions",
        data=data,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
    )
    
    try:
        with urllib.request.urlopen(req, timeout=60) as response:
            result = json.loads(response.read().decode())
            
            output = {
                "text": result["choices"][0]["message"]["content"],
                "model": result.get("model"),
                "usage": result.get("usage", {})
            }
            
            if include_citations:
                output["citations"] = result.get("citations", [])
                output["search_results"] = result.get("search_results", [])
            
            return output
            
    except urllib.error.HTTPError as e:
        error_body = e.read().decode()
        return {"error": f"HTTP {e.code}: {error_body}"}
    except Exception as e:
        return {"error": str(e)}

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="Query Perplexity API")
    parser.add_argument("query", help="The query to send")
    parser.add_argument("--model", default="sonar", 
                       choices=["sonar", "sonar-pro", "sonar-reasoning", "sonar-deep-research"],
                       help="Model to use")
    parser.add_argument("--max-tokens", type=int, default=1000,
                       help="Maximum tokens in response")
    parser.add_argument("--citations", action="store_true",
                       help="Include citations in output")
    parser.add_argument("--search-context", default="low",
                       choices=["low", "medium", "high"],
                       help="Search context size")
    
    args = parser.parse_args()
    
    result = query_perplexity(
        args.query,
        model=args.model,
        max_tokens=args.max_tokens,
        include_citations=args.citations,
        search_context=args.search_context
    )
    
    if "error" in result:
        print(f"Error: {result['error']}", file=sys.stderr)
        sys.exit(1)
    
    print(result["text"])
    
    if args.citations and result.get("citations"):
        print("\n--- Sources ---")
        for i, citation in enumerate(result["citations"][:5], 1):
            print(f"[{i}] {citation}")

if __name__ == "__main__":
    main()