test_agent.py

"""
Comprehensive Test Suite for the SQL Query Writer Agent

This test suite covers:
- Unit tests for individual components (imports, extraction, validation)
- Integration tests with 50+ query test cases
- Performance benchmarking

Run: python test_agent.py [--full] [--benchmark] [--verbose]

Options:
    --full      Run full test suite including LLM integration
    --benchmark Run performance benchmarks across different models
    --verbose   Show detailed output for each test
"""

import os
import sys
import time
import json
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass, field
from datetime import datetime


# =============================================================================
# TEST CASE DEFINITIONS
# =============================================================================

@dataclass
class TestCase:
    """A single test case for the SQL Query Writer Agent."""
    question: str
    category: str
    should_contain: List[str] = field(default_factory=list)
    should_not_contain: List[str] = field(default_factory=list)
    difficulty: str = "medium"  # easy, medium, hard
    description: str = ""


# Comprehensive test cases covering all query types
TEST_CASES = [
    # ==========================================================================
    # BASIC COUNT QUERIES (10 cases)
    # ==========================================================================
    TestCase(
        question="How many customers are there?",
        category="count",
        should_contain=["count", "customers"],
        difficulty="easy",
        description="Basic count query"
    ),
    TestCase(
        question="Count all products",
        category="count",
        should_contain=["count", "products"],
        difficulty="easy"
    ),
    TestCase(
        question="How many orders were placed?",
        category="count",
        should_contain=["count", "orders"],
        difficulty="easy"
    ),
    TestCase(
        question="How many stores do we have?",
        category="count",
        should_contain=["count", "stores"],
        difficulty="easy"
    ),
    TestCase(
        question="How many staff members are there?",
        category="count",
        should_contain=["count", "staffs"],
        difficulty="easy"
    ),
    TestCase(
        question="How many unique brands are there?",
        category="count",
        should_contain=["count", "distinct", "brands"],
        difficulty="medium"
    ),
    TestCase(
        question="How many orders were placed in 2018?",
        category="count",
        should_contain=["count", "orders", "2018"],
        difficulty="medium"
    ),
    TestCase(
        question="How many products are in each category?",
        category="count",
        should_contain=["count", "products", "group by"],
        difficulty="medium"
    ),
    TestCase(
        question="How many customers are from each state?",
        category="count",
        should_contain=["count", "customers", "state", "group by"],
        difficulty="medium"
    ),
    TestCase(
        question="What is the number of orders per store?",
        category="count",
        should_contain=["count", "orders", "store", "group by"],
        difficulty="medium"
    ),
    
    # ==========================================================================
    # RANKING QUERIES (10 cases)
    # ==========================================================================
    TestCase(
        question="What are the top 5 most expensive products?",
        category="ranking",
        should_contain=["products", "order by", "limit 5"],
        difficulty="easy"
    ),
    TestCase(
        question="Top 3 best selling products",
        category="ranking",
        should_contain=["products", "sum", "order by", "limit 3"],
        difficulty="medium"
    ),
    TestCase(
        question="Which customer has placed the most orders?",
        category="ranking",
        should_contain=["customers", "orders", "count", "order by", "limit"],
        difficulty="medium"
    ),
    TestCase(
        question="Top 10 cheapest products",
        category="ranking",
        should_contain=["products", "order by", "limit 10"],
        difficulty="easy"
    ),
    TestCase(
        question="Which brand has the highest total revenue?",
        category="ranking",
        should_contain=["brands", "sum", "order by", "limit"],
        difficulty="hard"
    ),
    TestCase(
        question="Top 5 stores by number of orders",
        category="ranking",
        should_contain=["stores", "orders", "count", "order by", "limit 5"],
        difficulty="medium"
    ),
    TestCase(
        question="Which staff member has processed the most orders?",
        category="ranking",
        should_contain=["staffs", "orders", "count", "order by", "limit"],
        difficulty="medium"
    ),
    TestCase(
        question="Top 5 customers by total spending",
        category="ranking",
        should_contain=["customers", "sum", "order by", "limit 5"],
        difficulty="hard"
    ),
    TestCase(
        question="What are the 3 least popular product categories?",
        category="ranking",
        should_contain=["categories", "count", "order by", "limit 3"],
        difficulty="hard"
    ),
    TestCase(
        question="Which 5 products have the lowest inventory?",
        category="ranking",
        should_contain=["products", "stocks", "order by", "limit 5"],
        difficulty="medium"
    ),
    
    # ==========================================================================
    # AGGREGATION QUERIES (10 cases)
    # ==========================================================================
    TestCase(
        question="What is the total revenue?",
        category="aggregation",
        should_contain=["sum", "order_items"],
        difficulty="medium"
    ),
    TestCase(
        question="What is the average product price?",
        category="aggregation",
        should_contain=["avg", "products"],
        difficulty="easy"
    ),
    TestCase(
        question="Total revenue by store",
        category="aggregation",
        should_contain=["sum", "stores", "group by"],
        difficulty="medium"
    ),
    TestCase(
        question="What is the total revenue by brand?",
        category="aggregation",
        should_contain=["sum", "brands", "group by"],
        difficulty="hard"
    ),
    TestCase(
        question="What is the average order value?",
        category="aggregation",
        should_contain=["avg", "order"],
        difficulty="hard"
    ),
    TestCase(
        question="What are the minimum and maximum product prices?",
        category="aggregation",
        should_contain=["min", "max", "products"],
        difficulty="easy"
    ),
    TestCase(
        question="Total quantity sold per product",
        category="aggregation",
        should_contain=["sum", "quantity", "products", "group by"],
        difficulty="medium"
    ),
    TestCase(
        question="Average discount given per order",
        category="aggregation",
        should_contain=["avg", "discount", "order"],
        difficulty="medium"
    ),
    TestCase(
        question="Sum of inventory across all stores",
        category="aggregation",
        should_contain=["sum", "stocks"],
        difficulty="easy"
    ),
    TestCase(
        question="Total revenue per category",
        category="aggregation",
        should_contain=["sum", "categories", "group by"],
        difficulty="hard"
    ),
    
    # ==========================================================================
    # JOIN QUERIES (10 cases)
    # ==========================================================================
    TestCase(
        question="List all staff members and their stores",
        category="join",
        should_contain=["staffs", "stores", "join"],
        difficulty="easy"
    ),
    TestCase(
        question="Show products with their brand names",
        category="join",
        should_contain=["products", "brands", "join"],
        difficulty="easy"
    ),
    TestCase(
        question="Show all orders with customer names",
        category="join",
        should_contain=["orders", "customers", "join"],
        difficulty="easy"
    ),
    TestCase(
        question="Show products with their category and brand",
        category="join",
        should_contain=["products", "categories", "brands", "join"],
        difficulty="medium"
    ),
    TestCase(
        question="List order items with product names and order dates",
        category="join",
        should_contain=["order_items", "products", "orders", "join"],
        difficulty="medium"
    ),
    TestCase(
        question="Show staff members and their managers",
        category="join",
        should_contain=["staffs", "join"],
        difficulty="hard",
        description="Self-join query"
    ),
    TestCase(
        question="List all products with their store inventory levels",
        category="join",
        should_contain=["products", "stocks", "stores", "join"],
        difficulty="medium"
    ),
    TestCase(
        question="Show orders with customer name, staff name, and store name",
        category="join",
        should_contain=["orders", "customers", "staffs", "stores", "join"],
        difficulty="hard"
    ),
    TestCase(
        question="List products with brand, category, and total quantity sold",
        category="join",
        should_contain=["products", "brands", "categories", "order_items", "join"],
        difficulty="hard"
    ),
    TestCase(
        question="Show each store's staff count and total inventory",
        category="join",
        should_contain=["stores", "staffs", "stocks", "join"],
        difficulty="hard"
    ),
    
    # ==========================================================================
    # FILTER QUERIES (10 cases)
    # ==========================================================================
    TestCase(
        question="Show all orders from 2018",
        category="filter",
        should_contain=["orders", "2018"],
        difficulty="easy"
    ),
    TestCase(
        question="Find products under $500",
        category="filter",
        should_contain=["products", "500"],
        difficulty="easy"
    ),
    TestCase(
        question="Find customers from New York",
        category="filter",
        should_contain=["customers", "ny"],
        difficulty="easy"
    ),
    TestCase(
        question="Find products priced between $1000 and $2000",
        category="filter",
        should_contain=["products", "between", "1000", "2000"],
        difficulty="medium"
    ),
    TestCase(
        question="Show orders that have been shipped",
        category="filter",
        should_contain=["orders", "shipped"],
        difficulty="medium"
    ),
    TestCase(
        question="Find pending orders",
        category="filter",
        should_contain=["orders", "status"],
        difficulty="medium"
    ),
    TestCase(
        question="List products from the 'Trek' brand",
        category="filter",
        should_contain=["products", "brands", "trek"],
        difficulty="medium"
    ),
    TestCase(
        question="Show orders from the first quarter of 2018",
        category="filter",
        should_contain=["orders", "2018"],
        difficulty="medium"
    ),
    TestCase(
        question="Find active staff members",
        category="filter",
        should_contain=["staffs", "active"],
        difficulty="easy"
    ),
    TestCase(
        question="Products with model year 2018 or later",
        category="filter",
        should_contain=["products", "model_year", "2018"],
        difficulty="medium"
    ),
    
    # ==========================================================================
    # NULL HANDLING QUERIES (5 cases)
    # ==========================================================================
    TestCase(
        question="Find orders that have not been shipped yet",
        category="null_handling",
        should_contain=["orders", "null"],
        difficulty="medium"
    ),
    TestCase(
        question="Which products have never been ordered?",
        category="null_handling",
        should_contain=["products", "order_items"],
        should_not_contain=[],
        difficulty="hard"
    ),
    TestCase(
        question="Find staff members without managers",
        category="null_handling",
        should_contain=["staffs", "manager", "null"],
        difficulty="medium"
    ),
    TestCase(
        question="Customers who have never placed an order",
        category="null_handling",
        should_contain=["customers", "orders"],
        difficulty="hard"
    ),
    TestCase(
        question="Products not available in any store",
        category="null_handling",
        should_contain=["products", "stocks"],
        difficulty="hard"
    ),
    
    # ==========================================================================
    # DATE RANGE QUERIES (5 cases)
    # ==========================================================================
    TestCase(
        question="Monthly revenue trend for 2018",
        category="date_range",
        should_contain=["orders", "order_items", "2018", "group by"],
        difficulty="hard"
    ),
    TestCase(
        question="Revenue by year",
        category="date_range",
        should_contain=["orders", "order_items", "group by"],
        difficulty="medium"
    ),
    TestCase(
        question="Orders placed in December 2017",
        category="date_range",
        should_contain=["orders", "2017-12"],
        difficulty="medium"
    ),
    TestCase(
        question="Compare Q1 and Q2 revenue for 2018",
        category="date_range",
        should_contain=["orders", "order_items", "2018"],
        difficulty="hard"
    ),
    TestCase(
        question="Daily order count for January 2018",
        category="date_range",
        should_contain=["orders", "count", "2018-01"],
        difficulty="hard"
    ),
    
    # ==========================================================================
    # SUBQUERY / COMPLEX QUERIES (5 cases)
    # ==========================================================================
    TestCase(
        question="Find products more expensive than the average price",
        category="subquery",
        should_contain=["products", "avg"],
        difficulty="medium"
    ),
    TestCase(
        question="Customers who have ordered more than once",
        category="subquery",
        should_contain=["customers", "orders", "count"],
        difficulty="medium"
    ),
    TestCase(
        question="Stores that have sold all product categories",
        category="subquery",
        should_contain=["stores", "categories"],
        difficulty="hard"
    ),
    TestCase(
        question="Products with above-average sales",
        category="subquery",
        should_contain=["products", "order_items", "avg"],
        difficulty="hard"
    ),
    TestCase(
        question="Staff members who have never processed an order",
        category="subquery",
        should_contain=["staffs", "orders"],
        difficulty="hard"
    ),
]


