"""
auth.py — Server-side Microsoft Entra ID (Azure AD) JWT validation.

Every protected route uses the @require_creditfair decorator which:
  1. Reads the Bearer token from the Authorization header.
  2. Fetches Microsoft's public JWKS keys (cached in-process).
  3. Verifies the token's RSA signature, expiry, audience and issuer.
  4. Checks that the upn / preferred_username / email claim ends with @creditfair.in.

This is the ONLY real security boundary. The client-side domain check in the
React app is UX only — a user can bypass it with DevTools. This check cannot
be bypassed because the token is cryptographically signed by Microsoft.
"""

import os
import functools
import requests
import jwt                          # PyJWT
from jwt.algorithms import RSAAlgorithm
from flask import request, jsonify, g

# ── Config ────────────────────────────────────────────────────────────────────
ALLOWED_EMAILS  = ['disha.doshi@creditfair.in', 'sunil@creditfair.in', 'raghava@blacksoil.co.in', 'gowri@blacksoil.co.in', 'mirza.karrar@blacksoil.co.in']
CLIENT_ID       = os.environ.get('OAUTH_APP_ID', '')
TENANT_ID       = os.environ.get('OAUTH_TENANT_ID', 'common')

# Microsoft's JWKS endpoint for the tenant
JWKS_URI = f'https://login.microsoftonline.com/{TENANT_ID}/discovery/v2.0/keys'

# Valid token issuers (v1 and v2 endpoints)
VALID_ISSUERS = [
    f'https://login.microsoftonline.com/{TENANT_ID}/v2.0',
    f'https://sts.windows.net/{TENANT_ID}/',
]

# ── JWKS cache (in-process, refreshed once per process lifetime) ──────────────
_jwks_cache = None

def _get_jwks():
    global _jwks_cache
    if _jwks_cache is None:
        resp = requests.get(JWKS_URI, timeout=10)
        resp.raise_for_status()
        _jwks_cache = {key['kid']: key for key in resp.json()['keys']}
    return _jwks_cache


def _get_signing_key(kid):
    """Return the RSA public key object for the given key ID."""
    jwks = _get_jwks()
    key_data = jwks.get(kid)
    if not key_data:
        # Key not in cache — refresh once in case Microsoft rotated keys
        global _jwks_cache
        _jwks_cache = None
        jwks = _get_jwks()
        key_data = jwks.get(kid)
    if not key_data:
        raise ValueError(f'Unknown signing key: {kid}')
    return RSAAlgorithm.from_jwk(key_data)


def _validate_token(token):
    """
    Decode and validate a Microsoft-issued access/id token.
    Returns the decoded claims dict on success, raises on failure.
    """
    # Peek at the header to get the key ID
    header = jwt.get_unverified_header(token)
    kid    = header.get('kid')
    if not kid:
        raise ValueError('Token missing kid header')

    public_key = _get_signing_key(kid)

    # Determine audience: access tokens use client_id, id tokens use client_id too
    audience = CLIENT_ID if CLIENT_ID else None

    decode_options = {'verify_exp': True}
    if not audience:
        # If CLIENT_ID is not configured, skip audience check (dev fallback)
        decode_options['verify_aud'] = False

    claims = jwt.decode(
        token,
        key=public_key,
        algorithms=['RS256'],
        audience=audience,
        options=decode_options,
    )

    # Verify issuer
    issuer = claims.get('iss', '')
    if TENANT_ID != 'common' and issuer not in VALID_ISSUERS:
        raise ValueError(f'Invalid token issuer: {issuer}')

    # Extract email / UPN claim (Microsoft uses several claim names)
    email = (
        claims.get('upn') or
        claims.get('preferred_username') or
        claims.get('email') or
        claims.get('unique_name') or
        ''
    ).lower()

    if email not in ALLOWED_EMAILS:
        raise PermissionError(
            f'Access denied: {email or "(no email claim)"} is not authorized'
        )

    claims['_verified_email'] = email
    return claims


# ── Decorator ─────────────────────────────────────────────────────────────────

def require_creditfair(fn):
    """
    Route decorator that validates the Microsoft JWT on every request.
    Sets g.user_email to the verified email on success.

    Usage:
        @bp.route('/api/policies', methods=['GET'])
        @require_creditfair
        def get_policies():
            ...
    """
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        auth_header = request.headers.get('Authorization', '')
        if not auth_header.startswith('Bearer '):
            return jsonify({'error': 'Unauthorized', 'message': 'Missing Bearer token'}), 401

        token = auth_header[len('Bearer '):]
        try:
            claims = _validate_token(token)
            g.user_email = claims['_verified_email']
        except PermissionError as exc:
            return jsonify({'error': 'Forbidden', 'message': str(exc)}), 403
        except Exception as exc:
            return jsonify({'error': 'Unauthorized', 'message': f'Token validation failed: {exc}'}), 401

        return fn(*args, **kwargs)
    return wrapper