#!/usr/bin/env python3
"""Native MCP OAuth example: discovery, consent, rotation, persisted restart.

Run in its own venv (requirements.txt). Tokens never enter stdout or arguments.
One process owns the private credential store for the lifetime of a connection.
"""
from __future__ import annotations

import argparse
import asyncio
from contextlib import contextmanager
import fcntl
import hashlib
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
import os
from pathlib import Path
import queue
import secrets
import tempfile
import threading
import time
from urllib.parse import parse_qs, urlsplit
import webbrowser

import httpx2
from pydantic import AnyUrl
from mcp import Client
from mcp.client.auth import AuthorizationCodeResult, OAuthClientProvider
from mcp.client.streamable_http import streamable_http_client
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken


class PrivateStorage:
    def __init__(self, path: Path, server_url: str):
        self.path, self.server_url = path, server_url
        self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)

    def read(self):
        try:
            fd = os.open(self.path, os.O_RDONLY | os.O_NOFOLLOW)
        except FileNotFoundError:
            return {'server_url': self.server_url}
        with os.fdopen(fd) as file:
            st = os.fstat(file.fileno())
            if st.st_uid != os.getuid() or st.st_mode & 0o077:
                raise RuntimeError('Credential file must belong to you and have permission 0600')
            data = json.load(file)
        if data.get('server_url') != self.server_url:
            raise RuntimeError('Credential store belongs to a different MCP endpoint')
        return data

    def write(self, field, value):
        data = self.read()
        data[field] = value
        self._replace(data)

    async def get_tokens(self):
        data = self.read()
        token = data.get('tokens')
        if not token:
            return None
        # expires_in is relative. Persist its absolute deadline so reopening the
        # process cannot accidentally give an expired access token a new hour.
        token['expires_in'] = max(0, int(data['expires_at'] - time.time()))
        return OAuthToken.model_validate(token)

    async def set_tokens(self, tokens):
        data = self.read()
        data.update(tokens=tokens.model_dump(mode='json'), expires_at=time.time() + (tokens.expires_in or 0))
        # Write tokens and deadline as one atomic record, never in two updates.
        self._replace(data)

    def _replace(self, data):
        fd, temp = tempfile.mkstemp(prefix='.oauth-', dir=self.path.parent)
        try:
            with os.fdopen(fd, 'w') as file:
                os.fchmod(file.fileno(), 0o600)
                json.dump(data, file)
                file.flush()
                os.fsync(file.fileno())
            os.replace(temp, self.path)
        finally:
            if os.path.exists(temp):
                os.unlink(temp)

    async def get_client_info(self):
        data = self.read().get('client_info')
        return OAuthClientInformationFull.model_validate(data) if data else None

    async def set_client_info(self, client_info):
        self.write('client_info', client_info.model_dump(mode='json'))

    @contextmanager
    def exclusive(self):
        fd = os.open(str(self.path) + '.lock', os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600)
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
            yield
        except BlockingIOError:
            raise RuntimeError('Another client is using this credential store. Use a separate store for a separate agent.') from None
        finally:
            os.close(fd)


class BrowserCallback:
    def __init__(self, port: int):
        self.result = queue.Queue(maxsize=1)
        self.expected_state = None
        parent = self
        class Handler(BaseHTTPRequestHandler):
            def log_message(self, *args):
                pass  # callback query includes the one-time authorization code

            def do_GET(self):
                params = parse_qs(urlsplit(self.path).query)
                valid = urlsplit(self.path).path == '/callback' and parent.expected_state and secrets.compare_digest(params.get('state', [''])[0], parent.expected_state)
                if not valid or any(len(v) != 1 for v in params.values()):
                    self.send_response(400); self.end_headers(); return
                self.send_response(200)
                self.send_header('Content-Type', 'text/plain; charset=utf-8')
                self.send_header('Cache-Control', 'no-store')
                self.send_header('Referrer-Policy', 'no-referrer')
                self.end_headers()
                self.wfile.write(b'Connection reviewed. You can return to your terminal.')
                try:
                    parent.result.put_nowait(params)
                except queue.Full:
                    pass
        self.server = ThreadingHTTPServer(('127.0.0.1', port), Handler)
        self.uri = f'http://127.0.0.1:{self.server.server_port}/callback'
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
        self.thread.start()

    async def open(self, url):
        self.expected_state = parse_qs(urlsplit(url).query)['state'][0]
        print('Review the connection in your browser: ' + url)
        await asyncio.to_thread(webbrowser.open, url)

    async def wait(self):
        try:
            params = await asyncio.to_thread(self.result.get, True, 600)
        except queue.Empty:
            raise RuntimeError('Connection review timed out; start again.') from None
        if 'error' in params:
            raise RuntimeError('The connection was declined.')
        return AuthorizationCodeResult(code=params['code'][0], state=params['state'][0], iss=params.get('iss', [None])[0])

    def close(self):
        self.server.shutdown(); self.server.server_close(); self.thread.join(timeout=2)


def provider(server, storage, callback, client_metadata_url=None):
    return OAuthClientProvider(server_url=server, client_metadata=OAuthClientMetadata(
        client_name='CiteDoor Python example', redirect_uris=[AnyUrl(callback.uri)],
        token_endpoint_auth_method='none', scope='mcp'), storage=storage,
        redirect_handler=callback.open, callback_handler=callback.wait, client_metadata_url=client_metadata_url)


async def connect(server, storage, callback, client_metadata_url=None):
    oauth = provider(server, storage, callback, client_metadata_url)
    async with httpx2.AsyncClient(auth=oauth, timeout=30, trust_env=False) as http:
        async with Client(streamable_http_client(server, http_client=http)) as client:
            result = await client.list_tools()
            print('Connected. Available tools: ' + ', '.join(tool.name for tool in result.tools))
            # OAuthClientProvider renews before later requests automatically.
            # Call tools here; a static Authorization header would bypass renewal.
            return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--server', default=os.environ.get('HIVEHALL_PUBLIC_URL', os.environ.get('CITEDOOR_PUBLIC_URL', '')).rstrip('/') + '/mcp'
                        if os.environ.get('HIVEHALL_PUBLIC_URL', os.environ.get('CITEDOOR_PUBLIC_URL')) else None,
                        help='MCP URL (or set HIVEHALL_PUBLIC_URL)')
    parser.add_argument('--callback-port', type=int, default=8858)
    parser.add_argument('--store', type=Path)
    parser.add_argument('--client-metadata-url', help='Optional public HTTPS client metadata document (CIMD)')
    args = parser.parse_args()
    if not args.server:
        parser.error('provide --server or set HIVEHALL_PUBLIC_URL')
    u = urlsplit(args.server)
    if u.scheme != 'https' and not (u.scheme == 'http' and u.hostname in ('127.0.0.1', '::1')):
        parser.error('Use HTTPS, or an explicit loopback IP for local development.')
    suffix = hashlib.sha256(args.server.encode()).hexdigest()[:24]
    path = args.store or Path.home() / '.local' / 'share' / 'citedoor' / ('oauth-' + suffix + '.json')
    storage = PrivateStorage(path, args.server)
    with storage.exclusive():
        callback = BrowserCallback(args.callback_port)
        try:
            asyncio.run(connect(args.server, storage, callback, args.client_metadata_url))
        finally:
            callback.close()


if __name__ == '__main__':
    try:
        main()
    except Exception as error:
        # OAuth/transport exceptions may contain token responses or callback codes.
        raise SystemExit('Connection failed (' + type(error).__name__ + '). Review permissions and retry.') from None