# =============================================================================
# TEST FUNCTIONS
# =============================================================================

def test_imports() -> bool:
    """Test that all required modules can be imported."""
    print("Testing imports...")
    try:
        from db.bike_store import BikeStoreDb, get_schema_info
        from src.schema import SchemaContext
        from src.prompts import PromptBuilder
        from src.validator import SQLValidator
        from src.conversation import ConversationManager
        from src.model_fallback import ensure_model_configured
        from src.utils import extract_sql, clean_sql, normalize_for_comparison
        from agent import QueryWriter
        print("  All imports OK")
        return True
    except ImportError as e:
        print(f"  Import error: {e}")
        return False


def test_schema_context() -> bool:
    """Test schema context loading and formatting."""
    print("\nTesting schema context...")
    from db.bike_store import BikeStoreDb
    from src.schema import SchemaContext
    
    db = BikeStoreDb('bike_store.db')
    ctx = SchemaContext('bike_store.db')
    
    print(f"  Found {len(ctx.schema)} tables")
    print(f"  Tables: {list(ctx.schema.keys())}")
    
    schema_text = ctx.format_for_prompt()
    print(f"  Schema prompt length: {len(schema_text)} chars")
    
    # Check that all expected tables are present
    expected_tables = {'brands', 'categories', 'customers', 'order_items', 
                       'orders', 'products', 'staffs', 'stocks', 'stores'}
    actual_tables = set(ctx.schema.keys())
    
    if expected_tables == actual_tables:
        print("  All expected tables present")
        return True
    else:
        missing = expected_tables - actual_tables
        extra = actual_tables - expected_tables
        if missing:
            print(f"  Missing tables: {missing}")
        if extra:
            print(f"  Extra tables: {extra}")
        return False


