"""
Test JWT token generation script for MCP Guardian.

This script generates valid JWT tokens for testing the protected endpoints
of the MCP Guardian demo application.
"""

import jwt
import time
from datetime import datetime, timedelta, timezone


# Private key corresponding to the public key in config.yaml
PRIVATE_KEY = """-----BEGIN PRIVATE KEY-----
MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDeyn9yHQ/T47GG
SdLBi3J46XeNjL1UJ0wdERxSIKbDyzyaXBMw5kFcVvjY3Mit1bpR+SIFHRZn8wON
txPHnLeEOubInfLI5/8txCrLgoq9f8SdNPQ1d8/ydvyc0HKm/85otvkVM3/RzfmN
nizxYF1r1I1QME60l/mzJgYBLv0Vc0oA3L4MX+fHYjhWKs805x8lkR9P8HNj8iI1
FTqjErR1DRHepbHLJkP3HHDkqmfr8tWbx16yUXtds+Aym8u4Zl5+rDN3ctX9Re07
aPOvusZxTUwivEgLkNhBDao/ie4mf/6Q4O32bwOV0pg6BqOJxUTw4x22dbcecr3X
G74V8JiLAgMBAAECggEAVMlBuT7gICZeEBoWL/wOw38uNVtjThxqdgFmAsdUT90P
ojr6PrjSFP8wUu+rG6Tj9IBYdxO7zDGKn6sVxviYSqVT4PNjSxwnW9Z05uA2RqnS
PHXDnnmLVN39gvw9NdrSRb2ki3H+6ex2VXPl7LJcu0W3xdWsjkPsDkahMGAehsOU
btR/nrR4GzaYwQ1UyjKsLSm0JVw8RufPjMV4gRNzQ/rbS4bTMXZP5Smu14Q9dhgn
O+3mkYUAoTtZ3blY2z3xLfFWAsthLFr0cH/IxCapKhkPLl/nI46EpSqnpqABcITb
oFAjEpuhoTaMl2qwo4HkmkB1IS4sgIhaLL14E5EtsQKBgQD4841Dh3lIpdYELM37
AexPQn3px3gbDWyAfpbNg4zA0yt2R8M1R60rLOkwxcigK3TJJc4cVsLKghCFHgAd
rXO9rtghxibfbO4qSx2FaB0HwTugf9VpbVLw1V6KYreOxO+O95KeWMcowrZiC1hF
Hkk/cPVzEKbLU67Z+D21AIldYwKBgQDlGVSgaZ4d+CUUVfhVYRT2Ihk9xkxA8cHY
4ynd4/sf4yPUaH/cPrp0q5i7mOG1aqV4XiYmWj969UdA/5MVLpESmiloBgte5wM6
OYxRh8Nha1PE6aRaJ1L0Lqui8bFFEYqQNN92tGI66N5XRkk0PfzrISM97sinIz6K
UxUgHbw0uQKBgQCQBn9UcRz2rZKKac1lCTDIjgRAarkOlY6PLC1cil209LZxQIUr
3uLNVrc2E0bFez6eXHPIq10YQXUDPl78zGtocq3xpCy7F88Vws06cAhTgzRW3338
91WDpT2DvCbLvx2EzbASUfQuna9cO4n+us3PmS83OYv1l91Rly9D9yKBVwKBgQCW
M1hHX1d6AJGW1pJcr6yFfm+cD6E4xwe3QZ/WST0Ad/Oqs+R4KmwmR8Z339DDFytr
Uhai+HI8FAlmNOlibshQMOBKwgO8caqtBsPXRDI1cZ9XL3NAZB2/18VRrDHzvX9h
8jg75CGuy7PR9lPSO33RTkWDukS1R6ceG0BzcOtOiQKBgAyaFkhgEKjgoOry3CA3
s7tbPDlJKohusTOyyVWdM1Ya7udmt9Wr6t/FKg/RGkq3zUWCJJShbPRdoeMlhuZn
B5Hb/tgYvVlVxNRfwfp7CXSEte3kN2rTCBhXM97vXspvWrVFhegme6VdAzLnQb2r
PeWmOPrH8GMMqpXCe35+8KC7
-----END PRIVATE KEY-----"""


