#!/usr/bin/env python3
"""Semantic extraction benchmark for RLM.

Tests RLM's ability to find semantically relevant messages that
CANNOT be found with keyword search alone. This forces RLM to use
llm_query() sub-calls for semantic understanding.

The key insight from the RLM paper: the LLM should partition data
and use recursive sub-calls to semantically label/filter content.
"""

from __future__ import annotations

import json
import sys
from dataclasses import dataclass
from pathlib import Path

from common import DEFAULT_DB, run_rlm_query


@dataclass
class SemanticQuery:
    """A query requiring semantic understanding, not keyword matching."""
    id: str
    question: str
    concept: str  # The semantic concept we're looking for
    # Keywords that should NOT appear in results (proves semantic, not keyword)
    avoid_keywords: list[str]
    # Keywords that indicate success (semantic equivalents)
    success_indicators: list[str]
    # Evaluation: did RLM use llm_query() sub-calls?
    requires_subcalls: bool = True


def get_semantic_queries() -> list[SemanticQuery]:
    """Queries that require semantic understanding, not keyword matching."""
    return [
        SemanticQuery(
            id="sem-stressed",
            question="Find messages where I or someone was feeling stressed, overwhelmed, or anxious. Don't just search for the word 'stressed' - look for messages that EXPRESS stress.",
            concept="emotional stress/anxiety",
            avoid_keywords=["stressed", "stress", "anxious", "anxiety"],
            success_indicators=["overwhelmed", "can't handle", "too much", "freaking out", "losing it", "burned out", "exhausted"],
        ),
        SemanticQuery(
            id="sem-excited",
            question="Find messages expressing genuine excitement or enthusiasm. Don't search for 'excited' - find messages that SHOW excitement.",
            concept="excitement/enthusiasm",
            avoid_keywords=["excited", "exciting"],
            success_indicators=["can't wait", "so pumped", "amazing", "awesome", "!!!", "omg", "yay"],
        ),
        SemanticQuery(
            id="sem-plans",
            question="Find messages where someone is making or discussing future plans. Look for intent to do something, not just mentions of 'plan'.",
            concept="future planning/intent",
            avoid_keywords=["plan", "plans", "planning"],
            success_indicators=["going to", "gonna", "let's", "should we", "want to", "thinking about", "how about"],
        ),
        SemanticQuery(
            id="sem-gratitude",
            question="Find messages expressing gratitude or appreciation. Don't search for 'thank' - find messages that convey genuine appreciation.",
            concept="gratitude/appreciation",
            avoid_keywords=["thank", "thanks", "grateful"],
            success_indicators=["appreciate", "meant a lot", "so kind", "really helped", "you're the best", "love you for"],
        ),
        SemanticQuery(
            id="sem-frustration",
            question="Find messages where someone is expressing frustration or annoyance without using those words.",
            concept="frustration/annoyance",
            avoid_keywords=["frustrated", "annoyed", "annoying", "frustrating"],
            success_indicators=["ugh", "seriously", "come on", "wtf", "can't believe", "ridiculous", "so stupid"],
        ),
        SemanticQuery(
            id="sem-asking-help",
            question="Find messages where someone is asking for help or assistance without explicitly saying 'help me'.",
            concept="requesting assistance",
            avoid_keywords=["help", "assist"],
            success_indicators=["could you", "would you mind", "do you know how", "any idea", "can you", "I need"],
        ),
        SemanticQuery(
            id="sem-apologizing",
            question="Find messages where someone is apologizing or expressing regret without just saying 'sorry'.",
            concept="apology/regret",
            avoid_keywords=["sorry", "apologize"],
            success_indicators=["my bad", "my fault", "shouldn't have", "feel bad", "I messed up", "forgive me"],
        ),
        SemanticQuery(
            id="sem-sarcasm",
            question="Find messages that appear to be sarcastic. This requires understanding tone and context.",
            concept="sarcasm/irony",
            avoid_keywords=[],  # No keywords to avoid - pure semantic
            success_indicators=["yeah right", "sure", "totally", "obviously", "wow", "great"],
        ),
    ]


@dataclass
class SemanticResult:
    """Result from evaluating a semantic query."""
    passed: bool
    score: float
    subcalls: int
    tokens: int
    used_subcalls: bool
    found_indicators: list[str]
    used_keywords: list[str]
    semantic_score: float