def test_sql_extraction() -> bool:
    """Test SQL extraction from various LLM response formats."""
    print("\nTesting SQL extraction...")
    from src.utils import extract_sql
    
    cases = [
        # Basic SQL
        ("SELECT * FROM customers", "SELECT * FROM customers"),
        # Markdown code block with sql tag
        ("```sql\nSELECT * FROM orders\n```", "SELECT * FROM orders"),
        # Markdown code block without tag
        ("Here's the query:\n```\nSELECT COUNT(*) FROM products\n```\nThis counts products.", 
         "SELECT COUNT(*) FROM products"),
        # SQL with trailing explanation
        ("SELECT * FROM stores\n\nThis query gets all stores.", "SELECT * FROM stores"),
        # Multi-line SQL
        ("SELECT p.product_name, b.brand_name\nFROM products p\nJOIN brands b ON p.brand_id = b.brand_id",
         "SELECT p.product_name, b.brand_name FROM products p JOIN brands b ON p.brand_id = b.brand_id"),
    ]
    
    all_pass = True
    for input_text, expected in cases:
        result = extract_sql(input_text)
        if result != expected:
            print(f"  FAIL: got '{result}' expected '{expected}'")
            all_pass = False
    
    if all_pass:
        print(f"  All {len(cases)} extraction tests passed")
    return all_pass


def test_validator() -> bool:
    """Test SQL validator syntax checking."""
    print("\nTesting SQL validator...")
    from src.validator import SQLValidator

    validator = SQLValidator('bike_store.db')
    
    # Test valid SQL
    valid_sql = "SELECT * FROM customers LIMIT 5"
    is_valid, err = validator.validate(valid_sql)
    print(f"  Valid SQL test: {'PASS' if is_valid else 'FAIL - ' + err}")
    
    # Test invalid SQL (nonexistent table)
    invalid_sql = "SELECT * FROM nonexistent_table"
    is_valid, err = validator.validate(invalid_sql)
    print(f"  Invalid table test: {'PASS' if not is_valid else 'FAIL - should have failed'}")
    
    # Test invalid SQL (syntax error)
    syntax_error_sql = "SELEC * FROM customers"
    is_valid, err = validator.validate(syntax_error_sql)
    print(f"  Syntax error test: {'PASS' if not is_valid else 'FAIL - should have failed'}")
    
    return True