def generate_test_token(
    user_id: str = "test-user-123",
    scopes: list = None,
    expires_in_hours: int = 24
) -> str:
    """
    Generate a test JWT token for MCP Guardian.
    
    Args:
        user_id: User identifier (sub claim)
        scopes: List of scopes the user has access to
        expires_in_hours: Token expiration time in hours
        
    Returns:
        Encoded JWT token string
    """
    if scopes is None:
        scopes = ["tools:calculator:execute", "tools:weather:read"]
    
    # Calculate expiration time
    exp_time = datetime.now(timezone.utc) + timedelta(hours=expires_in_hours)
    
    # Create token payload
    payload = {
        "sub": user_id,
        "aud": "mcp-tools-api",
        "scopes": scopes,
        "exp": int(exp_time.timestamp()),
        "iat": int(time.time()),
        "iss": "mcp-guardian-test"
    }
    
    # Generate token
    token = jwt.encode(
        payload,
        PRIVATE_KEY,
        algorithm="RS256"
    )
    
    return token


def generate_invalid_tokens():
    """
    Generate various invalid tokens for testing error handling.
    
    Returns:
        Dictionary of invalid tokens with descriptions
    """
    invalid_tokens = {}
    
    # Expired token
    expired_payload = {
        "sub": "test-user-123",
        "aud": "mcp-tools-api",
        "scopes": ["tools:calculator:execute"],
        "exp": int(time.time()) - 3600,  # Expired 1 hour ago
        "iat": int(time.time()) - 7200
    }
    invalid_tokens["expired"] = {
        "token": jwt.encode(expired_payload, PRIVATE_KEY, algorithm="RS256"),
        "description": "Token that expired 1 hour ago"
    }
    
    # Wrong audience
    wrong_aud_payload = {
        "sub": "test-user-123",
        "aud": "wrong-audience",
        "scopes": ["tools:calculator:execute"],
        "exp": int(time.time()) + 3600
    }
    invalid_tokens["wrong_audience"] = {
        "token": jwt.encode(wrong_aud_payload, PRIVATE_KEY, algorithm="RS256"),
        "description": "Token with wrong audience"
    }
    
    # Insufficient scopes
    insufficient_scopes_payload = {
        "sub": "test-user-123",
        "aud": "mcp-tools-api",
        "scopes": ["tools:weather:read"],  # Missing calculator scope
        "exp": int(time.time()) + 3600
    }
    invalid_tokens["insufficient_scopes"] = {
        "token": jwt.encode(insufficient_scopes_payload, PRIVATE_KEY, algorithm="RS256"),
        "description": "Token with insufficient scopes (missing tools:calculator:execute)"
    }
    
    return invalid_tokens


def main():
    """Main function to demonstrate token generation."""
    print("MCP Guardian - Test Token Generator")
    print("=" * 50)
    
    # Generate valid token
    valid_token = generate_test_token()
    print("\n✅ Valid Token (for testing protected endpoints):")
    print(f"Token: {valid_token}")
    print("\nUse this token in your requests:")
    print(f"Authorization: Bearer {valid_token}")
    
    # Generate invalid tokens
    invalid_tokens = generate_invalid_tokens()
    print("\n❌ Invalid Tokens (for testing error handling):")
    for token_type, token_info in invalid_tokens.items():
        print(f"\n{token_type.upper()}:")
        print(f"Description: {token_info['description']}")
        print(f"Token: {token_info['token']}")
    
    print("\n" + "=" * 50)
    print("Testing Instructions:")
    print("1. Start the demo server: python main.py")
    print("2. Test public endpoint: curl http://localhost:8000/")
    print("3. Test protected endpoint with valid token:")
    print(f"   curl -X POST http://localhost:8000/tools/calculate \\")
    print(f"        -H 'Authorization: Bearer {valid_token}' \\")
    print(f"        -H 'Content-Type: application/json' \\")
    print(f"        -d '{{\"expression\": \"5 * 3\"}}'")
    print("4. Test with invalid tokens to verify error handling")


if __name__ == "__main__":
    main()