def evaluate_semantic_result(
    output: str,
    query: SemanticQuery,
    subcalls: int,
    tokens: int,
) -> SemanticResult:
    """Evaluate if RLM found semantically relevant results."""

    output_lower = output.lower()

    # Check if RLM used keyword shortcuts (bad)
    used_keywords = []
    for kw in query.avoid_keywords:
        if kw.lower() in output_lower:
            used_keywords.append(kw)

    # Check for semantic success indicators
    found_indicators = []
    for indicator in query.success_indicators:
        if indicator.lower() in output_lower:
            found_indicators.append(indicator)

    # Did RLM use sub-calls? (Critical for semantic understanding)
    used_subcalls = subcalls > 0

    # Scoring
    semantic_score = len(found_indicators) / max(len(query.success_indicators), 1)
    keyword_penalty = len(used_keywords) / max(len(query.avoid_keywords), 1) if query.avoid_keywords else 0
    subcall_bonus = 1.0 if used_subcalls else 0.0

    # Final score: reward semantic matches and subcalls, penalize keyword shortcuts
    final_score = (semantic_score * 0.5) + (subcall_bonus * 0.3) - (keyword_penalty * 0.2)
    final_score = max(0, min(1, final_score))

    passed = final_score >= 0.3 and (used_subcalls or semantic_score > 0.3)

    return SemanticResult(
        passed=passed,
        score=final_score,
        subcalls=subcalls,
        tokens=tokens,
        used_subcalls=used_subcalls,
        found_indicators=found_indicators,
        used_keywords=used_keywords,
        semantic_score=semantic_score,
    )


def run_benchmark(db_path: str | None = None, verbose: bool = False) -> dict:
    """Run the semantic extraction benchmark."""
    db_path = db_path or str(DEFAULT_DB)

    print("=== Semantic Extraction Benchmark ===")
    print(f"Database: {db_path}")
    print("\nThis tests RLM's ability to find semantically relevant content")
    print("that CANNOT be found with keyword search alone.\n")

    queries = get_semantic_queries()
    print(f"Queries: {len(queries)}")

    results = []
    total_subcalls = 0
    total_tokens = 0
    passed = 0

    for q in queries:
        print(f"\n  [{q.id}] {q.concept}")
        print(f"    Question: {q.question[:60]}...")

        output, subcalls, tokens = run_rlm_query(q.question, db_path, verbose, benchmark_mode=True)
        total_subcalls = total_subcalls + subcalls
        total_tokens = total_tokens + tokens

        eval_result = evaluate_semantic_result(output, q, subcalls, tokens)

        if eval_result.passed:
            passed = passed + 1
            status = "✓"
        else:
            status = "✗"

        print(f"    {status} Score: {eval_result.score:.2f} | Sub-calls: {subcalls} | Tokens: {tokens:,}")
        if eval_result.found_indicators:
            print(f"      Found semantic indicators: {eval_result.found_indicators[:3]}")
        if eval_result.used_keywords:
            print(f"      ⚠ Used keyword shortcuts: {eval_result.used_keywords}")
        if not eval_result.used_subcalls:
            print("      ⚠ No llm_query() sub-calls - not using semantic analysis")

        results.append({
            "id": q.id,
            "concept": q.concept,
            "question": q.question,
            "passed": eval_result.passed,
            "score": eval_result.score,
            "subcalls": eval_result.subcalls,
            "tokens": eval_result.tokens,
            "used_subcalls": eval_result.used_subcalls,
            "found_indicators": eval_result.found_indicators,
            "used_keywords": eval_result.used_keywords,
            "semantic_score": eval_result.semantic_score,
            "output_preview": output[:500],
        })

    # Summary
    accuracy = passed / len(queries) * 100
    avg_subcalls = total_subcalls / len(queries)
    avg_tokens = total_tokens / len(queries)

    print("\n=== Results ===")
    print(f"Passed: {passed}/{len(queries)} ({accuracy:.1f}%)")
    print(f"Total sub-calls: {total_subcalls}")
    print(f"Avg sub-calls per query: {avg_subcalls:.1f}")
    print(f"Total tokens: {total_tokens:,}")
    print(f"Avg tokens per query: {avg_tokens:,.0f}")

    # Key insight
    queries_with_subcalls = sum(1 for r in results if r["used_subcalls"])
    print(f"\nQueries using llm_query(): {queries_with_subcalls}/{len(queries)}")

    if queries_with_subcalls < len(queries) / 2:
        print("\n⚠ RLM is NOT using semantic sub-calls effectively!")
        print("  It's falling back to keyword/FTS search instead of")
        print("  using llm_query() to semantically analyze message content.")

    return {
        "accuracy": accuracy,
        "passed": passed,
        "total": len(queries),
        "total_subcalls": total_subcalls,
        "avg_subcalls": avg_subcalls,
        "total_tokens": total_tokens,
        "avg_tokens": avg_tokens,
        "queries_using_subcalls": queries_with_subcalls,
        "results": results,
    }


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Run semantic extraction benchmark")
    parser.add_argument("--db", help="Path to shadow.db", default=str(DEFAULT_DB))
    parser.add_argument("--verbose", "-v", action="store_true")
    parser.add_argument("--output", "-o", help="Save results to JSON")

    args = parser.parse_args()

    benchmark_results = run_benchmark(args.db, args.verbose)

    if args.output:
        with open(args.output, "w") as f:
            json.dump(benchmark_results, f, indent=2, default=str)
        print(f"\nResults saved to: {args.output}")