def test_semantic_validator() -> bool:
    """Test semantic validation of queries."""
    print("\nTesting semantic validator...")
    from src.validator import SQLValidator

    validator = SQLValidator('bike_store.db')
    
    test_cases = [
        # (question, sql, should_pass)
        ("How many customers?", "SELECT COUNT(*) FROM customers", True),
        ("How many customers?", "SELECT * FROM customers", False),  # Missing COUNT
        ("Top 5 products", "SELECT * FROM products ORDER BY list_price DESC LIMIT 5", True),
        ("Top 5 products", "SELECT * FROM products", False),  # Missing LIMIT
        ("Total revenue", "SELECT SUM(quantity * list_price) FROM order_items", True),
        ("Total revenue", "SELECT * FROM order_items", False),  # Missing SUM
        (
            "Which customer got the most discounts?",
            "SELECT c.first_name, c.last_name FROM customers c JOIN orders o ON c.customer_id = o.customer_id "
            "JOIN order_items oi ON o.order_id = oi.order_id GROUP BY c.first_name, c.last_name "
            "ORDER BY SUM(oi.quantity * oi.list_price * oi.discount) DESC LIMIT 1",
            False,
        ),
    ]
    
    all_pass = True
    for question, sql, should_pass in test_cases:
        is_valid, err = validator.validate_semantic(sql, question)
        passed = (is_valid == should_pass)
        status = "PASS" if passed else f"FAIL (expected {'valid' if should_pass else 'invalid'})"
        print(f"  '{question[:30]}...' -> {status}")
        if not passed:
            all_pass = False
    
    return all_pass


def test_prompt_builder() -> bool:
    """Test prompt builder query classification."""
    print("\nTesting prompt builder...")
    from src.prompts import PromptBuilder
    
    pb = PromptBuilder("TEST SCHEMA")
    
    cases = [
        ("How many customers are there?", "count"),
        ("Top 5 products", "ranking"),
        ("Total revenue", "aggregation"),
        ("List all staff", "join"),
        ("Products never ordered", "null_handling"),
        ("Monthly revenue trend", "date_range"),
        ("Find customers from NY", "filter"),
    ]
    
    all_pass = True
    for q, expected in cases:
        result = pb.classify_query(q)
        passed = result == expected
        status = "OK" if passed else f"got {result}"
        print(f"  '{q}' -> {status}")
        if not passed:
            all_pass = False
    
    return all_pass


