"""Automatic HF session reuse through a state-checked, silent OIDC flow."""
from __future__ import annotations

import base64
from dataclasses import dataclass
import hmac
import json
import os
import re
import secrets
import time
from urllib.parse import urlsplit

from fastapi import FastAPI, HTTPException, Request
from itsdangerous import BadSignature, TimestampSigner, URLSafeTimedSerializer
from starlette.middleware.sessions import SessionMiddleware
from starlette.requests import HTTPConnection
from starlette.responses import HTMLResponse, JSONResponse

from hf_space.history_auth import identity_from_oauth, optional_identity
from hf_space.history_bridge import (AppSessionMiddleware, AppSessions, BootstrapGuards, Bridges,
                                    bridge_body, popup_bootstrap, private_headers)


AUTH_PATH = '/api/history/auth'
FLOW_KEY = '_supra_history_oauth_flow'
ATTEMPT_KEY = '_supra_history_oauth_silent_at'
FLOW_LIFETIME = 600
CHECK_COOKIE = '__Host-supra-history-oauth-check'


@dataclass(frozen=True)
class OAuthSettings:
    client_id: str
    client_secret: str
    origin: str
    provider: str = 'https://huggingface.co'
    scopes: str = 'openid profile'

    @classmethod
    def from_environment(cls):
        # Do not inherit Gradio's local mocked OAuth/server-token behavior.
        if not os.environ.get('SPACE_ID'):
            return None
        client_id = os.environ.get('OAUTH_CLIENT_ID')
        client_secret = os.environ.get('OAUTH_CLIENT_SECRET')
        host = os.environ.get('SPACE_HOST', '').split(',')[0].strip()
        provider = os.environ.get('OPENID_PROVIDER_URL', '').rstrip('/')
        if not client_id or not client_secret or not host or provider != 'https://huggingface.co':
            return None
        parsed = urlsplit('https://' + host)
        if (parsed.netloc != host or not parsed.hostname or parsed.path
                or parsed.username is not None or parsed.password is not None):
            return None
        # Personal history only needs these two baseline profile scopes.
        return cls(client_id, client_secret, 'https://' + host)


class SessionWriteGuard:
    """Keep unchanged requests from overwriting concurrent OAuth cookie updates.

    Uses the existing middleware's signer, name, and expiry rules. It does not
    create or authenticate another session format.
    """

    def __init__(self, app, *, secret_key, session_cookie='session', max_age=14 * 24 * 60 * 60):
        self.app = app
        self.signer = TimestampSigner(str(secret_key))
        self.session_cookie = session_cookie
        self.max_age = max_age

    async def __call__(self, scope, receive, send):
        if scope['type'] != 'http':
            await self.app(scope, receive, send)
            return
        initial = {}
        cookie = HTTPConnection(scope).cookies.get(self.session_cookie)
        valid_initial = cookie is None
        if cookie is not None:
            try:
                data = self.signer.unsign(cookie.encode('utf-8'), max_age=self.max_age)
                value = json.loads(base64.b64decode(data))
                if isinstance(value, dict):
                    initial = value
                    valid_initial = True
            except (BadSignature, ValueError, TypeError, UnicodeError):
                pass

        async def guarded_send(message):
            if message['type'] == 'http.response.start' and valid_initial:
                current = scope.get('session')
                if isinstance(current, dict) and current == initial:
                    cookie_prefix = (self.session_cookie + '=').encode('latin-1')
                    message = dict(message)
                    message['headers'] = [
                        (name, value) for name, value in message.get('headers', [])
                        if not (name.lower() == b'set-cookie' and value.lstrip().startswith(cookie_prefix))
                    ]
            await send(message)

        await self.app(scope, receive, guarded_send)


def install_session_write_guard(app: FastAPI):
    if getattr(app.state, 'supra_session_write_guard', False):
        return
    for middleware in app.user_middleware:
        if middleware.cls is SessionMiddleware:
            settings = middleware.kwargs
            app.add_middleware(SessionWriteGuard, secret_key=settings['secret_key'],
                               session_cookie=settings.get('session_cookie', 'session'),
                               max_age=settings.get('max_age', 14 * 24 * 60 * 60))
            app.state.supra_session_write_guard = True
            return
    raise RuntimeError('Automatic account connection requires the existing OAuth session middleware.')


