"""
Test that error responses include prompt and token diagnostics.
"""

import sys
sys.path.insert(0, 'src')

from ai_med_extract.utils.unified_model_manager import ModelError, check_token_limits
from ai_med_extract.services.error_handler import create_error_response

def test_error_includes_prompt_and_tokens():
    """Test that ModelError preserves prompt and token info in error responses"""
    print("Testing error response includes prompt and token diagnostics...\n")
    
    # Create a test prompt
    test_prompt = "This is a test patient summary prompt. " * 100  # ~4000 chars
    model_name = "microsoft/Phi-3-mini-4k-instruct"
    
    # Get token info
    token_check = check_token_limits(test_prompt, model_name, reserve_for_output=2048)
    
    # Create a ModelError with prompt and token info
    error = ModelError(
        model_name=model_name,
        error_type="token_limit_exceeded",
        details="Test error",
        prompt=test_prompt,
        token_info=token_check
    )
    
    # Create error response
    response = create_error_response(error)
    
    # Verify response includes prompt and token info
    print("Error Response Structure:")
    print(f"  error_type: {response['error_type']}")
    print(f"  error_code: {response['error_code']}")
    print(f"  status: {response['status']}")
    print(f"\nDetails included:")
    
    assert "details" in response, "Response missing 'details' field"
    details = response["details"]
    
    # Check for prompt information
    if "prompt_preview" in details:
        print(f"  [OK] prompt_preview: {details['prompt_preview'][:80]}...")
        print(f"  [OK] prompt_length: {details['prompt_length']} characters")
        print(f"  [OK] full_prompt: Available ({len(details.get('full_prompt', ''))} chars)")
    else:
        print("  [MISSING] prompt information missing")
        
    # Check for token diagnostics
    if "token_diagnostics" in details:
        token_diag = details["token_diagnostics"]
        print(f"\n  Token Diagnostics:")
        print(f"    - estimated_tokens: {token_diag.get('estimated_tokens')}")
        print(f"    - max_tokens: {token_diag.get('max_tokens')}")
        print(f"    - available_for_input: {token_diag.get('available_for_input')}")
        print(f"    - usage_percentage: {token_diag.get('usage_percentage'):.1f}%")
        print(f"    - within_limit: {token_diag.get('within_limit')}")
    else:
        print("  [MISSING] token_diagnostics missing")
    
    # Assertions
    assert "prompt_preview" in details, "Missing prompt_preview"
    assert "prompt_length" in details, "Missing prompt_length"
    assert "full_prompt" in details, "Missing full_prompt"
    assert "token_diagnostics" in details, "Missing token_diagnostics"
    
    assert details["prompt_length"] == len(test_prompt), "Prompt length mismatch"
    assert details["full_prompt"] == test_prompt, "Full prompt mismatch"
    
    token_diag = details["token_diagnostics"]
    assert "estimated_tokens" in token_diag, "Missing estimated_tokens"
    assert "max_tokens" in token_diag, "Missing max_tokens"
    assert "available_for_input" in token_diag, "Missing available_for_input"
    
    print("\n[PASS] Error response includes all prompt and token diagnostics!")
    return True

if __name__ == "__main__":
    print("="*60)
    print("Error Response Diagnostics Test")
    print("="*60 + "\n")
    
    try:
        test_error_includes_prompt_and_tokens()
        print("\n" + "="*60)
        print("[SUCCESS] All tests passed!")
        print("="*60)
        print("\nError responses now include:")
        print("  - Full prompt for debugging")
        print("  - Prompt preview (first 500 chars)")
        print("  - Prompt length in characters")
        print("  - Complete token diagnostics")
        
    except AssertionError as e:
        print(f"\n[FAILED] {e}")
        sys.exit(1)
    except Exception as e:
        print(f"\n[ERROR] {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)