def test_conversation_manager() -> bool:
    """Test follow-up detection and standalone rewrite behavior."""
    print("\nTesting conversation manager...")
    from src.conversation import ConversationManager

    conv = ConversationManager(enabled=True, max_turns=6, use_llm_rewrite=False, preview_rows=2)

    # Seed context with a customer-focused turn.
    conv.record_turn(
        original_question="Which customer has the longest name?",
        effective_question="Which customer has the longest name?",
        sql="SELECT first_name, last_name FROM customers ORDER BY length(first_name)+length(last_name) DESC LIMIT 1",
        results=[("Christopher", "Richardson")],
        query_type="ranking",
        tables=["customers"],
    )

    follow_up = "What did he buy?"
    detected = conv.detect_follow_up(follow_up)
    rewritten, meta = conv.rewrite_query(follow_up, llm_client=None, model=None)

    detect_pass = detected is True
    rewrite_pass = rewritten.lower() == "what products did Christopher Richardson buy?".lower()
    print(f"  Follow-up detection: {'PASS' if detect_pass else 'FAIL'}")
    print(f"  Heuristic rewrite: {'PASS' if rewrite_pass else 'FAIL'}")

    # Multi-turn pronoun resolution should keep person/product references.
    conv.record_turn(
        original_question="What did he buy?",
        effective_question="What products did Christopher Richardson buy?",
        sql="SELECT DISTINCT p.product_name FROM customers c JOIN orders o ON c.customer_id = o.customer_id JOIN order_items oi ON o.order_id = oi.order_id JOIN products p ON oi.product_id = p.product_id WHERE c.first_name = 'Christopher' AND c.last_name = 'Richardson' ORDER BY p.product_name",
        results=[("Trek Domane AL 2 Women's - 2018",)],
        query_type="join",
        tables=["customers", "orders", "order_items", "products"],
    )
    conv.record_turn(
        original_question="what shop was that from",
        effective_question="What shop sold the Trek Domane AL 2 Women's - 2018 bike?",
        sql="SELECT s.store_name FROM products p JOIN stocks sk ON p.product_id = sk.product_id JOIN stores s ON sk.store_id = s.store_id WHERE p.product_name = 'Trek Domane AL 2 Women''s - 2018'",
        results=[("Santa Cruz Bikes",), ("Baldwin Bikes",), ("Rowlett Bikes",)],
        query_type="join",
        tables=["products", "stocks", "stores"],
    )
    chained_follow_up = "which one did he buy it from"
    chained_rewrite, _ = conv.rewrite_query(chained_follow_up, llm_client=None, model=None)
    chained_pass = chained_rewrite.lower() == "which store did Christopher Richardson buy Trek Domane AL 2 Women's - 2018 from?".lower()
    print(f"  Multi-turn person/product rewrite: {'PASS' if chained_pass else 'FAIL'}")

    conv.record_turn(
        original_question="which one did he buy it from",
        effective_question="Which store did Christopher Richardson buy Trek Domane AL 2 Women's - 2018 from?",
        sql="SELECT s.store_name FROM customers c JOIN orders o ON c.customer_id = o.customer_id JOIN stores s ON o.store_id = s.store_id JOIN order_items oi ON o.order_id = oi.order_id JOIN products p ON oi.product_id = p.product_id WHERE c.first_name = 'Christopher' AND c.last_name = 'Richardson' AND p.product_name = 'Trek Domane AL 2 Women''s - 2018' ORDER BY o.order_date DESC LIMIT 1",
        results=[("Santa Cruz Bikes",)],
        query_type="join",
        tables=["customers", "orders", "stores", "order_items", "products"],
    )
    store_follow_up = "how many sales did that shop get"
    store_rewrite, _ = conv.rewrite_query(store_follow_up, llm_client=None, model=None)
    store_pass = store_rewrite.lower() == "how many sales did Santa Cruz Bikes get?".lower()
    print(f"  Demonstrative shop rewrite: {'PASS' if store_pass else 'FAIL'}")

    store_alias_follow_up = "how many sales did that outlet get"
    store_alias_rewrite, _ = conv.rewrite_query(store_alias_follow_up, llm_client=None, model=None)
    store_alias_pass = store_alias_rewrite.lower() == "how many sales did Santa Cruz Bikes get?".lower()
    print(f"  Demonstrative outlet rewrite: {'PASS' if store_alias_pass else 'FAIL'}")

    which_alias_follow_up = "which branch did he buy it from"
    which_alias_rewrite, _ = conv.rewrite_query(which_alias_follow_up, llm_client=None, model=None)
    which_alias_pass = which_alias_rewrite.lower() == "which store did Christopher Richardson buy Trek Domane AL 2 Women's - 2018 from?".lower()
    print(f"  Which-branch rewrite: {'PASS' if which_alias_pass else 'FAIL'}")

    conv.record_turn(
        original_question="Which customer got the most discounts?",
        effective_question="Which customer got the most discounts?",
        sql="SELECT c.first_name, c.last_name, SUM(oi.quantity * oi.list_price * oi.discount) AS total_discount FROM customers c JOIN orders o ON c.customer_id = o.customer_id JOIN order_items oi ON o.order_id = oi.order_id GROUP BY c.customer_id, c.first_name, c.last_name ORDER BY total_discount DESC LIMIT 1",
        results=[("Abby", "Gamble", 4697.8838)],
        query_type="ranking",
        tables=["customers", "orders", "order_items"],
    )
    discount_follow_up = "how much did she receive?"
    discount_rewrite, _ = conv.rewrite_query(discount_follow_up, llm_client=None, model=None)
    discount_pass = discount_rewrite.lower() == "what is the total discount amount received by Abby Gamble?".lower()
    print(f"  Discount follow-up rewrite: {'PASS' if discount_pass else 'FAIL'}")

    # Unresolvable follow-up should not fabricate context.
    conv.clear()
    unresolved, unresolved_meta = conv.rewrite_query("What about him?", llm_client=None, model=None)
    unresolved_pass = unresolved == "What about him?" and unresolved_meta.get("rewritten") is False
    print(f"  Unresolved safety fallback: {'PASS' if unresolved_pass else 'FAIL'}")

    return (
        detect_pass
        and rewrite_pass
        and chained_pass
        and store_pass
        and store_alias_pass
        and which_alias_pass
        and discount_pass
        and unresolved_pass
    )


def test_query_intent_filter() -> bool:
    """Test CLI intent filter allows realistic DB questions."""
    print("\nTesting query intent filter...")
    from main import is_valid_query

    cases = [
        ("What shop sold the Trek Domane AL 2 Women's - 2018 bike?", True),
        ("What did Christopher Richardson buy?", True),
        ("hello there", False),
        ("who are you", False),
    ]

    all_pass = True
    for text, expected in cases:
        result = is_valid_query(text)
        passed = result == expected
        status = "PASS" if passed else f"FAIL (got {result}, expected {expected})"
        print(f"  '{text}' -> {status}")
        if not passed:
            all_pass = False

    return all_pass


