import sys
import json
import time
import argparse
from typing import List, Optional
from perplexity import Perplexity

DEFAULT_MODELS = [
    "openai/gpt-5.4",
    "anthropic/claude-sonnet-4-6",
    "google/gemini-3.1-flash-lite",
    "xai/grok-4.20-non-reasoning",
    "perplexity/sonar",
]


def compare_models(prompt: str, models: List[str]) -> List[dict]:
    """Send the same prompt to each model and collect metrics."""
    client = Perplexity()
    results = []

    for model in models:
        print(f"  Querying {model}...")
        try:
            start = time.time()
            response = client.responses.create(
                model=model,
                input=prompt,
                max_output_tokens=1024,
            )
            elapsed = time.time() - start

            output_text = response.output_text
            total_cost = response.usage.cost.total_cost
            input_tokens = response.usage.input_tokens
            output_tokens = response.usage.output_tokens

            results.append({
                "model": model,
                "status": "success",
                "latency_s": round(elapsed, 2),
                "response_length": len(output_text),
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "cost_usd": total_cost,
                "preview": output_text[:120].replace("\n", " "),
            })
        except Exception as e:
            results.append({
                "model": model,
                "status": "error",
                "error": str(e),
                "latency_s": None,
                "response_length": 0,
                "input_tokens": 0,
                "output_tokens": 0,
                "cost_usd": None,
                "preview": "",
            })

    return results


def run_fallback(prompt: str, models: List[str]) -> dict:
    """Send a single request with a model fallback chain."""
    client = Perplexity()

    print(f"  Sending request with fallback chain: {models}")
    start = time.time()
    response = client.responses.create(
        models=models,
        input=prompt,
        max_output_tokens=1024,
    )
    elapsed = time.time() - start

    return {
        "requested_models": models,
        "model_used": response.model,
        "latency_s": round(elapsed, 2),
        "response_length": len(response.output_text),
        "cost_usd": response.usage.cost.total_cost,
        "preview": response.output_text[:200].replace("\n", " "),
    }


def format_table(results: List[dict]) -> str:
    """Format comparison results as a text table."""
    # Sort by latency (successful responses first)
    successful = [r for r in results if r["status"] == "success"]
    failed = [r for r in results if r["status"] != "success"]
    successful.sort(key=lambda r: r["latency_s"])

    lines = []
    header = f"{'Model':<42} {'Latency':>8} {'Length':>8} {'Tokens':>8} {'Cost':>10}"
    lines.append(header)
    lines.append("-" * len(header))

    for r in successful:
        tokens = f"{r['input_tokens']}+{r['output_tokens']}"
        cost = f"${r['cost_usd']:.5f}"
        lines.append(
            f"{r['model']:<42} {r['latency_s']:>7.2f}s {r['response_length']:>8} {tokens:>8} {cost:>10}"
        )

    for r in failed:
        lines.append(f"{r['model']:<42} {'FAILED':>8} {'-':>8} {'-':>8} {'-':>10}")

    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(
        description="Multi-Provider Model Comparison"
    )
    parser.add_argument("prompt", nargs="?", help="The prompt to send")
    parser.add_argument("--file", help="Read prompt from a file")
    parser.add_argument(
        "--models",
        nargs="+",
        default=DEFAULT_MODELS,
        help="Models to compare",
    )
    parser.add_argument(
        "--fallback",
        action="store_true",
        help="Use model fallback chain instead of comparing",
    )
    parser.add_argument(
        "--json", action="store_true", help="Output results as JSON"
    )
    args = parser.parse_args()

    # Resolve prompt
    if args.file:
        with open(args.file, "r") as f:
            prompt = f.read().strip()
    elif args.prompt:
        prompt = args.prompt
    else:
        print("Error: Provide a prompt or use --file.", file=sys.stderr)
        sys.exit(1)

    print(f"Prompt: {prompt[:80]}{'...' if len(prompt) > 80 else ''}\n")

    if args.fallback:
        print("Running model fallback chain...\n")
        result = run_fallback(prompt, args.models)
        if args.json:
            print(json.dumps(result, indent=2))
        else:
            print(f"Fallback chain: {' -> '.join(result['requested_models'])}")
            print(f"Model used:     {result['model_used']}")
            print(f"Latency:        {result['latency_s']}s")
            print(f"Response length: {result['response_length']} chars")
            print(f"Cost:           ${result['cost_usd']:.5f}")
            print(f"\nPreview: {result['preview']}")
    else:
        print(f"Comparing {len(args.models)} models...\n")
        results = compare_models(prompt, args.models)
        if args.json:
            print(json.dumps(results, indent=2))
        else:
            print(format_table(results))
            print(f"\nComparison complete. {len(results)} models evaluated.")


if __name__ == "__main__":
    main()