#!python
"""
psst-openai: Send psst-compressed prompts to OpenAI API
Usage: psst-openai "Your prompt here" [--model gpt-4] [--system "System prompt"] [--glossary path/to/glossary.json]
"""

import os
import sys
import argparse
import json
from typing import Dict, List, Optional, Union
try:
    import requests
except ImportError:
    print("Error: requests library not found. Install with: pip install requests")
    sys.exit(1)

from psst_compiler import PsstCompiler

# Default configuration
DEFAULT_MODEL = "gpt-3.5-turbo"
DEFAULT_API_KEY_ENV = "OPENAI_API_KEY"
DEFAULT_GLOSSARY = "core_glossary.json"

def get_api_key() -> str:
    """Get API key from environment variable."""
    api_key = os.environ.get(DEFAULT_API_KEY_ENV)
    if not api_key:
        print(f"❌ Error: OpenAI API key not found.")
        print(f"   Set the {DEFAULT_API_KEY_ENV} environment variable:")
        print(f"   export {DEFAULT_API_KEY_ENV}=your_api_key_here")
        sys.exit(1)
    return api_key

def call_openai_api(
    prompt: str, 
    model: str = DEFAULT_MODEL,
    system_prompt: Optional[str] = None,
    api_key: Optional[str] = None
) -> Dict:
    """Send prompt to OpenAI API and return response."""
    if api_key is None:
        api_key = get_api_key()
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})
    
    payload = {
        "model": model,
        "messages": messages
    }
    
    try:
        response = requests.post(
            "https://api.openai.com/v1/chat/completions",
            headers=headers,
            json=payload,
            timeout=30
        )
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"❌ API call failed: {e}")
        if hasattr(e, "response") and hasattr(e.response, "text"):
            print(f"   Response: {e.response.text}")
        sys.exit(1)

def compress_prompt(prompt: str, glossary_path: str = DEFAULT_GLOSSARY) -> tuple[str, int]:
    """Compress prompt using psst compiler. Returns (compressed_prompt, savings)."""
    try:
        compiler = PsstCompiler(glossary_path)
        original_length = len(prompt)
        compressed = compiler.compress(prompt)
        compressed_length = len(compressed)
        savings = original_length - compressed_length
        return compressed, savings
    except Exception as e:
        print(f"❌ Compression failed: {e}")
        sys.exit(1)

def read_prompt_from_stdin() -> str:
    """Read prompt from stdin if available."""
    if not sys.stdin.isatty():
        return sys.stdin.read().strip()
    return ""

def format_response(response_data: Dict, show_usage: bool = False) -> str:
    """Format the OpenAI response for display."""
    try:
        content = response_data["choices"][0]["message"]["content"]
        result = content.strip()
        
        if show_usage and "usage" in response_data:
            usage = response_data["usage"]
            result += f"\n\n📊 Token Usage:"
            result += f"\n   Prompt: {usage.get('prompt_tokens', 0)} tokens"
            result += f"\n   Response: {usage.get('completion_tokens', 0)} tokens"
            result += f"\n   Total: {usage.get('total_tokens', 0)} tokens"
        
        return result
    except (KeyError, IndexError) as e:
        print(f"❌ Failed to extract response content: {e}")
        print("Raw response:", json.dumps(response_data, indent=2))
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(
        description="Send psst-compressed prompts to OpenAI API",
        epilog="""
Examples:
  psst-openai "Explain quantum computing"
  cat prompt.txt | psst-openai
  psst-openai --system "You are a helpful assistant" "What is AI?"
  psst-openai --show-tokens --model gpt-4 "Complex analysis task"
        """,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument("prompt", nargs="?", help="Prompt text (omit to read from stdin)")
    parser.add_argument("--model", "-m", default=DEFAULT_MODEL, 
                        help=f"OpenAI model to use (default: {DEFAULT_MODEL})")
    parser.add_argument("--system", "-s", help="System prompt")
    parser.add_argument("--glossary", "-g", default=DEFAULT_GLOSSARY,
                        help=f"Glossary file to use (default: {DEFAULT_GLOSSARY})")
    parser.add_argument("--show-tokens", "-t", action="store_true", 
                        help="Show compression and token usage information")
    parser.add_argument("--raw", "-r", action="store_true",
                        help="Don't compress the prompt (send as is)")
    parser.add_argument("--json", "-j", action="store_true",
                        help="Return raw JSON response")
    parser.add_argument("--usage", "-u", action="store_true",
                        help="Show token usage statistics")
    
    args = parser.parse_args()
    
    # Read from stdin if no prompt argument and stdin has data
    stdin_prompt = read_prompt_from_stdin()
    if args.prompt is None and not stdin_prompt:
        parser.print_help()
        sys.exit(1)
    
    # Combine arguments and stdin
    prompt = args.prompt or stdin_prompt
    
    if not prompt.strip():
        print("❌ Error: Empty prompt provided")
        sys.exit(1)
    
    # Compress prompt unless raw mode
    if not args.raw:
        compressed, savings = compress_prompt(prompt, args.glossary)
        
        if args.show_tokens:
            original_length = len(prompt)
            compressed_length = len(compressed)
            print(f"🗜️  Compression Results:")
            print(f"   Original: {original_length} characters")
            print(f"   Compressed: {compressed_length} characters")
            print(f"   Savings: {savings} characters ({savings/original_length*100:.1f}% reduction)")
            print(f"   Using model: {args.model}")
            print("=" * 50)
        
        prompt_to_send = compressed
    else:
        prompt_to_send = prompt
        if args.show_tokens:
            print(f"📝 Sending uncompressed prompt ({len(prompt)} characters)")
            print(f"   Using model: {args.model}")
            print("=" * 50)
    
    # Send to API
    print("🤖 Calling OpenAI API...", end="", flush=True)
    response = call_openai_api(
        prompt=prompt_to_send,
        model=args.model,
        system_prompt=args.system
    )
    print(" ✅")
    print()
    
    # Process response
    if args.json:
        print(json.dumps(response, indent=2))
    else:
        formatted_response = format_response(response, args.usage)
        print(formatted_response)

if __name__ == "__main__":
    main() 