def test_model_fallback_helpers() -> bool:
    """Test local model fallback helper logic without network dependency."""
    print("\nTesting model fallback helpers...")
    from src.model_fallback import choose_preferred_model, ensure_model_configured

    selected = choose_preferred_model(["mistral:latest", "llama3.2:latest"])
    select_pass = selected == "llama3.2:latest"
    print(f"  Preferred model selection: {'PASS' if select_pass else 'FAIL'}")

    # Deterministic path: preconfigured model should short-circuit.
    original_model = os.environ.get("OLLAMA_MODEL")
    original_host = os.environ.get("OLLAMA_HOST")
    original_carleton_host = os.environ.get("CARLETON_OLLAMA_HOST")
    original_carleton_default = os.environ.get("CARLETON_DEFAULT_MODEL")

    os.environ["OLLAMA_MODEL"] = "unit-test-model"
    ok, msg = ensure_model_configured(interactive=False)
    configured_pass = ok and "unit-test-model" in msg
    print(f"  Configured model short-circuit: {'PASS' if configured_pass else 'FAIL'}")

    # Carleton host configured + no model -> should default to llama3.3
    os.environ.pop("OLLAMA_MODEL", None)
    os.environ["OLLAMA_HOST"] = "https://rcsllm.carleton.ca/rcsapi"
    os.environ["CARLETON_OLLAMA_HOST"] = "https://rcsllm.carleton.ca/rcsapi"
    os.environ["CARLETON_DEFAULT_MODEL"] = "llama3.3"
    ok, msg = ensure_model_configured(interactive=False)
    remote_default_pass = ok and (os.environ.get("OLLAMA_MODEL") == "llama3.3")
    print(f"  Remote default model policy: {'PASS' if remote_default_pass else 'FAIL'}")

    # Restore env state
    if original_model is None:
        os.environ.pop("OLLAMA_MODEL", None)
    else:
        os.environ["OLLAMA_MODEL"] = original_model

    if original_host is None:
        os.environ.pop("OLLAMA_HOST", None)
    else:
        os.environ["OLLAMA_HOST"] = original_host

    if original_carleton_host is None:
        os.environ.pop("CARLETON_OLLAMA_HOST", None)
    else:
        os.environ["CARLETON_OLLAMA_HOST"] = original_carleton_host

    if original_carleton_default is None:
        os.environ.pop("CARLETON_DEFAULT_MODEL", None)
    else:
        os.environ["CARLETON_DEFAULT_MODEL"] = original_carleton_default

    return select_pass and configured_pass and remote_default_pass


def test_unanswerable_and_guardrails() -> bool:
    """Test refusal detection and final SQL contract guardrails without LLM calls."""
    print("\nTesting unanswerable detection and guardrails...")
    from agent import QueryWriter

    class StubSchemaContext:
        def __init__(self, matched: bool):
            self.matched = matched

        def route_tables(self, question: str, fallback_to_all: bool = True):
            if self.matched:
                return ['orders', 'customers'], True
            if fallback_to_all:
                return ['brands', 'categories', 'customers', 'order_items', 'orders',
                        'products', 'staffs', 'stocks', 'stores'], False
            return [], False

        def get_table_names(self):
            return ['brands', 'categories', 'customers', 'order_items', 'orders',
                    'products', 'staffs', 'stocks', 'stores']

    class GuardrailValidator:
        def __init__(self, syntax_ok: bool = True, quick_ok: bool = True):
            self.syntax_ok = syntax_ok
            self.quick_ok = quick_ok
            self.quick_calls = 0

        def validate(self, sql: str):
            return self.syntax_ok, "" if self.syntax_ok else "syntax error"

        def execute_quick(self, sql: str):
            self.quick_calls += 1
            return self.quick_ok, "" if self.quick_ok else "runtime error"

    all_pass = True

    # Case 1: obvious off-topic should refuse immediately
    writer = QueryWriter.__new__(QueryWriter)
    writer.schema_ctx = StubSchemaContext(matched=True)
    writer._llm_unanswerable_check = lambda q: False
    case_1 = writer._is_unanswerable("What is the capital of France?")
    print(f"  Off-topic refusal test: {'PASS' if case_1 else 'FAIL'}")
    all_pass = all_pass and case_1

    # Case 2: no schema match + LLM says unanswerable
    writer = QueryWriter.__new__(QueryWriter)
    writer.schema_ctx = StubSchemaContext(matched=False)
    writer._llm_unanswerable_check = lambda q: True
    case_2 = writer._is_unanswerable("Who won the world cup in 1998?")
    print(f"  LLM fallback unanswerable test: {'PASS' if case_2 else 'FAIL'}")
    all_pass = all_pass and case_2

    # Case 3: no schema match + on-topic cues + LLM unavailable -> do not refuse
    writer = QueryWriter.__new__(QueryWriter)
    writer.schema_ctx = StubSchemaContext(matched=False)
    writer._llm_unanswerable_check = lambda q: None
    case_3 = writer._is_unanswerable("Show city and state for each store")
    print(f"  Conservative fallback test: {'PASS' if not case_3 else 'FAIL'}")
    all_pass = all_pass and (not case_3)

    # Case 4: single-word prompts should be refused as underspecified
    writer = QueryWriter.__new__(QueryWriter)
    writer.schema_ctx = StubSchemaContext(matched=True)
    writer._llm_unanswerable_check = lambda q: False
    case_4 = writer._is_unanswerable("Orders")
    print(f"  Single-word refusal test: {'PASS' if case_4 else 'FAIL'}")
    all_pass = all_pass and case_4

    # Guardrail contract checks
    writer = QueryWriter.__new__(QueryWriter)
    writer.execution_validation = True
    writer.validator = GuardrailValidator(syntax_ok=True, quick_ok=True)

    good_sql = writer._enforce_output_contract("SELECT * FROM customers LIMIT 1")
    good_pass = good_sql.startswith("SELECT")
    print(f"  Guardrail valid SQL test: {'PASS' if good_pass else 'FAIL'}")
    all_pass = all_pass and good_pass

    bad_prefix = writer._enforce_output_contract("DROP TABLE customers")
    prefix_pass = bad_prefix == writer.REFUSAL_SQL
    print(f"  Guardrail prefix test: {'PASS' if prefix_pass else 'FAIL'}")
    all_pass = all_pass and prefix_pass

    writer.validator = GuardrailValidator(syntax_ok=False, quick_ok=True)
    bad_syntax = writer._enforce_output_contract("SELECT * FROM customers")
    syntax_pass = bad_syntax == writer.REFUSAL_SQL
    print(f"  Guardrail syntax test: {'PASS' if syntax_pass else 'FAIL'}")
    all_pass = all_pass and syntax_pass

    writer.validator = GuardrailValidator(syntax_ok=True, quick_ok=False)
    bad_runtime = writer._enforce_output_contract("SELECT * FROM customers")
    runtime_pass = bad_runtime == writer.REFUSAL_SQL
    print(f"  Guardrail runtime test: {'PASS' if runtime_pass else 'FAIL'}")
    all_pass = all_pass and runtime_pass

    return all_pass


