#!/usr/bin/env python3
"""
Benchmark script for ColBERT token truncation optimization.

Tests the impact of truncating ColBERT query tokens on:
1. Query latency (speedup)
2. Result quality (overlap with full-token results)

KEY INSIGHT: This benchmark tests with DSL-prefixed queries because:
- Class components have symbols embedded: "ᵬ Name: Button..."
- DSL symbols at the START of queries should be preserved by truncation
- Truncating to first N% of tokens keeps these high-signal symbols

Usage:
    python scripts/test_colbert_truncation.py

The script:
1. Pulls real patterns from mcp_gchat_cards_v7 collection (with DSL in relationship_text)
2. Generates DSL-prefixed test queries using ModuleWrapper symbols
3. For each query, runs with 100%, 50%, 25%, 10% tokens
4. Reports latency and result overlap metrics
"""

import os
import sys
import time
from dataclasses import dataclass
from typing import Any

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config.settings import settings


@dataclass
class BenchmarkResult:
    """Result from a single benchmark run."""

    description: str
    token_ratio: float
    full_tokens: int
    truncated_tokens: int
    query_latency_ms: float
    result_ids: list[str]
    result_scores: list[float]


@dataclass
class ComparisonMetrics:
    """Metrics comparing truncated results to full-token baseline."""

    token_ratio: float
    avg_speedup_percent: float
    avg_result_overlap_top5: float
    avg_result_overlap_top3: float
    avg_score_correlation: float
    sample_count: int


def get_module_wrapper():
    """Get the card_framework ModuleWrapper for DSL symbol access."""
    try:
        from gchat.card_framework_wrapper import get_card_framework_wrapper

        return get_card_framework_wrapper()
    except Exception as e:
        print(f"Failed to get ModuleWrapper: {e}")
        return None


def get_test_descriptions(limit: int = 30) -> list[dict[str, Any]]:
    """
    Pull real patterns from the Qdrant collection WITH DSL notation.

    Returns list of dicts with:
        - card_description: The description text
        - relationship_text: DSL notation (e.g., "§[δ, Ƀ[ᵬ×2]] | Section...")
        - dsl_query: DSL-prefixed query for testing truncation
        - id: Point ID
    """
    try:
        from qdrant_client import models

        from config.qdrant_client import get_qdrant_client

        client = get_qdrant_client()
        collection_name = settings.card_collection

        # Query for instance_pattern points that have relationship_text (DSL)
        results, _ = client.scroll(
            collection_name=collection_name,
            scroll_filter=models.Filter(
                must=[
                    models.FieldCondition(
                        key="type", match=models.MatchValue(value="instance_pattern")
                    )
                ]
            ),
            limit=limit * 2,  # Fetch extra to filter
            with_payload=[
                "card_description",
                "relationship_text",
                "structure_description",
            ],
        )

        descriptions = []
        for point in results:
            desc = point.payload.get("card_description", "")
            dsl = point.payload.get("relationship_text", "")

            if desc and len(desc) > 10 and dsl:
                # Create DSL-prefixed query: "§[δ, Ƀ] Build a card with..."
                # This puts the high-signal DSL symbols at the START
                dsl_query = f"{dsl} {desc}"

                descriptions.append(
                    {
                        "card_description": desc,
                        "relationship_text": dsl,
                        "dsl_query": dsl_query,
                        "id": str(point.id),
                        "structure_description": point.payload.get(
                            "structure_description", ""
                        ),
                    }
                )

        # Sort by DSL query length to get variety
        descriptions.sort(key=lambda x: len(x["dsl_query"]))

        # Sample evenly across lengths
        if len(descriptions) > limit:
            step = len(descriptions) // limit
            descriptions = descriptions[::step][:limit]

        print(f"Loaded {len(descriptions)} test patterns with DSL from Qdrant")
        if descriptions:
            print(f"  Sample DSL: {descriptions[0]['relationship_text'][:60]}...")
        return descriptions

    except Exception as e:
        print(f"Failed to load descriptions from Qdrant: {e}")
        return []