def _completion(settings: OAuthSettings, request: Request, error: str | None = None, *, connected=None):
    identity = optional_identity(request)
    authenticated = identity is not None if connected is None else connected
    message = {'type': 'supra-history-auth', 'authenticated': authenticated}
    if error and not authenticated:
        message['error'] = error if error in ('login_required', 'consent_required', 'access_denied') else 'unavailable'
    nonce = secrets.token_urlsafe(24)
    # The payload never contains tokens, claims, a username, or caller URLs.
    payload = json.dumps(message, separators=(',', ':'))
    origin = json.dumps(settings.origin)
    html = f'''<!doctype html><html><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1"><title>Supra-Tex</title>
<style nonce="{nonce}">body{{background:#101010;color:#ddd;font:16px Arial,sans-serif;padding:32px}}a{{color:#f2cd70}}</style></head>
<body><p>{'Account connected. You can close this window and return to Supra-Tex.' if authenticated else 'You can close this window and return to Supra-Tex.'}</p>
<script nonce="{nonce}">(() => {{const message={payload};const origin={origin};if(window.parent!==window)window.parent.postMessage(message,origin);if(window.opener)window.opener.postMessage(message,origin);window.close();}})();</script></body></html>'''
    response = HTMLResponse(html, headers={
        'Cache-Control': 'private, no-store', 'Pragma': 'no-cache',
        'Content-Security-Policy': f"default-src 'none'; script-src 'nonce-{nonce}'; style-src 'nonce-{nonce}'; base-uri 'none'; form-action 'none'; frame-ancestors 'self' https://huggingface.co",
        'Referrer-Policy': 'no-referrer', 'X-Content-Type-Options': 'nosniff',
    })
    if request.scope.get('supra_history_clear_check_state'):
        response.headers.append('set-cookie', f'{CHECK_COOKIE}=; Path=/; Max-Age=0; HttpOnly; Secure; SameSite=None; Partitioned')
    return response