def test_fast_path_execution_validation() -> bool:
    """Test that fast path runs execution-based validation when enabled."""
    print("\nTesting fast-path execution validation...")
    from agent import QueryWriter

    class StubPromptBuilder:
        @staticmethod
        def build_minimal_prompt(query_type=None, tables=None):
            return "TEST PROMPT"

    class StubFastValidator:
        def __init__(self):
            self.exec_calls = 0

        def validate(self, sql: str):
            return True, ""

        def validate_semantic(self, sql: str, question: str):
            return True, ""

        def execute_and_validate(self, sql: str, question: str):
            self.exec_calls += 1
            return True, [], ""

    writer = QueryWriter.__new__(QueryWriter)
    writer.prompt_builder = StubPromptBuilder()
    writer.validator = StubFastValidator()
    writer.execution_validation = True
    writer._call_llm = lambda system_prompt, user_prompt, temperature=0: "SELECT COUNT(*) FROM customers"
    writer._retry_with_error = lambda question, bad_sql, error, system_prompt: "SELECT COUNT(*) FROM customers"

    result = writer._fast_generate(
        prompt="How many customers are there?",
        query_type="count",
        tables=["customers"]
    )

    passed = writer.validator.exec_calls > 0 and result.sql.startswith("SELECT")
    print(f"  Fast-path execution validation call test: {'PASS' if passed else 'FAIL'}")
    return passed


def run_test_case(agent, test_case: TestCase, verbose: bool = False) -> Dict:
    """Run a single test case and return results."""
    start_time = time.time()
    result = {
        'question': test_case.question,
        'category': test_case.category,
        'difficulty': test_case.difficulty,
        'passed': False,
        'sql': None,
        'error': None,
        'execution_time': 0,
        'checks': []
    }
    
    try:
        sql = agent.generate_query(test_case.question)
        result['sql'] = sql
        result['execution_time'] = time.time() - start_time
        
        # Check syntax
        is_valid, error = agent.validator.validate(sql)
        if not is_valid:
            result['error'] = f"Syntax error: {error}"
            return result
        
        # Check should_contain keywords
        sql_lower = sql.lower()
        for keyword in test_case.should_contain:
            if keyword.lower() not in sql_lower:
                result['checks'].append(f"Missing: {keyword}")
        
        # Check should_not_contain keywords
        for keyword in test_case.should_not_contain:
            if keyword.lower() in sql_lower:
                result['checks'].append(f"Should not contain: {keyword}")
        
        # Query is valid and passes all checks
        if not result['checks']:
            result['passed'] = True
        
    except Exception as e:
        result['error'] = str(e)
        result['execution_time'] = time.time() - start_time
    
    return result


def run_full_test_suite(verbose: bool = False) -> Dict:
    """Run the complete test suite against the agent."""
    print("\n" + "=" * 60)
    print("FULL TEST SUITE")
    print("=" * 60)
    
    from agent import QueryWriter
    from db.bike_store import BikeStoreDb
    
    host = os.getenv('OLLAMA_HOST', 'http://localhost:11434')
    model = os.getenv('OLLAMA_MODEL', 'llama3.2')
    print(f"Using: {host} with model {model}")
    
    # Initialize
    print("\nInitializing database and agent...")
    db = BikeStoreDb('bike_store.db')
    agent = QueryWriter('bike_store.db')
    
    # Run tests
    results = {
        'total': len(TEST_CASES),
        'passed': 0,
        'failed': 0,
        'by_category': {},
        'by_difficulty': {},
        'details': []
    }
    
    print(f"\nRunning {len(TEST_CASES)} test cases...\n")
    
    for i, test_case in enumerate(TEST_CASES, 1):
        result = run_test_case(agent, test_case, verbose)
        results['details'].append(result)
        
        if result['passed']:
            results['passed'] += 1
            status = "PASS"
        else:
            results['failed'] += 1
            status = "FAIL"
        
        # Track by category
        cat = test_case.category
        if cat not in results['by_category']:
            results['by_category'][cat] = {'passed': 0, 'failed': 0}
        if result['passed']:
            results['by_category'][cat]['passed'] += 1
        else:
            results['by_category'][cat]['failed'] += 1
        
        # Track by difficulty
        diff = test_case.difficulty
        if diff not in results['by_difficulty']:
            results['by_difficulty'][diff] = {'passed': 0, 'failed': 0}
        if result['passed']:
            results['by_difficulty'][diff]['passed'] += 1
        else:
            results['by_difficulty'][diff]['failed'] += 1
        
        # Print progress
        if verbose or not result['passed']:
            print(f"[{i}/{len(TEST_CASES)}] {status}: {test_case.question[:50]}...")
            if not result['passed'] and result['error']:
                print(f"         Error: {result['error'][:80]}")
            if not result['passed'] and result['checks']:
                print(f"         Checks: {result['checks']}")
        else:
            print(f"[{i}/{len(TEST_CASES)}] {status}: {test_case.question[:50]}...")
    
    # Print summary
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    print(f"Total: {results['total']} | Passed: {results['passed']} | Failed: {results['failed']}")
    print(f"Success Rate: {results['passed']/results['total']*100:.1f}%")
    
    print("\nBy Category:")
    for cat, stats in sorted(results['by_category'].items()):
        total = stats['passed'] + stats['failed']
        pct = stats['passed']/total*100 if total > 0 else 0
        print(f"  {cat}: {stats['passed']}/{total} ({pct:.0f}%)")
    
    print("\nBy Difficulty:")
    for diff in ['easy', 'medium', 'hard']:
        if diff in results['by_difficulty']:
            stats = results['by_difficulty'][diff]
            total = stats['passed'] + stats['failed']
            pct = stats['passed']/total*100 if total > 0 else 0
            print(f"  {diff}: {stats['passed']}/{total} ({pct:.0f}%)")
    
    return results


