"""
Test script to verify Phi-3.5-mini-instruct token limit configuration
"""

import sys
import os

# Add the parent directory to the path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from ai_med_extract.utils.model_config import get_model_token_limit

def test_phi35_token_limits():
    """Test that Phi-3.5 models are correctly configured with 128K context"""
    
    test_cases = [
        ("microsoft/Phi-3.5-mini-instruct", 131072, "Phi-3.5-mini-instruct"),
        ("microsoft/Phi-3.5-MoE-instruct", 131072, "Phi-3.5-MoE-instruct"),
        ("microsoft/Phi-3.5-vision-instruct", 131072, "Phi-3.5-vision-instruct"),
        ("microsoft/Phi-3-mini-4k-instruct", 8192, "Phi-3-mini-4k (should be 8K)"),
        ("microsoft/Phi-3-mini-128k-instruct", 131072, "Phi-3-mini-128k"),
    ]
    
    print("Testing Phi-3.5 Token Limit Configuration")
    print("=" * 60)
    
    all_passed = True
    for model_name, expected_limit, description in test_cases:
        actual_limit = get_model_token_limit(model_name)
        status = "[PASS]" if actual_limit == expected_limit else "[FAIL]"
        
        if actual_limit != expected_limit:
            all_passed = False
            
        print(f"{status} | {description}")
        print(f"       Model: {model_name}")
        print(f"       Expected: {expected_limit:,} tokens")
        print(f"       Actual:   {actual_limit:,} tokens")
        print()
    
    print("=" * 60)
    if all_passed:
        print("[SUCCESS] All tests passed!")
        return 0
    else:
        print("[ERROR] Some tests failed!")
        return 1

if __name__ == "__main__":
    exit_code = test_phi35_token_limits()
    sys.exit(exit_code)