def generate_synthetic_dsl_queries(wrapper, count: int = 10) -> list[dict[str, Any]]:
    """
    Generate synthetic DSL-prefixed test queries using ModuleWrapper symbols.

    Args:
        wrapper: ModuleWrapper instance with symbol_mapping
        count: Number of queries to generate

    Returns:
        List of test query dicts
    """
    if not wrapper:
        return []

    symbols = getattr(wrapper, "symbol_mapping", {})
    if not symbols:
        print("No symbol mapping available")
        return []

    print(f"Using {len(symbols)} DSL symbols from ModuleWrapper")

    # Common component combinations for testing
    test_patterns = [
        (["Section", "DecoratedText"], "Show a simple text message"),
        (
            ["Section", "DecoratedText", "ButtonList", "Button"],
            "Card with action buttons",
        ),
        (
            ["Section", "DecoratedText", "DecoratedText", "DecoratedText"],
            "Card with multiple text items",
        ),
        (["Section", "Grid", "GridItem"], "Grid layout with items"),
        (["Section", "DecoratedText", "Image"], "Card with image and text"),
        (["Section", "ButtonList", "Button", "Button"], "Card with two buttons"),
        (["Section", "ChipList", "Chip"], "Card with selectable chips"),
        (["Section", "Columns", "Column"], "Multi-column layout"),
        (["Section", "DecoratedText", "Divider", "ButtonList"], "Card with divider"),
        (
            ["Section", "TextParagraph", "ButtonList", "Button"],
            "Simple text with button",
        ),
    ]

    queries = []
    for i, (components, description) in enumerate(test_patterns[:count]):
        # Build DSL notation using wrapper
        dsl = wrapper.build_dsl_from_paths(components, "")

        # Create DSL-prefixed query
        dsl_query = f"{dsl} {description}"

        queries.append(
            {
                "card_description": description,
                "relationship_text": dsl,
                "dsl_query": dsl_query,
                "id": f"synthetic_{i}",
                "components": components,
            }
        )

    print(f"Generated {len(queries)} synthetic DSL queries")
    if queries:
        print(f"  Sample: {queries[0]['dsl_query'][:70]}...")

    return queries


def benchmark_query(
    feedback_loop, query: str, token_ratio: float, show_tokens: bool = False
) -> BenchmarkResult:
    """
    Run a single benchmark query with the given token ratio.

    Args:
        feedback_loop: FeedbackLoop instance
        query: The query string (should be DSL-prefixed for best results)
        token_ratio: Fraction of tokens to use (0.0-1.0)
        show_tokens: If True, print token details for debugging
    """
    # Get embedder to count tokens
    embedder = feedback_loop._get_embedder()
    if not embedder:
        raise RuntimeError("Could not get ColBERT embedder")

    # Count full tokens and show what gets truncated
    vectors_raw = list(embedder.query_embed(query))[0]
    full_token_count = len(vectors_raw)

    # Calculate truncated count
    if token_ratio < 1.0:
        truncated_count = max(1, int(full_token_count * token_ratio))
    else:
        truncated_count = full_token_count

    if show_tokens and token_ratio == 1.0:
        # Show what tokens are in the query (useful for understanding DSL placement)
        print(f"    Full query ({full_token_count} tokens): {query[:80]}...")

    # Time the query
    start_time = time.perf_counter()
    class_results, content_patterns, _ = feedback_loop.query_with_feedback(
        component_query=query,
        description=query,
        limit=10,
        use_negative_feedback=False,  # Disable for cleaner benchmark
        token_ratio=token_ratio,
    )
    elapsed_ms = (time.perf_counter() - start_time) * 1000

    # Combine results (prioritize content patterns, then class results)
    all_results = content_patterns + class_results
    result_ids = [str(r.get("id", "")) for r in all_results[:10]]
    result_scores = [r.get("score", 0.0) for r in all_results[:10]]

    return BenchmarkResult(
        description=query[:50] + "..." if len(query) > 50 else query,
        token_ratio=token_ratio,
        full_tokens=full_token_count,
        truncated_tokens=truncated_count,
        query_latency_ms=elapsed_ms,
        result_ids=result_ids,
        result_scores=result_scores,
    )


def calculate_overlap(
    baseline_ids: list[str], test_ids: list[str], top_n: int
) -> float:
    """Calculate overlap ratio for top N results."""
    if not baseline_ids or not test_ids:
        return 0.0

    baseline_set = set(baseline_ids[:top_n])
    test_set = set(test_ids[:top_n])

    if not baseline_set:
        return 1.0  # No baseline to compare

    overlap = len(baseline_set & test_set)
    return overlap / len(baseline_set)