def run_benchmark(models: List[str] = None) -> Dict:
    """Benchmark the agent across different models."""
    print("\n" + "=" * 60)
    print("MODEL BENCHMARK")
    print("=" * 60)
    
    if models is None:
        models = ['llama3.2', 'llama3.3', 'qwen2.5-coder']
    
    from agent import QueryWriter
    from db.bike_store import BikeStoreDb
    
    # Sample test cases for benchmarking (subset for speed)
    benchmark_cases = [tc for tc in TEST_CASES if tc.difficulty in ['easy', 'medium']][:20]
    
    results = {}
    original_model = os.getenv('OLLAMA_MODEL', 'llama3.2')
    
    for model in models:
        print(f"\nBenchmarking model: {model}")
        os.environ['OLLAMA_MODEL'] = model
        
        try:
            db = BikeStoreDb('bike_store.db')
            agent = QueryWriter('bike_store.db')
            
            model_results = {
                'passed': 0,
                'failed': 0,
                'total_time': 0,
                'avg_time': 0
            }
            
            for tc in benchmark_cases:
                result = run_test_case(agent, tc, verbose=False)
                if result['passed']:
                    model_results['passed'] += 1
                else:
                    model_results['failed'] += 1
                model_results['total_time'] += result['execution_time']
            
            model_results['avg_time'] = model_results['total_time'] / len(benchmark_cases)
            results[model] = model_results
            
            total = model_results['passed'] + model_results['failed']
            pct = model_results['passed']/total*100
            print(f"  Results: {model_results['passed']}/{total} ({pct:.0f}%)")
            print(f"  Avg time: {model_results['avg_time']:.2f}s")
            
        except Exception as e:
            print(f"  Error with model {model}: {e}")
            results[model] = {'error': str(e)}
    
    # Restore original model
    os.environ['OLLAMA_MODEL'] = original_model
    
    # Print comparison
    print("\n" + "-" * 40)
    print("COMPARISON")
    print("-" * 40)
    for model, res in results.items():
        if 'error' in res:
            print(f"{model}: ERROR")
        else:
            total = res['passed'] + res['failed']
            pct = res['passed']/total*100
            print(f"{model}: {pct:.0f}% accuracy, {res['avg_time']:.2f}s avg")
    
    return results


# =============================================================================
# MAIN
# =============================================================================

if __name__ == "__main__":
    print("=" * 60)
    print("SQL Query Writer Agent - Comprehensive Test Suite")
    print("=" * 60)
    
    args = sys.argv[1:]
    verbose = '--verbose' in args or '-v' in args
    
    # Always run basic tests
    if not test_imports():
        sys.exit(1)
    
    basic_checks = [
        test_sql_extraction(),
        test_conversation_manager(),
        test_query_intent_filter(),
        test_model_fallback_helpers(),
        test_unanswerable_and_guardrails(),
        test_fast_path_execution_validation(),
    ]
    if not all(basic_checks):
        print("\nOne or more basic tests failed.")
        sys.exit(1)
    
    # Full tests require database and LLM
    if '--full' in args:
        test_schema_context()
        test_validator()
        test_semantic_validator()
        test_prompt_builder()
        
        results = run_full_test_suite(verbose=verbose)
        
        # Save results to file
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        results_file = f"test_results_{timestamp}.json"
        with open(results_file, 'w') as f:
            json.dump(results, f, indent=2, default=str)
        print(f"\nResults saved to {results_file}")
    
    elif '--benchmark' in args:
        test_schema_context()
        
        # Parse models from command line or use defaults
        models = None
        for i, arg in enumerate(args):
            if arg == '--models' and i + 1 < len(args):
                models = args[i + 1].split(',')
        
        run_benchmark(models)
    
    else:
        print("\nBasic tests passed!")
        print("\nOptions:")
        print("  --full       Run full test suite with 50+ queries")
        print("  --benchmark  Benchmark different models")
        print("  --verbose    Show detailed output")
        print("  --models     Comma-separated list of models (with --benchmark)")
        print("\n(Full tests require Kaggle credentials and Ollama)")