#!/usr/bin/env python3
"""Step 2: Exchange authorization code for tokens using saved state."""
import sys, os, json, asyncio

sys.path.insert(0, "/usr/local/lib/hermes-agent")

SERVER_NAME = sys.argv[1] if len(sys.argv) > 1 else "cloudflare"
CALLBACK_URL = sys.argv[2] if len(sys.argv) > 2 else ""

if not CALLBACK_URL:
    print("Usage: python3 oauth-step2.py <server_name> <callback_url>")
    sys.exit(1)

# Load state from step 1
state_path = f"/tmp/mcp_oauth_{SERVER_NAME}.json"
if not os.path.exists(state_path):
    print(f"  ✗ State file not found. Run oauth-step1.py first.")
    sys.exit(1)

with open(state_path) as f:
    state_data = json.load(f)

from urllib.parse import urlparse, parse_qs
parsed = urlparse(CALLBACK_URL)
cb_params = parse_qs(parsed.query)
received_code = cb_params.get("code", [None])[0]
received_state = cb_params.get("state", [None])[0]

if not received_code:
    print("  ✗ No 'code' parameter in callback URL")
    sys.exit(1)

if received_state != state_data["state"]:
    print(f"  ✗ State mismatch! Expected {state_data['state']}, got {received_state}")
    sys.exit(1)

print(f"  ✓ Callback verified (state matches)")
print(f"  Code: {received_code[:40]}...")

from tools.mcp_oauth import HermesTokenStorage

storage = HermesTokenStorage(SERVER_NAME)
client_id = state_data["client_id"]
token_endpoint = state_data["token_endpoint"]
redirect_uri = state_data["redirect_uri"]
code_verifier = state_data["code_verifier"]

import httpx

async def exchange():
    async with httpx.AsyncClient(timeout=15.0) as client:
        data = {
            "grant_type": "authorization_code",
            "code": received_code,
            "redirect_uri": redirect_uri,
            "client_id": client_id,
            "code_verifier": code_verifier,
        }
        
        resp = await client.post(token_endpoint, data=data, headers={"Accept": "application/json"})
        
        if resp.status_code != 200:
            print(f"  ✗ Token exchange failed: HTTP {resp.status_code}")
            print(f"  Response: {resp.text[:500]}")
            return False
        
        token_data = resp.json()
        print(f"  ✓ Token received!")
        
        from mcp.shared.auth import OAuthToken, OAuthClientInformationFull
        token = OAuthToken(
            access_token=token_data.get("access_token", ""),
            token_type=token_data.get("token_type", "Bearer"),
            expires_in=token_data.get("expires_in"),
            refresh_token=token_data.get("refresh_token"),
            scope=token_data.get("scope"),
        )
        await storage.set_tokens(token)
        
        client_info = OAuthClientInformationFull(
            client_id=client_id,
            redirect_uris=[redirect_uri],
        )
        await storage.set_client_info(client_info)
        
        print(f"  ✓ Tokens saved to ~/.hermes/mcp-tokens/{SERVER_NAME}.json")
        return True

success = asyncio.run(exchange())

if success:
    print(f"\n  ✓ '{SERVER_NAME}' authenticated successfully!")
    # Clean up state file
    os.remove(state_path)
else:
    print(f"\n  ✗ '{SERVER_NAME}' authentication failed")
    sys.exit(1)