# server:

from fastmcp import FastMCP
from fastmcp.server.dependencies import get_http_headers

mcp = FastMCP("TestServer")

@mcp.tool
def test_token():
    """Test tool that reads HTTP headers"""
    headers = get_http_headers()
    token = headers.get("x-forwarded-access-token", "NONE")
    return {"token": token}

if __name__ == "__main__":
    mcp.run(transport="streamable-http", host="0.0.0.0", port=3000)


# reverse proxy to modify the header:
# run it with `python reverse_proxy.py`
#!/usr/bin/env python3
"""
MCP Reverse Proxy in Python
Similar to the Go OAuth proxy but simplified for header transformation
"""

import os
import requests
from flask import Flask, request, Response, stream_with_context
from werkzeug.datastructures import Headers

app = Flask(__name__)

# Configuration
MCP_SERVER_URL = os.getenv("MCP_SERVER_URL", "http://localhost:3000").rstrip('/')

def transform_headers(original_headers):
    """Transform incoming headers for forwarding to MCP server"""
    import uuid
    
    # Start with completely clean headers
    new_headers = {}
    
    # Explicitly copy only the headers we want, excluding problematic ones
    excluded_headers = {'authorization', 'host', 'content-length', 'transfer-encoding', 'x-forwarded-access-token'}
    
    # First pass: copy all non-excluded headers
    for key, value in original_headers.items():
        if key.lower() not in excluded_headers:
            new_headers[key] = value
    
    # Second pass: apply transformations
    token_found = False
    for key, value in original_headers.items():
        if key.lower() == 'x-forwarded-access-token':
            token_found = True
            new_uuid = str(uuid.uuid4())
            new_headers['X-Forwarded-Access-Token'] = new_uuid
            print(f"TRANSFORMED: {key}: {value} -> X-Forwarded-Access-Token: {new_uuid}")
            break
    
    if not token_found:
        print("WARNING: x-forwarded-access-token header not found in original headers")
    
    print("Final forwarded headers:", dict(new_headers))
    return new_headers

@app.route('/mcp', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
@app.route('/mcp/<path:subpath>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
def proxy_to_mcp(subpath=''):
    """Proxy requests to MCP server with header transformation"""
    
    # Build target URL
    if subpath:
        target_url = f"{MCP_SERVER_URL}/mcp/{subpath}"
    else:
        target_url = f"{MCP_SERVER_URL}/mcp/"
    
    # Add query parameters if any
    if request.query_string:
        target_url += f"?{request.query_string.decode()}"
    
    # Transform headers
    forwarded_headers = transform_headers(request.headers)
    
    # Debug logging
    print(f"Proxying {request.method} {request.url} -> {target_url}")
    print(f"Original headers: {dict(request.headers)}")
    print(f"Forwarded headers: {dict(forwarded_headers)}")
    
    # Additional debug: check the exact headers being sent
    print("EXACT HEADERS BEING SENT TO DOWNSTREAM:")
    for k, v in forwarded_headers.items():
        print(f"  {k}: {v}")
    
    try:
        # Get request data
        request_data = request.get_data()
        print(f"Request data length: {len(request_data) if request_data else 0}")
        
        # Forward the request
        response = requests.request(
            method=request.method,
            url=target_url,
            headers=forwarded_headers,
            data=request_data,
            cookies=request.cookies,
            allow_redirects=False,
            stream=True,
            timeout=30
        )
        
        print(f"Response status: {response.status_code}")
        print(f"Response headers: {dict(response.headers)}")
        
        # Handle response headers - filter out problematic ones
        response_headers = {}
        excluded_response_headers = {'transfer-encoding', 'content-encoding', 'content-length', 'connection'}
        
        for key, value in response.headers.items():
            if key.lower() not in excluded_response_headers:
                response_headers[key] = value
        
        # Handle redirects by modifying the location header
        if response.status_code in [301, 302, 307, 308] and 'location' in response_headers:
            original_location = response_headers['location']
            # Replace the target server URL with our proxy URL
            if original_location.startswith(MCP_SERVER_URL):
                new_location = original_location.replace(MCP_SERVER_URL, 'http://localhost:8080', 1)
                response_headers['location'] = new_location
                print(f"REDIRECT INTERCEPTED: {original_location} -> {new_location}")
        
        # Create response with original status code
        def generate():
            for chunk in response.iter_content(chunk_size=8192):
                yield chunk
        
        flask_response = Response(
            stream_with_context(generate()),
            status=response.status_code,
            headers=response_headers
        )
        
        return flask_response
        
    except requests.exceptions.RequestException as e:
        print(f"Proxy error: {e}")
        return Response(
            f"Proxy error: {str(e)}", 
            status=502,
            headers={'Content-Type': 'text/plain'}
        )

@app.route('/health')
def health_check():
    """Health check endpoint"""
    return {'status': 'ok', 'target': MCP_SERVER_URL}

if __name__ == '__main__':
    print(f"Starting MCP Reverse Proxy on http://localhost:8080")
    print(f"Forwarding to: {MCP_SERVER_URL}")
    app.run(host='0.0.0.0', port=8080, debug=True)


# client -- Send Multiple Requests with Different Tokens:

import asyncio
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport

async def test_multiple_requests():
    async with Client(transport=StreamableHttpTransport(
        "http://127.0.0.1:8080/mcp/",
        headers={"x-forwarded-access-token": "TOKEN_1"},
    )) as client:
        # Request 1
        result1 = await client.call_tool("test_token", {})
        print(f"Request 1 result: {result1.content[0].text}") 

        # the subsequent request within the same session, but the token should be updated by the reverse proxy
        result2 = await client.call_tool("test_token", {})
        print(f"Request 2 result: {result2.content[0].text}") 

asyncio.run(test_multiple_requests())