#!/usr/bin/env python3
"""
Quick test script to verify the system works with small models.
"""

import os
import sys
from langchain_models import langchain_models_registry
from custom_evaluator import custom_evaluator

def test_smallest_model():
    """Test with the smallest available model."""
    print("🚀 Testing with smallest model (DistilGPT-2)...")
    
    # Get the smallest model
    model_config = langchain_models_registry.get_model_config("DistilGPT-2")
    if not model_config:
        print("❌ DistilGPT-2 model not found")
        return False
    
    print(f"📋 Model: {model_config.name}")
    print(f"📋 Model ID: {model_config.model_id}")
    
    try:
        # Create the model
        print("📥 Creating model...")
        model = langchain_models_registry.create_langchain_model(model_config)
        print("✅ Model created successfully")
        
        # Test SQL generation
        print("🔍 Testing SQL generation...")
        prompt_template = """
You are an expert SQL developer.

Database Schema:
{schema}

Question: {question}

Generate a SQL query:
"""
        
        schema = "-- NYC Taxi Dataset\nCREATE TABLE trips (id INT, fare_amount FLOAT, total_amount FLOAT);"
        question = "How many trips are there?"
        
        result = langchain_models_registry.generate_sql(
            model_config, prompt_template, schema, question
        )
        
        print(f"📝 Generated SQL: {result}")
        
        if result and len(result) > 10:
            print("✅ SQL generation successful!")
            return True
        else:
            print("⚠️ SQL generation produced short result")
            return False
            
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

if __name__ == "__main__":
    success = test_smallest_model()
    if success:
        print("\n🎉 System is working! Ready to run full evaluation.")
    else:
        print("\n❌ System needs fixes.")
    sys.exit(0 if success else 1)