#!/usr/bin/env python3
"""Generate a benchmark corpus with real GUIDs from shadow_corpus_test.db."""

import json
import os
import sqlite3
from pathlib import Path

# Default database path
DEFAULT_DB = Path.home() / "Library/Application Support/Recall/shadow_corpus_test.db"


def get_guids_matching(conn: sqlite3.Connection, pattern: str, limit: int = 5) -> list[str]:
    """Find GUIDs of messages matching via FTS5 (same as RLM search)."""
    # Use FTS5 MATCH to match how RLM searches, with JOIN to get guid and timestamp
    cursor = conn.execute(
        """SELECT m.guid FROM messages_fts
           JOIN messages m ON messages_fts.rowid = m.rowid
           WHERE messages_fts MATCH ?
           ORDER BY m.timestamp DESC LIMIT ?""",
        (pattern, limit)
    )
    return [row[0] for row in cursor.fetchall()]


def get_guids_from_sender(conn: sqlite3.Connection, sender_pattern: str, limit: int = 5) -> list[str]:
    """Find GUIDs of messages from a specific sender pattern."""
    cursor = conn.execute(
        "SELECT guid FROM messages WHERE sender LIKE ? ORDER BY timestamp DESC LIMIT ?",
        (f"%{sender_pattern}%", limit)
    )
    return [row[0] for row in cursor.fetchall()]


def main():
    db_path = os.environ.get("SHADOW_DB_PATH", str(DEFAULT_DB))

    if not Path(db_path).exists():
        print(f"Error: Database not found at {db_path}")
        return

    conn = sqlite3.connect(db_path)

    # Build queries with real GUIDs
    queries = []

    # Keyword queries - use full query text for FTS matching
    for query_text in [
        "pizza",
        "dinner plans",
        "concert tickets",
        "flight booking",
        "meeting schedule",
        "birthday",
        "coffee",
        "gym workout",
    ]:
        # Use first word as ID, full text for FTS
        term = query_text.split()[0]
        guids = get_guids_matching(conn, query_text, limit=5)
        if guids:
            queries.append({
                "id": f"kw-{term}",
                "text": query_text,
                "category": "keyword",
                "expected_guids": guids
            })
            print(f"kw-{term}: Found {len(guids)} messages")

    # Semantic queries (broader patterns) - use full query text for FTS matching
    semantic_queries = [
        "food delivery OR restaurant",
        "travel plans",
        "work project",
        "family gathering",
        "health OR fitness",
    ]
    for query_text in semantic_queries:
        # Use first word as ID
        term = query_text.split()[0]
        guids = get_guids_matching(conn, query_text, limit=5)
        if guids:
            queries.append({
                "id": f"sem-{term}",
                "text": query_text,
                "category": "semantic",
                "expected_guids": guids
            })
            print(f"sem-{term}: Found {len(guids)} messages")

    conn.close()

    # Create corpus structure (no messages - we use external DB)
    corpus = {
        "version": "2.0",
        "description": "Real corpus from shadow_corpus_test.db for RLM benchmarks",
        "external_db": True,
        "messages": [],  # Empty - using external DB
        "chats": [],
        "queries": queries
    }

    # Write corpus
    output_path = Path(__file__).parent.parent / "benchmarks" / "real_corpus.json"
    output_path.parent.mkdir(exist_ok=True)

    with open(output_path, "w") as f:
        json.dump(corpus, f, indent=2)

    print(f"\nGenerated {len(queries)} queries")
    print(f"Written to: {output_path}")


if __name__ == "__main__":
    main()
