#!/usr/bin/env python3
"""
RLM Soak Test - Random question loop to find prompt and sandbox errors.

Usage:
    python soak_test.py --count 50 --verbose
    python soak_test.py --category person --count 20
    python soak_test.py --loop  # Run continuously until stopped
"""

import argparse
import json
import os
import random
import re
import subprocess
import sys
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Literal

# Question categories
CATEGORIES = [
    "person",
    "time",
    "semantic",
    "aggregation",
    "app",
    "edge",
]

# Diverse question bank covering different patterns
QUESTION_BANK: dict[str, list[str]] = {
    "person": [
        "Who did I talk to about dinner plans?",
        "What did Sarah say last week?",
        "Find messages from my mom",
        "Who mentioned the gym recently?",
        "What conversations did I have with John?",
        "Who was I texting late at night?",
        "Find my chat with the person who mentioned Netflix",
        "Who did I make plans with last weekend?",
        "What did my coworkers say about the project?",
        "Find messages where someone asked for my help",
        "Who sent me a link recently?",
        "What did the person from Tinder say?",
        "Find conversations about meeting up",
        "Who mentioned coffee this week?",
        "What did my friends say about the party?",
    ],
    "time": [
        "What did I talk about yesterday?",
        "Summarize my conversations from last week",
        "What messages did I get this morning?",
        "Find conversations from January",
        "What was I discussing last month?",
        "Show me messages from the weekend",
        "What happened in my chats on Monday?",
        "Find my oldest conversation",
        "What's the most recent message I received?",
        "Summarize my chats from the past 3 days",
        "What did people say on my birthday?",
        "Find messages from last night",
    ],
    "semantic": [
        "Who seemed stressed or anxious recently?",
        "Find messages where someone was excited about something",
        "Who was making future plans?",
        "Find conversations showing gratitude or appreciation",
        "Who seemed frustrated or annoyed?",
        "Find messages where someone was asking for advice",
        "Who apologized to me recently?",
        "Find sarcastic or joking messages",
        "Who was venting about work?",
        "Find messages with good news",
        "Who seemed sad or upset?",
        "Find flirty or romantic conversations",
        "Who was complaining about something?",
        "Find messages where someone changed their mind",
    ],
    "aggregation": [
        "How many messages did I send this month?",
        "Who do I text the most?",
        "What percentage of my chats are group chats?",
        "How many people have I messaged this year?",
        "What's my most active conversation?",
        "How many messages mention food?",
        "Who sends me the longest messages?",
        "What time of day do I text the most?",
        "How many unread conversations do I have?",
        "What's the average length of my messages?",
        "How many messages did I receive vs send?",
        "Which day of the week am I most active?",
    ],
    "app": [
        "What did someone from Hinge say?",
        "Find my Tinder conversations",
        "What dating app messages do I have?",
        "Who did I match with on Bumble?",
        "Find messages about Uber or Lyft",
        "What DoorDash or delivery conversations do I have?",
        "Find messages about Venmo or payments",
        "What Spotify or music links were shared?",
        "Find Instagram or social media mentions",
        "What flight or travel bookings were discussed?",
    ],
    "edge": [
        "Find messages with emoji",
        "What about O'Brien's message?",
        "Find messages containing 'can't' or 'won't'",
        "Search for messages with quotes like \"hello\"",
        "Find messages with phone numbers",
        "What about messages in Spanish?",
        "Find very short messages (1-2 words)",
        "What messages have URLs?",
        "Find messages with ALL CAPS",
        "Search for messages with '...' or ellipsis",
        "Find messages that are just punctuation like '???' or '!!!'",
        "What about messages with @ mentions?",
    ],
}

FailureType = Literal[
    "sandbox_error",
    "sql_error",
    "empty_response",
    "phone_number_leak",
    "unknown_contact",
    "timeout",
    "execution_error",
    "model_error",
]


@dataclass
class TestResult:
    question: str
    category: str
    passed: bool
    failure_type: FailureType | None = None
    error_message: str | None = None
    answer: str | None = None
    duration_seconds: float = 0.0
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())