def calculate_score_correlation(
    baseline_scores: list[float], test_scores: list[float]
) -> float:
    """Calculate simple correlation between score rankings."""
    if not baseline_scores or not test_scores:
        return 0.0

    # Normalize scores
    min_len = min(len(baseline_scores), len(test_scores))
    if min_len == 0:
        return 0.0

    # Simple rank correlation: check if relative ordering is preserved
    correct_order = 0
    total_pairs = 0

    for i in range(min_len):
        for j in range(i + 1, min_len):
            total_pairs += 1
            # Check if same relative ordering
            baseline_order = baseline_scores[i] >= baseline_scores[j]
            test_order = test_scores[i] >= test_scores[j]
            if baseline_order == test_order:
                correct_order += 1

    return correct_order / total_pairs if total_pairs > 0 else 1.0


def run_benchmark(test_data: list[dict[str, Any]]) -> dict[float, ComparisonMetrics]:
    """
    Run full benchmark across all DSL queries and token ratios.

    Args:
        test_data: List of dicts with 'dsl_query', 'relationship_text', 'card_description'
    """
    from gchat.feedback_loop import get_feedback_loop

    feedback_loop = get_feedback_loop()

    token_ratios = [1.0, 0.5, 0.25, 0.1]
    results_by_ratio: dict[float, list[BenchmarkResult]] = {r: [] for r in token_ratios}
    baseline_results: dict[str, BenchmarkResult] = {}  # query -> baseline result

    print("\nRunning benchmarks with DSL-prefixed queries...")
    print("-" * 70)

    for i, data in enumerate(test_data):
        # Use DSL-prefixed query for testing truncation effectiveness
        dsl_query = data.get("dsl_query", data.get("card_description", ""))
        dsl_part = data.get("relationship_text", "")[:30]
        desc_part = data.get("card_description", "")[:40]

        print(f'\n[{i + 1}/{len(test_data)}] DSL: {dsl_part}... + "{desc_part}..."')

        # Run for each token ratio
        for ratio in token_ratios:
            try:
                # Show token details only for first query at baseline
                show_tokens = i == 0 and ratio == 1.0
                result = benchmark_query(feedback_loop, dsl_query, ratio, show_tokens)
                results_by_ratio[ratio].append(result)

                if ratio == 1.0:
                    baseline_results[result.description] = result

                status = (
                    "baseline"
                    if ratio == 1.0
                    else f"{result.truncated_tokens}/{result.full_tokens} tokens"
                )
                print(
                    f"  {ratio:.0%}: {result.query_latency_ms:6.1f}ms, "
                    f"{len(result.result_ids)} results ({status})"
                )

            except Exception as e:
                print(f"  {ratio:.0%}: FAILED - {e}")

    # Calculate comparison metrics
    print("\n" + "=" * 70)
    print("RESULTS SUMMARY")
    print("=" * 70)

    metrics: dict[float, ComparisonMetrics] = {}

    for ratio in token_ratios:
        ratio_results = results_by_ratio[ratio]
        if not ratio_results:
            continue

        speedups = []
        overlaps_top5 = []
        overlaps_top3 = []
        correlations = []

        for result in ratio_results:
            baseline = baseline_results.get(
                result.description[:50] + "..."
                if len(result.description) > 50
                else result.description
            )
            if not baseline:
                # Find baseline by matching truncated description
                for desc, base in baseline_results.items():
                    if desc.startswith(result.description[:30]):
                        baseline = base
                        break

            if baseline and baseline.query_latency_ms > 0:
                # Speedup (negative means slower)
                speedup = (
                    (baseline.query_latency_ms - result.query_latency_ms)
                    / baseline.query_latency_ms
                    * 100
                )
                speedups.append(speedup)

                # Overlap
                overlaps_top5.append(
                    calculate_overlap(baseline.result_ids, result.result_ids, 5)
                )
                overlaps_top3.append(
                    calculate_overlap(baseline.result_ids, result.result_ids, 3)
                )

                # Score correlation
                correlations.append(
                    calculate_score_correlation(
                        baseline.result_scores, result.result_scores
                    )
                )

        metrics[ratio] = ComparisonMetrics(
            token_ratio=ratio,
            avg_speedup_percent=sum(speedups) / len(speedups) if speedups else 0,
            avg_result_overlap_top5=(
                sum(overlaps_top5) / len(overlaps_top5) if overlaps_top5 else 0
            ),
            avg_result_overlap_top3=(
                sum(overlaps_top3) / len(overlaps_top3) if overlaps_top3 else 0
            ),
            avg_score_correlation=(
                sum(correlations) / len(correlations) if correlations else 0
            ),
            sample_count=len(ratio_results),
        )

    return metrics


