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

Tests RLM's ability to answer questions with verifiable ground truth.
Ground truth is computed via SQL, then compared to RLM's answers.
"""

from __future__ import annotations

import json
import sqlite3
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, TypeVar

from common import (
    DEFAULT_DB,
    run_rlm_query,
    extract_number,
    extract_percentage,
    compare_count,
    compare_exact,
    compare_contains,
)

# Type variables for generic extractors and comparators
T = TypeVar('T')
GroundTruth = int | float | str


@dataclass
class BenchmarkQuery:
    """A factual extraction query with ground truth."""
    id: str
    question: str
    category: str  # count, ranking, temporal, aggregation, discovery
    ground_truth_sql: str
    extract_answer: Callable[[str], T | None]
    compare: Callable[[T | None, GroundTruth], tuple[bool, str]]


def get_queries() -> list[BenchmarkQuery]:
    """Define the benchmark queries."""
    return [
        # === COUNTING TASKS ===
        BenchmarkQuery(
            id="count-total",
            question="How many total messages are in the database?",
            category="count",
            ground_truth_sql="SELECT COUNT(*) FROM messages",
            extract_answer=extract_number,
            compare=lambda rlm, truth: compare_count(rlm, int(truth), tolerance=0.05),
        ),
        BenchmarkQuery(
            id="count-pizza",
            question="How many messages mention pizza?",
            category="count",
            ground_truth_sql="SELECT COUNT(*) FROM messages WHERE text LIKE '%pizza%'",
            extract_answer=extract_number,
            compare=lambda rlm, truth: compare_count(rlm, int(truth), tolerance=0.15),
        ),
        BenchmarkQuery(
            id="count-from-me",
            question="How many messages did I send (is_from_me = 1)?",
            category="count",
            ground_truth_sql="SELECT COUNT(*) FROM messages WHERE is_from_me = 1",
            extract_answer=extract_number,
            compare=lambda rlm, truth: compare_count(rlm, int(truth), tolerance=0.1),
        ),
        BenchmarkQuery(
            id="count-group-chats",
            question="How many messages are in group chats?",
            category="count",
            ground_truth_sql="SELECT COUNT(*) FROM messages WHERE is_group = 1",
            extract_answer=extract_number,
            compare=lambda rlm, truth: compare_count(rlm, int(truth), tolerance=0.1),
        ),
        BenchmarkQuery(
            id="count-unique-senders",
            question="How many different senders are there in messages where is_from_me = 0?",
            category="count",
            ground_truth_sql="SELECT COUNT(DISTINCT sender) FROM messages WHERE is_from_me = 0",
            extract_answer=extract_number,
            compare=lambda rlm, truth: compare_count(rlm, int(truth), tolerance=0.1),
        ),

        # === RANKING TASKS ===
        BenchmarkQuery(
            id="rank-top-sender",
            question="Who sends me the most messages (excluding myself, where is_from_me = 0)? Give the sender value.",
            category="ranking",
            ground_truth_sql="""SELECT sender FROM messages WHERE is_from_me = 0
                      GROUP BY sender ORDER BY COUNT(*) DESC LIMIT 1""",
            extract_answer=lambda text: text.strip(),
            compare=compare_contains,
        ),
        BenchmarkQuery(
            id="rank-top-chat",
            question="What is my most active chat by message count? Give the chat name.",
            category="ranking",
            ground_truth_sql="""SELECT chat_name FROM messages
                      GROUP BY chat_id ORDER BY COUNT(*) DESC LIMIT 1""",
            extract_answer=lambda text: text.strip(),
            compare=compare_contains,
        ),

        # === AGGREGATION TASKS ===
        BenchmarkQuery(
            id="agg-group-pct",
            question="What percentage of my messages are in group chats? Give a number.",
            category="aggregation",
            ground_truth_sql="""SELECT ROUND(100.0 * SUM(CASE WHEN is_group = 1 THEN 1 ELSE 0 END) / COUNT(*), 1)
                      FROM messages""",
            extract_answer=extract_percentage,
            compare=lambda rlm, truth: compare_count(rlm, float(truth), tolerance=0.15) if rlm else (False, "No percentage found"),
        ),
        BenchmarkQuery(
            id="agg-sent-pct",
            question="What percentage of messages did I send vs receive?",
            category="aggregation",
            ground_truth_sql="""SELECT ROUND(100.0 * SUM(CASE WHEN is_from_me = 1 THEN 1 ELSE 0 END) / COUNT(*), 1)
                      FROM messages""",
            extract_answer=extract_percentage,
            compare=lambda rlm, truth: compare_count(rlm, float(truth), tolerance=0.15) if rlm else (False, "No percentage found"),
        ),

        # === DISCOVERY TASKS ===
        BenchmarkQuery(
            id="disc-mentions-coffee",
            question="Do any messages mention coffee? Answer yes or no.",
            category="discovery",
            ground_truth_sql="SELECT CASE WHEN COUNT(*) > 0 THEN 'yes' ELSE 'no' END FROM messages WHERE text LIKE '%coffee%'",
            extract_answer=lambda text: "yes" if ("yes" in text.lower() or "found" in text.lower()) else "no",
            compare=compare_exact,
        ),
        BenchmarkQuery(
            id="disc-mentions-airport",
            question="Do any messages mention an airport? Answer yes or no.",
            category="discovery",
            ground_truth_sql="SELECT CASE WHEN COUNT(*) > 0 THEN 'yes' ELSE 'no' END FROM messages WHERE text LIKE '%airport%'",
            extract_answer=lambda text: "yes" if ("yes" in text.lower() or "found" in text.lower()) else "no",
            compare=compare_exact,
        ),
    ]


def compute_ground_truth(conn: sqlite3.Connection, queries: list[BenchmarkQuery]) -> dict[str, GroundTruth]:
    """Compute ground truth for all queries."""
    results: dict[str, GroundTruth] = {}
    for q in queries:
        cursor = conn.execute(q.ground_truth_sql)
        row = cursor.fetchone()
        results[q.id] = row[0] if row else 0
    return results


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

    print("=== Factual Extraction Benchmark ===")
    print(f"Database: {db_path}")

    # Get queries and compute ground truth
    queries = get_queries()
    conn = sqlite3.connect(db_path)
    ground_truth = compute_ground_truth(conn, queries)
    conn.close()

    print(f"Queries: {len(queries)}")
    print("\nGround Truth:")
    for q in queries:
        print(f"  {q.id}: {ground_truth[q.id]}")

    print("\nRunning RLM queries...")

    results = []
    correct = 0
    total_tokens = 0
    total_subcalls = 0

    for q in queries:
        print(f"\n  [{q.category}] {q.id}: {q.question[:50]}...")

        # Run RLM
        rlm_response, subcalls, tokens = run_rlm_query(q.question, db_path, verbose, benchmark_mode=True)
        total_tokens = total_tokens + tokens
        total_subcalls = total_subcalls + subcalls

        # Extract answer from RLM response
        rlm_answer = q.extract_answer(rlm_response)

        # Compare to ground truth
        truth = ground_truth[q.id]
        passed, reason = q.compare(rlm_answer, truth)

        if passed:
            correct = correct + 1
            print(f"    ✓ {reason} | Tokens: {tokens:,}")
        else:
            print(f"    ✗ {reason} | Tokens: {tokens:,}")

        results.append({
            "id": q.id,
            "question": q.question,
            "category": q.category,
            "ground_truth": truth,
            "rlm_answer": rlm_answer,
            "rlm_response": rlm_response[:500],
            "passed": passed,
            "reason": reason,
            "tokens": tokens,
            "subcalls": subcalls,
        })

    # Summary
    accuracy = correct / len(queries) * 100
    avg_tokens = total_tokens / len(queries)
    print("\n=== Results ===")
    print(f"Correct: {correct}/{len(queries)} ({accuracy:.1f}%)")
    print(f"Total tokens: {total_tokens:,}")
    print(f"Avg tokens per query: {avg_tokens:,.0f}")
    print(f"Total sub-calls: {total_subcalls}")

    # By category
    categories: dict[str, dict[str, int]] = {}
    for r in results:
        cat = r["category"]
        if cat not in categories:
            categories[cat] = {"correct": 0, "total": 0}
        categories[cat]["total"] = categories[cat]["total"] + 1
        if r["passed"]:
            categories[cat]["correct"] = categories[cat]["correct"] + 1

    print("\nBy Category:")
    for cat, stats in sorted(categories.items()):
        pct = stats["correct"] / stats["total"] * 100
        print(f"  {cat}: {stats['correct']}/{stats['total']} ({pct:.0f}%)")

    return {
        "accuracy": accuracy,
        "correct": correct,
        "total": len(queries),
        "total_tokens": total_tokens,
        "avg_tokens": avg_tokens,
        "total_subcalls": total_subcalls,
        "by_category": categories,
        "results": results,
    }


if __name__ == "__main__":
    import argparse

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

    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}")