@dataclass
class SoakReport:
    total_questions: int = 0
    total_passed: int = 0
    total_failed: int = 0
    by_category: dict[str, dict[str, int]] = field(default_factory=dict)
    by_failure_type: dict[str, int] = field(default_factory=dict)
    failures: list[dict] = field(default_factory=list)
    start_time: str = field(default_factory=lambda: datetime.now().isoformat())
    end_time: str | None = None
    duration_seconds: float = 0.0

    def to_dict(self) -> dict:
        return {
            "total_questions": self.total_questions,
            "total_passed": self.total_passed,
            "total_failed": self.total_failed,
            "pass_rate": f"{(self.total_passed / self.total_questions * 100):.1f}%" if self.total_questions > 0 else "N/A",
            "by_category": self.by_category,
            "by_failure_type": self.by_failure_type,
            "failures": self.failures,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_seconds": self.duration_seconds,
        }


# Patterns for failure detection
PHONE_NUMBER_PATTERN = re.compile(r"\+?1?[-.\s]?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}")
UNKNOWN_CONTACT_PATTERNS = [
    "unknown contact",
    "unknown sender",
    "the user",
    "the sender",
    "phone number",
    "unidentified",
]
SQL_ERROR_PATTERNS = [
    "operationalerror",
    "syntax error",
    "no such column",
    "no such table",
    "near \"",
]
SANDBOX_ERROR_PATTERNS = [
    "sandboxerror",
    "syntaxerror",
    "nameerror",
    "attributeerror",
    "typeerror",
    "keyerror",
    "indexerror",
    "valueerror",
]
MODEL_ERROR_PATTERNS = [
    "modelrelayerror",
    "api error",
    "rate limit",
    "token limit",
    "context length",
]


def classify_failure(
    exit_code: int,
    stdout: str,
    stderr: str,
    timed_out: bool,
    question: str = "",
) -> tuple[FailureType | None, str | None]:
    """Classify the type of failure based on output."""

    if timed_out:
        return "timeout", "Query timed out"

    combined = (stdout + stderr).lower()

    # Check for sandbox/Python errors
    for pattern in SANDBOX_ERROR_PATTERNS:
        if pattern in combined:
            # Extract the actual error message
            for line in (stdout + stderr).split("\n"):
                if pattern.lower() in line.lower():
                    return "sandbox_error", line.strip()
            return "sandbox_error", f"Sandbox error detected: {pattern}"

    # Check for SQL errors
    for pattern in SQL_ERROR_PATTERNS:
        if pattern in combined:
            for line in (stdout + stderr).split("\n"):
                if pattern in line.lower():
                    return "sql_error", line.strip()
            return "sql_error", f"SQL error detected: {pattern}"

    # Check for model/API errors
    for pattern in MODEL_ERROR_PATTERNS:
        if pattern in combined:
            return "model_error", f"Model error detected: {pattern}"

    # Check for execution errors (non-zero exit, not caught above)
    if exit_code != 0:
        error_msg = stderr.strip() if stderr.strip() else "Unknown execution error"
        return "execution_error", error_msg

    # Check answer quality (only if we got an answer)
    answer = stdout.strip()

    if not answer:
        return "empty_response", "No answer returned"

    # Check for phone number leaks (contact resolution failure)
    # Skip this check if the question explicitly asks about phone numbers
    question_lower = question.lower()
    asks_for_phone_numbers = "phone number" in question_lower or "phone numbers" in question_lower
    if not asks_for_phone_numbers and PHONE_NUMBER_PATTERN.search(answer):
        return "phone_number_leak", "Phone number found in response instead of contact name"

    # Check for unknown contact indicators
    # Skip patterns that match the question intent (e.g., "phone number" when asking about phone numbers)
    answer_lower = answer.lower()
    for pattern in UNKNOWN_CONTACT_PATTERNS:
        # Skip "phone number" pattern if question asks about phone numbers
        if pattern == "phone number" and asks_for_phone_numbers:
            continue
        if pattern in answer_lower:
            return "unknown_contact", f"Contact resolution issue: '{pattern}' found in response"

    return None, None


def run_query(
    question: str,
    db_path: str,
    contacts_db_path: str | None,
    timeout_seconds: int = 60,
    verbose: bool = False,
) -> tuple[int, str, str, bool]:
    """Run an RLM query and return (exit_code, stdout, stderr, timed_out)."""

    cmd = [
        "uv", "run", "--directory", str(Path(__file__).parent),
        "rcl-rlm", question, "--db", db_path,
    ]

    if contacts_db_path:
        cmd.extend(["--contacts-db", contacts_db_path])

    if verbose:
        cmd.append("--verbose")

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout_seconds,
            env={**os.environ, "PYTHONUNBUFFERED": "1"},
        )
        return result.returncode, result.stdout, result.stderr, False
    except subprocess.TimeoutExpired:
        return 1, "", f"Timeout after {timeout_seconds} seconds", True