def print_metrics_table(metrics: dict[float, ComparisonMetrics]) -> None:
    """Print a formatted metrics table."""
    print("\n" + "-" * 70)
    print(
        f"{'Token Ratio':<12} {'Speedup':<10} {'Top-5 Overlap':<14} "
        f"{'Top-3 Overlap':<14} {'Score Corr.':<12} {'Samples':<8}"
    )
    print("-" * 70)

    for ratio in sorted(metrics.keys(), reverse=True):
        m = metrics[ratio]
        speedup_str = f"{m.avg_speedup_percent:+.1f}%" if ratio != 1.0 else "baseline"
        print(
            f"{m.token_ratio:<12.0%} {speedup_str:<10} {m.avg_result_overlap_top5:<14.1%} "
            f"{m.avg_result_overlap_top3:<14.1%} {m.avg_score_correlation:<12.1%} "
            f"{m.sample_count:<8}"
        )

    print("-" * 70)


def print_recommendations(metrics: dict[float, ComparisonMetrics]) -> None:
    """Print optimization recommendations based on metrics."""
    print("\nRECOMMENDATIONS:")
    print("-" * 70)

    for ratio in [0.5, 0.25, 0.1]:
        if ratio not in metrics:
            continue

        m = metrics[ratio]
        speedup = m.avg_speedup_percent
        overlap = m.avg_result_overlap_top5

        if overlap >= 0.9 and speedup >= 20:
            verdict = "RECOMMENDED"
            reason = f"{speedup:.0f}% faster with {overlap:.0%} result quality"
        elif overlap >= 0.8 and speedup >= 40:
            verdict = "CONSIDER"
            reason = f"Good trade-off: {speedup:.0f}% faster, {overlap:.0%} quality"
        elif overlap < 0.7:
            verdict = "NOT RECOMMENDED"
            reason = f"Quality too low ({overlap:.0%})"
        else:
            verdict = "MARGINAL"
            reason = (
                f"Speedup ({speedup:.0f}%) may not justify quality loss ({overlap:.0%})"
            )

        print(f"  {ratio:.0%} tokens: {verdict}")
        print(f"    {reason}")


def main():
    """Main entry point."""
    print("=" * 70)
    print("ColBERT Token Truncation Benchmark (DSL-Prefixed Queries)")
    print("=" * 70)
    print(f"\nCollection: {settings.card_collection}")
    print("\nThis benchmark tests with DSL symbols at the START of queries.")
    print("Hypothesis: Truncation preserves high-signal DSL symbols in first tokens.")

    # Get ModuleWrapper for DSL symbol access
    wrapper = get_module_wrapper()

    # Try to load real patterns with DSL from Qdrant
    test_data = get_test_descriptions(limit=20)

    if not test_data:
        print("\nNo real patterns found. Generating synthetic DSL queries...")
        test_data = generate_synthetic_dsl_queries(wrapper, count=10)

    if not test_data:
        print(
            "\nERROR: Could not generate test data. Check ModuleWrapper availability."
        )
        return

    if len(test_data) < 5:
        print(f"\nWARNING: Only {len(test_data)} test queries available.")
        print("Results may not be statistically significant.")

    # Show symbol mapping sample
    if wrapper:
        symbols = getattr(wrapper, "symbol_mapping", {})
        if symbols:
            sample_symbols = list(symbols.items())[:5]
            print(f"\nDSL Symbol samples: {dict(sample_symbols)}")

    # Run benchmark
    metrics = run_benchmark(test_data)

    # Print results
    print_metrics_table(metrics)
    print_recommendations(metrics)

    print("\n" + "=" * 70)
    print("INTERPRETATION:")
    print("-" * 70)
    print("If DSL symbols are in the first tokens, truncation should:")
    print("  - Preserve the §, δ, ᵬ symbols that strongly match class embeddings")
    print("  - Maintain high overlap even at 25-50% token ratios")
    print("  - Show better quality than plain text truncation")
    print("=" * 70)

    print("\nDone!")


if __name__ == "__main__":
    main()