def install_oauth_routes(app: FastAPI, *, settings: OAuthSettings | None = None, oauth_client=None,
                         bridges=None, sessions=None):
    previous = getattr(app.state, 'supra_history_oauth_endpoints', None)
    if previous is not None:
        return previous
    settings = settings or OAuthSettings.from_environment()
    if settings is None:
        return None
    from authlib.integrations.base_client import OAuthError
    if oauth_client is None:
        from authlib.integrations.starlette_client import OAuth
        oauth = OAuth()
        oauth.register('supra_history', client_id=settings.client_id, client_secret=settings.client_secret,
                       client_kwargs={'scope': settings.scopes, 'code_challenge_method': 'S256', 'timeout': 20},
                       server_metadata_url=settings.provider + '/.well-known/openid-configuration')
        oauth_client = oauth.supra_history
    install_session_write_guard(app)
    bridges = bridges if bridges is not None else Bridges()
    sessions = sessions if sessions is not None else AppSessions(settings.client_secret)
    check_signer = URLSafeTimedSerializer(settings.client_secret, salt='supra-history-oauth-check-v1')
    guards = BootstrapGuards(settings.client_secret)
    app.add_middleware(AppSessionMiddleware, sessions=sessions)
    app.state.supra_history_bridges = bridges
    app.state.supra_history_app_sessions = sessions

    async def authorize(request: Request, silent=False, bridge_id=None, refresh=False, check=None):
        # Older Gradio sessions include a full ID token/profile. Once invalid,
        # retaining them alongside PKCE state can overflow the browser's 4 KiB
        # cookie limit and lose the new authorization state.
        request.session.pop('oauth_info', None)
        now = time.time()
        if silent:
            try:
                previous_attempt = float(request.session.get(ATTEMPT_KEY, 0))
            except (ValueError, TypeError):
                previous_attempt = now
            if now - previous_attempt < (2 if refresh else 30):
                return _completion(settings, request, 'unavailable', connected=False if refresh else None)
            request.session[ATTEMPT_KEY] = now
        state = secrets.token_urlsafe(32)
        nonce = secrets.token_urlsafe(32)
        for name in list(request.session):
            if name.startswith('_state_supra_history_'):
                request.session.pop(name, None)
        request.session[FLOW_KEY] = {'state': state, 'created_at': now, 'silent': silent,
                                     'refresh': refresh,
                                     **({'check': check} if check else {}),
                                     **({'bridge_id': bridge_id} if bridge_id else {})}
        current_app = request.scope.get('supra_history_app_identity')
        if refresh and current_app is not None:
            # A delayed failure from this check must not clear a newer login
            # that another tab completed while the provider was responding.
            request.session[FLOW_KEY]['app_sid'] = current_app.session_key
            request.session[FLOW_KEY]['app_check'] = (request.scope.get('supra_history_app_claims') or {}).get('check')
        params = {'state': state, 'nonce': nonce}
        if silent:
            params['prompt'] = 'none'
        try:
            response = await oauth_client.authorize_redirect(request, settings.origin + AUTH_PATH + '/callback', **params)
            response.headers['Cache-Control'] = 'private, no-store'
            response.headers['Referrer-Policy'] = 'no-referrer'
            return response
        except Exception:
            # Network/provider details must not leak into browser messages.
            request.session.pop(FLOW_KEY, None)
            if bridge_id:
                bridges.finish(bridge_id)
            return _completion(settings, request, 'unavailable', connected=False if refresh or bridge_id else None)

    @app.get(AUTH_PATH + '/start')
    async def start(request: Request, silent: bool = False, refresh: bool = False, check: str | None = None):
        if not refresh and optional_identity(request) is not None:
            return _completion(settings, request)
        if refresh:
            if not check or not re.fullmatch(r'[A-Za-z0-9_-]{32,128}', check):
                raise HTTPException(400, 'Invalid account verification request.')
            # Keep the verified app session during the provider round trip so
            # other tabs do not lose it. This cookie is not fresh HF proof:
            # the frontend waits for the new check key before showing data or
            # dispatching a new GPU request.
            request.session.pop('oauth_info', None)
        response = await authorize(request, silent, refresh=refresh, check=check if refresh else None)
        if refresh:
            if response.status_code == 302:
                # Retain Authlib's existing PKCE/nonce state in the same browser
                # partition as the app. Silent checks must also work when an
                # unpartitioned Gradio session cookie is blocked by the browser.
                state = request.session[FLOW_KEY]['state']
                saved = {key: value for key, value in request.session.items()
                         if key in (FLOW_KEY, ATTEMPT_KEY, '_state_supra_history_' + state)}
                value = check_signer.dumps(saved)
                if len(value) > 3500:
                    raise HTTPException(503, 'Account verification could not start. Please sign in again.')
                response.headers.append('set-cookie', f'{CHECK_COOKIE}={value}; Path=/; Max-Age={FLOW_LIFETIME}; HttpOnly; Secure; SameSite=None; Partitioned')
        return response

    @app.get(AUTH_PATH + '/popup')
    async def popup(request: Request):
        response = popup_bootstrap(settings.origin, request)
        if response.status_code == 200:
            guards.issue(response)
        return response

    @app.post(AUTH_PATH + '/bridge')
    async def create_bridge(request: Request):
        value = await bridge_body(request, None, settings.origin)
        result = bridges.create(value.get('bridge_id'), value.get('start_secret'))
        result.update(popup_url=AUTH_PATH + '/popup', start_url=AUTH_PATH + '/bridge/start',
                      poll_url=AUTH_PATH + '/bridge/poll')
        return JSONResponse(result, headers=private_headers())

    @app.post(AUTH_PATH + '/bridge/start')
    async def bridge_start(request: Request):
        value = await bridge_body(request, ('bridge_id', 'start_secret'), settings.origin)
        guards.require(request)
        bridges.start(value['bridge_id'], value['start_secret'])
        guards.require(request, consume=True)
        # Always identify the current HF account. An older application cookie
        # must never replace the account returned by this new OAuth operation.
        response = await authorize(request, bridge_id=value['bridge_id'])
        location = response.headers.get('location', '')
        destination = urlsplit(location)
        provider = urlsplit(settings.provider)
        if (response.status_code != 302 or destination.scheme != 'https'
                or destination.netloc != provider.netloc):
            bridges.finish(value['bridge_id'])
            raise HTTPException(503, 'Sign-in could not start. Please try again shortly.')
        result = JSONResponse({'authorize_url': location}, headers=private_headers())
        guards.clear(result)
        return result

    @app.post(AUTH_PATH + '/bridge/poll')
    async def bridge_poll(request: Request):
        value = await bridge_body(request, ('bridge_id', 'poll_secret'), settings.origin)
        status, claims = bridges.poll(value['bridge_id'], value['poll_secret'])
        response = JSONResponse({'status': status, 'authenticated': status == 'complete',
                                 'started': bridges.started(value['bridge_id'])}, headers=private_headers())
        if status == 'complete':
            sessions.set_cookie(response, claims)
        return response

    @app.get(AUTH_PATH + '/callback')
    async def callback(request: Request):
        state = request.query_params.get('state', '')
        flow = request.session.get(FLOW_KEY)
        if not isinstance(flow, dict) or flow.get('state') != state:
            try:
                saved = check_signer.loads(request.cookies.get(CHECK_COOKIE, ''), max_age=FLOW_LIFETIME)
                candidate = saved.get(FLOW_KEY) if isinstance(saved, dict) else None
                if (isinstance(candidate, dict) and candidate.get('refresh') is True
                        and isinstance(candidate.get('state'), str) and len(state) < 256
                        and hmac.compare_digest(state, candidate['state'])):
                    request.session.update(saved)
                    flow = candidate
            except (BadSignature, ValueError, TypeError):
                pass
        valid = isinstance(flow, dict) and isinstance(flow.get('state'), str)
        if valid:
            try:
                age = time.time() - float(flow['created_at'])
                valid = 0 <= age < FLOW_LIFETIME and len(state) < 256 and hmac.compare_digest(state, flow['state'])
            except (KeyError, ValueError, TypeError):
                valid = False
        # Authlib raises provider error responses before checking state. Check
        # here first, including denied/failed callbacks, and preserve newer flows.
        if not valid:
            return _completion(settings, request, 'unavailable')
        request.scope['supra_history_clear_check_state'] = bool(flow.get('refresh'))
        bridge_id = flow.get('bridge_id')
        fresh = bool(bridge_id or flow.get('refresh'))
        try:
            token = await oauth_client.authorize_access_token(request)
            if not isinstance(token, dict) or not token.get('id_token') or not isinstance(token.get('userinfo'), dict):
                return _completion(settings, request, 'unavailable', connected=False if fresh else None)
            profile = token['userinfo']
            verified = {key: profile[key] for key in ('sub', 'preferred_username', 'name', 'picture', 'profile') if key in profile}
            candidate = {'access_token': token.get('access_token'), 'expires_at': token.get('expires_at'),
                         'scope': token.get('scope', 'openid profile'), 'token_type': token.get('token_type', 'Bearer'),
                         'userinfo': verified}
            # Validate the same identity format as all private API endpoints
            # before replacing a previous, still valid account session.
            old = request.session.get('oauth_info')
            request.session['oauth_info'] = candidate
            identity = identity_from_oauth(candidate)
            if identity is None:
                if old is None:
                    request.session.pop('oauth_info', None)
                else:
                    request.session['oauth_info'] = old
                return _completion(settings, request, 'unavailable', connected=False if fresh else None)
            if bridge_id:
                try:
                    expires = min(float(candidate['expires_at']), float(profile['exp']))
                except (KeyError, ValueError, TypeError, OverflowError):
                    bridges.finish(bridge_id)
                    return _completion(settings, request, 'unavailable', connected=False)
                claims = sessions.claims(identity, expires)
                claims['check'] = bridge_id
                if sessions.validate(claims) is None or not bridges.finish(bridge_id, claims):
                    return _completion(settings, request, 'unavailable', connected=False)
            response = _completion(settings, request, connected=True)
            if not bridge_id:
                # Silent OAuth runs in the original page's embedding partition.
                # Replace its application cookie as well as Gradio's session;
                # otherwise the old app cookie would shadow the new HF account.
                claims = sessions.claims(identity, candidate['expires_at'],
                                         previous=request.scope.get('supra_history_app_claims') if flow.get('refresh') else None)
                if flow.get('check'):
                    claims['check'] = flow['check']
                sessions.set_cookie(response, claims)
            return response
        except OAuthError as exc:
            error = getattr(exc, 'error', 'unavailable')
            response = _completion(settings, request, error, connected=False if fresh else None)
            if flow.get('refresh') and error in ('login_required', 'consent_required', 'access_denied'):
                current_app = request.scope.get('supra_history_app_identity')
                current_check = (request.scope.get('supra_history_app_claims') or {}).get('check')
                if current_app is None or (current_app.session_key == flow.get('app_sid')
                                           and current_check == flow.get('app_check')):
                    request.session.pop('oauth_info', None)
                    sessions.clear_cookie(response)
            return response
        except Exception:
            return _completion(settings, request, 'unavailable', connected=False if fresh else None)
        finally:
            # Errors may leave Authlib's temporary state behind; bound its cookie
            # size and prevent a replay without affecting an older OAuth login.
            await oauth_client.framework.clear_state_data(request.session, state)
            request.session.pop(FLOW_KEY, None)
            if bridge_id:
                bridges.finish(bridge_id)

    endpoints = {'login_url': AUTH_PATH + '/start', 'silent_login_url': AUTH_PATH + '/start?silent=1',
                 'bridge_url': AUTH_PATH + '/bridge', 'popup_url': AUTH_PATH + '/popup'}
    app.state.supra_history_oauth_endpoints = endpoints
    return endpoints