#!/usr/bin/env python3
"""Prompt compression benchmark: compare text before/after compression.

Simulates LLMLingua-style token scoring. In production, swap for:
- LLMLingua-2: fast constant-latency token selector
- LongLLMLingua: question-aware compression for RAG
Requires: pip install llmlingua (for real usage)
"""

from __future__ import annotations

import re
from collections import Counter
from typing import Optional


# ── Simulated compression (swap for LLMLingua in prod) ───────────────────


def tokenize(text: str) -> list[str]:
    """Simple whitespace tokenizer. In prod: use the model's actual tokenizer."""
    return text.split()


def detokenize(tokens: list[str]) -> str:
    return " ".join(tokens)


def compress_simple(
    text: str,
    question: str = "",
    ratio: float = 0.5,
    protect_phrases: Optional[list[str]] = None,
) -> tuple[str, dict]:
    """Simulate question-aware compression: keep tokens most relevant to the question.

    Production: LongLLMLingua scores each token by perplexity/info-content
    conditioned on the question, keeping the most informative.
    """
    tokens = tokenize(text)
    n_keep = max(1, int(len(tokens) * ratio))

    protected: set[int] = set()
    if protect_phrases:
        for phrase in protect_phrases:
            for m in re.finditer(re.escape(phrase), text, re.IGNORECASE):
                start_tok = len(tokenize(text[:m.start()]))
                end_tok = len(tokenize(text[:m.end()]))
                for i in range(start_tok, end_tok):
                    protected.add(i)

    # Simple relevance scoring: tokens matching the question score higher
    q_words = set(w.lower() for w in tokenize(question)) if question else set()

    scored = []
    for i, tok in enumerate(tokens):
        score = 2.0 if i in protected else (1.5 if tok.lower() in q_words else 1.0)
        scored.append((i, tok, score))

    # Keep top-scoring tokens (or protected + best of rest)
    protected_tokens = [s for s in scored if s[0] in protected]
    rest = [s for s in scored if s[0] not in protected]
    rest_sorted = sorted(rest, key=lambda x: x[2], reverse=True)

    keep = protected_tokens + rest_sorted[:max(0, n_keep - len(protected_tokens))]
    keep.sort(key=lambda x: x[0])  # restore original order

    compressed = detokenize([k[1] for k in keep])

    stats = {
        "original_tokens": len(tokens),
        "compressed_tokens": len(keep),
        "compression_ratio": len(keep) / max(1, len(tokens)),
        "protected_tokens": len(protected_tokens),
        "question_aware": bool(question),
    }
    return compressed, stats


# ── Benchmark ─────────────────────────────────────────────────────────────


def main() -> None:
    # Simulated RAG context (retrieved chunks)
    context = (
        "The PostgreSQL database system version 14.2 was deployed on server svr-prod-01 "
        "at 2025-03-12T09:00:00Z by the DevOps team. The migration involved 42 tables and "
        "approximately 1.2 million rows of user data. Connection pooling is configured via "
        "pgbouncer with a pool size of 20 and a max_connections limit of 100. The query "
        "planner shows index scans on the users and orders tables, with a hash join for "
        "complex reporting queries. Monitoring dashboards at grafana.example.com track "
        "RED metrics for request rate, error rate, and duration p95."
    )

    queries = [
        ("How many tables were migrated?", 0.4),  # ratio
        ("What is the max_connections limit?", 0.3),
        ("Where are the monitoring dashboards?", 0.4),
        ("Explain the full system architecture", 0.8),  # broad query, keep more
    ]

    print("=== Prompt Compression Benchmark ===\n")

    for question, ratio in queries:
        compressed, stats = compress_simple(
            context, question=question, ratio=ratio,
            protect_phrases=["PostgreSQL", "pgbouncer"],
        )
        print(f"Q: {question}")
        print(f"  Ratio: {stats['compression_ratio']:.0%} ({stats['original_tokens']} -> {stats['compressed_tokens']} tokens)")
        print(f"  Compressed: {compressed[:120]}")
        print()

    # Cost calculation
    print("=== Cost model (example) ===")
    print(f"  Context tokens (raw): ~100 tokens")
    print(f"  At $3/M input tokens, 2000 req/mo:")
    monthly_raw = (100 * 2000 * 3) / 1_000_000
    monthly_4x = (25 * 2000 * 3) / 1_000_000
    print(f"    No compression: ${monthly_raw:.2f}/mo")
    print(f"    With 4x compression: ${monthly_4x:.2f}/mo (${monthly_raw - monthly_4x:.2f}/mo saved)")

    print("\nProduction rules:")
    print("  1. Never compress the user's instruction/task text")
    print("  2. Question-aware compression (LongLLMLingua) for RAG")
    print("  3. Start at 2-4x ratio, tune up only with eval")
    print("  4. Measure faithfulness before/after (RAGAS, etc.)")
    print("  5. Choose fast compressor (LLMLingua-2) for latency-sensitive")


if __name__ == "__main__":
    main()