#!/usr/bin/env python3
"""
MCP client wrapper for Bedrock AgentCore with Cognito authentication.
This script acts as an MCP stdio proxy that forwards requests to the AgentCore runtime.
"""
import sys
import json
import boto3
import httpx
from urllib.parse import quote
from botocore.exceptions import ClientError
import os
import requests


# Client Credentials
CLIENT_ID = os.getenv("CLIENT_ID")
CLIENT_SECRET = os.getenv("CLIENT_SECRET")

# Runtime Variables
MCP_URL = os.getenv("MCP_URL")
TOKEN_URL = os.getenv("TOKEN_URL")
REGION = "ap-southeast-2"

# Initialize Cognito client
cognito_client = boto3.client('cognito-idp', region_name=REGION)


def get_cognito_token():
    """Get an access token from Cognito"""
    
    token_data = {
        "grant_type": "client_credentials",
        "client_id": CLIENT_ID,
        "client_secret": CLIENT_SECRET
    }
    
    headers = {
        "Content-Type": "application/x-www-form-urlencoded"
    }

    try:
        response = requests.post(TOKEN_URL, data=token_data, headers=headers)
        response.raise_for_status()
        access_token = response.json()["access_token"]

        return access_token
        
    except ClientError as e:
        print(json.dumps({"error": f"Cognito auth failed: {str(e)}"}), file=sys.stderr)
        raise


def call_agentcore(method, params=None):
    """Call the AgentCore runtime with proper authentication."""
    access_token = get_cognito_token()
    
    payload = {
        "jsonrpc": "2.0",
        "id": 1,
        "method": method,
        "params": params or {}
    }
    
    
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json, text/event-stream',
        'Authorization': f'Bearer {access_token}'
    }
    
    # Debug: log the request
    print(f"DEBUG: Calling AgentCore with method={method}, params={params}", file=sys.stderr)
    
    response = httpx.post(MCP_URL, json=payload, headers=headers, timeout=30.0)
    response.raise_for_status()
    
    # Parse the response (handle event-stream format)
    raw = response.text
    
    # Debug: log the raw response
    print(f"DEBUG: Raw response: {raw[:500]}", file=sys.stderr)
    
    if '{' in raw:
        json_data = json.loads(raw[raw.find('{'):])
        result = json_data.get('result')
        
        # Debug: log the extracted result
        print(f"DEBUG: Extracted result: {result}", file=sys.stderr)
        
        return result
    
    return None


def handle_request(request):
    """Handle incoming MCP requests and proxy to AgentCore."""
    method = request.get('method')
    params = request.get('params', {})
    request_id = request.get('id')
    
    try:
        result = call_agentcore(method, params)
        
        # For notifications (no id), don't send a response
        if request_id is None:
            return None
            
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": result
        }
    except Exception as e:
        # For notifications (no id), don't send error response
        if request_id is None:
            return None
            
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": -32603,
                "message": str(e)
            }
        }


def main():
    """Main loop for MCP stdio protocol."""
    for line in sys.stdin:
        try:
            line = line.strip()
            if not line:
                continue
                
            request = json.loads(line)
            response = handle_request(request)
            
            # Only print response if it's not None (notifications don't get responses)
            if response is not None:
                print(json.dumps(response), flush=True)
        except json.JSONDecodeError:
            continue
        except Exception as e:
            error_response = {
                "jsonrpc": "2.0",
                "id": None,
                "error": {
                    "code": -32700,
                    "message": f"Parse error: {str(e)}"
                }
            }
            print(json.dumps(error_response), flush=True)


if __name__ == "__main__":
    main()