def run_single_test(
    question: str,
    category: str,
    db_path: str,
    contacts_db_path: str | None,
    timeout_seconds: int,
    verbose: bool,
) -> TestResult:
    """Run a single test question and return the result."""

    start = time.time()
    exit_code, stdout, stderr, timed_out = run_query(
        question, db_path, contacts_db_path, timeout_seconds, verbose
    )
    duration = time.time() - start

    failure_type, error_message = classify_failure(exit_code, stdout, stderr, timed_out, question)

    return TestResult(
        question=question,
        category=category,
        passed=failure_type is None,
        failure_type=failure_type,
        error_message=error_message,
        answer=stdout.strip() if failure_type is None else None,
        duration_seconds=round(duration, 2),
    )


def select_questions(
    count: int,
    category: str | None = None,
    shuffle: bool = True,
) -> list[tuple[str, str]]:
    """Select questions to run. Returns list of (question, category) tuples."""

    if category:
        if category not in QUESTION_BANK:
            raise ValueError(f"Unknown category: {category}. Valid: {CATEGORIES}")
        questions = [(q, category) for q in QUESTION_BANK[category]]
    else:
        questions = [
            (q, cat)
            for cat, qs in QUESTION_BANK.items()
            for q in qs
        ]

    if shuffle:
        random.shuffle(questions)

    # If count exceeds available questions, cycle through
    if count > len(questions):
        multiplier = (count // len(questions)) + 1
        questions = questions * multiplier

    return questions[:count]


def run_soak_test(
    count: int,
    category: str | None,
    db_path: str,
    contacts_db_path: str | None,
    timeout_seconds: int,
    verbose: bool,
    loop: bool,
    output_file: str | None,
) -> SoakReport:
    """Run the soak test and return a report."""

    report = SoakReport()

    # Initialize category stats
    for cat in CATEGORIES:
        report.by_category[cat] = {"passed": 0, "failed": 0}

    iteration = 0
    try:
        while True:
            iteration += 1
            questions = select_questions(count, category)

            if loop:
                print(f"\n{'='*60}")
                print(f"Iteration {iteration} - Running {len(questions)} questions")
                print(f"{'='*60}")

            for i, (question, cat) in enumerate(questions, 1):
                report.total_questions += 1

                print(f"\n[{i}/{len(questions)}] ({cat}) {question}")

                result = run_single_test(
                    question, cat, db_path, contacts_db_path, timeout_seconds, verbose
                )

                if result.passed:
                    report.total_passed += 1
                    report.by_category[cat]["passed"] += 1
                    print(f"  PASS ({result.duration_seconds}s)")
                    if verbose and result.answer:
                        # Show first 200 chars of answer
                        preview = result.answer[:200] + "..." if len(result.answer) > 200 else result.answer
                        print(f"  Answer: {preview}")
                else:
                    report.total_failed += 1
                    report.by_category[cat]["failed"] += 1
                    report.by_failure_type[result.failure_type] = (
                        report.by_failure_type.get(result.failure_type, 0) + 1
                    )
                    report.failures.append({
                        "question": result.question,
                        "category": result.category,
                        "failure_type": result.failure_type,
                        "error_message": result.error_message,
                        "duration_seconds": result.duration_seconds,
                        "timestamp": result.timestamp,
                    })
                    print(f"  FAIL [{result.failure_type}] ({result.duration_seconds}s)")
                    print(f"  Error: {result.error_message}")

            # Print interim summary
            print(f"\n{'='*60}")
            print(f"Progress: {report.total_passed}/{report.total_questions} passed "
                  f"({report.total_passed/report.total_questions*100:.1f}%)")
            if report.by_failure_type:
                print(f"Failures by type: {dict(report.by_failure_type)}")
            print(f"{'='*60}")

            if not loop:
                break

    except KeyboardInterrupt:
        print("\n\nSoak test interrupted by user.")

    report.end_time = datetime.now().isoformat()
    start_dt = datetime.fromisoformat(report.start_time)
    end_dt = datetime.fromisoformat(report.end_time)
    report.duration_seconds = round((end_dt - start_dt).total_seconds(), 2)

    # Save report if output file specified
    if output_file:
        with open(output_file, "w") as f:
            json.dump(report.to_dict(), f, indent=2)
        print(f"\nReport saved to: {output_file}")

    return report


def print_final_report(report: SoakReport) -> None:
    """Print a summary of the soak test results."""

    print("\n" + "=" * 60)
    print("SOAK TEST REPORT")
    print("=" * 60)
    print(f"Duration: {report.duration_seconds}s")
    print(f"Total questions: {report.total_questions}")
    print(f"Passed: {report.total_passed}")
    print(f"Failed: {report.total_failed}")
    if report.total_questions > 0:
        print(f"Pass rate: {report.total_passed / report.total_questions * 100:.1f}%")

    print("\nBy category:")
    for cat, stats in sorted(report.by_category.items()):
        total = stats["passed"] + stats["failed"]
        if total > 0:
            rate = stats["passed"] / total * 100
            print(f"  {cat}: {stats['passed']}/{total} ({rate:.0f}%)")

    if report.by_failure_type:
        print("\nFailures by type:")
        for ftype, count in sorted(report.by_failure_type.items(), key=lambda x: -x[1]):
            print(f"  {ftype}: {count}")

    if report.failures:
        print(f"\nRecent failures (showing up to 10):")
        for failure in report.failures[-10:]:
            print(f"  [{failure['failure_type']}] {failure['question'][:60]}...")
            print(f"    -> {failure['error_message'][:80]}")

    print("=" * 60)


def main():
    parser = argparse.ArgumentParser(
        description="RLM Soak Test - Find prompt and sandbox errors"
    )
    parser.add_argument(
        "--count", "-n",
        type=int,
        default=20,
        help="Number of questions to run (default: 20)",
    )
    parser.add_argument(
        "--category", "-c",
        choices=CATEGORIES,
        help="Run only questions from this category",
    )
    parser.add_argument(
        "--db",
        default=os.environ.get(
            "SHADOW_DB_PATH",
            os.path.expanduser("~/Library/Application Support/RecallRLM/shadow.db")
        ),
        help="Path to shadow.db",
    )
    parser.add_argument(
        "--contacts-db",
        default=None,
        help="Path to contacts.db (default: derived from --db path)",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=60,
        help="Timeout per query in seconds (default: 60)",
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Show verbose output including answers",
    )
    parser.add_argument(
        "--loop",
        action="store_true",
        help="Run continuously until stopped (Ctrl+C)",
    )
    parser.add_argument(
        "--output", "-o",
        help="Save JSON report to this file",
    )

    args = parser.parse_args()

    # Verify DB exists
    if not os.path.exists(args.db):
        print(f"Error: shadow.db not found at {args.db}", file=sys.stderr)
        print("Set SHADOW_DB_PATH or use --db to specify the path.", file=sys.stderr)
        sys.exit(1)

    # Check for API key
    if not os.environ.get("MODELRELAY_API_KEY"):
        print("Error: MODELRELAY_API_KEY environment variable not set.", file=sys.stderr)
        sys.exit(1)

    # Derive contacts path from db path if not explicitly provided
    if args.contacts_db is None:
        # Check env var first, then derive from --db path
        if env_contacts := os.environ.get("CONTACTS_DB_PATH"):
            contacts_db_path = env_contacts
        else:
            contacts_db_path = str(Path(args.db).with_name("contacts.db"))
    else:
        contacts_db_path = args.contacts_db

    contacts_db = contacts_db_path if os.path.exists(contacts_db_path) else None
    if not contacts_db:
        print(f"Warning: contacts.db not found at {contacts_db_path}, running without contacts")

    print(f"RLM Soak Test")
    print(f"  Database: {args.db}")
    print(f"  Contacts: {contacts_db or 'None'}")
    print(f"  Questions: {args.count}")
    print(f"  Category: {args.category or 'all'}")
    print(f"  Timeout: {args.timeout}s")
    print(f"  Loop: {args.loop}")

    report = run_soak_test(
        count=args.count,
        category=args.category,
        db_path=args.db,
        contacts_db_path=contacts_db,
        timeout_seconds=args.timeout,
        verbose=args.verbose,
        loop=args.loop,
        output_file=args.output,
    )

    print_final_report(report)

    # Exit with error code if any failures
    sys.exit(0 if report.total_failed == 0 else 1)


if __name__ == "__main__":
    main()
