#!/usr/bin/env python3
"""
OAuth login for Google Workspace MCP servers (Sheets, Drive) in non-TTY.

Usage:
  1. systemctl --user stop hermes-gateway.service
  2. python3 /root/workspace/google-workspace-mcp-login.py
  3. For each server, open the URL in your browser, authorize, paste callback URL
  4. systemctl --user start hermes-gateway.service
"""
import sys
import os
import json
import hashlib
import base64
import secrets
import string
import asyncio
from urllib.parse import urlencode, urlparse, parse_qs

sys.path.insert(0, "/usr/local/lib/hermes-agent")
sys.path.insert(0, "/usr/local/lib/hermes-agent/venv/lib/python3.11/site-packages")

from hermes_cli.config import load_config
from tools.mcp_oauth import HermesTokenStorage, _configure_callback_port
from tools.mcp_oauth_manager import get_manager

import httpx

# ── Google OAuth settings ──────────────────────────────────────────────
CLIENT_ID = "1090672727072-4buk3522j52pqut02m22rcnto31938jl.apps.googleusercontent.com"
CLIENT_SECRET = "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"
AUTH_URI = "https://accounts.google.com/o/oauth2/auth"
TOKEN_URI = "https://oauth2.googleapis.com/token"

# Servers to configure
SERVERS = [
    {
        "name": "google-sheets",
        "scopes": [
            "https://www.googleapis.com/auth/spreadsheets",
            "https://www.googleapis.com/auth/drive.file",
        ],
    },
    {
        "name": "google-drive",
        "scopes": [
            "https://www.googleapis.com/auth/drive.readonly",
            "https://www.googleapis.com/auth/drive.file",
        ],
    },
]
# ────────────────────────────────────────────────────────────────────────


async def exchange_code(
    server_name: str,
    code: str,
    code_verifier: str,
    redirect_uri: str,
) -> bool:
    """Exchange authorization code for tokens and save them."""
    from mcp.shared.auth import OAuthToken, OAuthClientInformationFull

    storage = HermesTokenStorage(server_name)

    async with httpx.AsyncClient(timeout=15.0) as client:
        data = {
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": redirect_uri,
            "client_id": CLIENT_ID,
            "client_secret": CLIENT_SECRET,
            "code_verifier": code_verifier,
        }

        resp = await client.post(TOKEN_URI, 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!")

        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)

        # Write complete client info
        client_info = OAuthClientInformationFull(
            client_id=CLIENT_ID,
            client_secret=CLIENT_SECRET,
            redirect_uris=[redirect_uri],
            grant_types=["authorization_code", "refresh_token"],
            response_types=["code"],
            token_endpoint_auth_method="client_secret_post",
        )
        await storage.set_client_info(client_info)

        print(f"  ✓ Tokens saved to ~/.hermes/mcp-tokens/{server_name}.json")
        return True


def main():
    for server in SERVERS:
        name = server["name"]
        scopes = server["scopes"]

        print(f"\n{'=' * 70}")
        print(f"  Server: {name}")
        print(f"  Scopes: {' '.join(scopes)}")
        print(f"{'=' * 70}")

        # Clear old tokens
        storage = HermesTokenStorage(name)
        for p in [storage._tokens_path(), storage._client_info_path(), storage._meta_path()]:
            if p.exists():
                p.unlink()
        get_manager().remove(name)

        # Configure callback port
        cfg = {}
        _configure_callback_port(cfg, storage)
        resolved_port = cfg.get("_resolved_port", 38255)
        redirect_uri = f"http://127.0.0.1:{resolved_port}/callback"

        # Generate PKCE
        code_verifier = "".join(
            secrets.choice(string.ascii_letters + string.digits + "-._~")
            for _ in range(128)
        )
        digest = hashlib.sha256(code_verifier.encode()).digest()
        code_challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=")
        state = secrets.token_urlsafe(32)

        # Build auth URL
        params = {
            "response_type": "code",
            "client_id": CLIENT_ID,
            "redirect_uri": redirect_uri,
            "scope": " ".join(scopes),
            "state": state,
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
            "access_type": "offline",
            "prompt": "consent",
        }
        auth_url = f"{AUTH_URI}?{urlencode(params)}"

        print(f"\n  OPEN THIS URL IN YOUR BROWSER:\n")
        print(f"  {auth_url}")
        print(f"\n  After authorizing, paste the FULL callback URL and press Enter:")
        print(f"  (or type 'skip' to skip this server)")

        # Read callback URL from stdin
        try:
            callback_url = sys.stdin.readline().strip()
        except (EOFError, KeyboardInterrupt):
            print("\n  Skipping.")
            continue

        if not callback_url or callback_url.lower() in ("skip", "s", "q", "quit"):
            print("  Skipped.")
            continue

        # Parse callback
        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")
            continue

        if received_state != state:
            print(f"  ✗ State mismatch! Expected {state}, got {received_state}")
            continue

        print(f"  ✓ Callback verified")

        # Exchange code
        success = asyncio.run(
            exchange_code(name, received_code, code_verifier, redirect_uri)
        )

        if success:
            print(f"  ✓ '{name}' authenticated successfully!")
        else:
            print(f"  ✗ '{name}' authentication failed")

    print(f"\n{'=' * 70}")
    print(f"  Done. Restart the gateway:")
    print(f"    systemctl --user restart hermes-gateway.service")
    print(f"{'=' * 70}")


if __name__ == "__main__":
    main()