from my_modules.app.setup import cache, LIMITER
from my_modules.app.constens import API_GROUP
from my_modules.app.logger import logger

from quart import Blueprint, redirect, url_for, session, request, render_template, abort, make_response, current_app

from authlib.jose import JsonWebKey, JoseError, jwt as authlib_jwt
from authlib.integrations.httpx_client import AsyncOAuth2Client
from jwt import InvalidTokenError

import httpx, uuid, os

auth_login_bp = Blueprint('auth_login', __name__)

# OAuth
OIDC_CLIENT_ID = os.getenv('OIDC_CLIENT_ID', '')
OIDC_CLIENT_SECRET = os.getenv('OIDC_CLIENT_SECRET', '')
OIDC_METADATA_URL = os.getenv('OIDC_METADATA_URL', '')
REDIRECT_URI_SCHEME = os.getenv('REDIRECT_URI_SCHEME', 'http')

async def get_oidc_metadata():
  async with httpx.AsyncClient() as client:
    response = await client.get(OIDC_METADATA_URL)
    response.raise_for_status()
    return response.json()

@auth_login_bp.route('/login', methods=['GET'])
@auth_login_bp.route('/auth', methods=['GET'])
@LIMITER.limit("5 per minute; 30 per hour")
async def login():
  metadata = await get_oidc_metadata()

  auth_id = request.cookies.get('auth_id') or str(uuid.uuid4())
  nonce = str(uuid.uuid4())

  await cache.set(f"oauth:nonce:{auth_id}", nonce, ttl=3603)

  client = AsyncOAuth2Client(
    client_id=OIDC_CLIENT_ID,
    client_secret=OIDC_CLIENT_SECRET,
    redirect_uri=url_for('auth_login.auth_callback', _external=True, _scheme=REDIRECT_URI_SCHEME),
    scope='openid profile email',
  )

  uri, state = client.create_authorization_url(
    metadata['authorization_endpoint'],
    nonce=nonce,
    state=auth_id,
  )

  response = await make_response(redirect(uri))
  response.set_cookie('auth_id', auth_id, max_age=3600, httponly=True, secure=True, samesite='Lax')
  return response

@auth_login_bp.route('/logout', methods=['GET'])
@auth_login_bp.route('/auth/logout', methods=['GET'])
@LIMITER.limit("10 per minute")
async def logout():
  user = session.get('user')
  if user:
    await logger.info(f'logging out user: {user}')
    session.pop('user', None)

  session.clear()
  return redirect(url_for('side_main.index'))

@auth_login_bp.route('/auth/callback', methods=['GET'])
@LIMITER.limit("5 per minute")
async def auth_callback():
  try:
    auth_id = request.args.get('state')
    code = request.args.get('code')
    nonce = await cache.get(f"oauth:nonce:{auth_id}")

    if not nonce:
      await logger.error('Nonce not found')
      return await render_template('views/api/token.htm', error='Nonce not found'), 400

    await cache.delete(f"oauth:nonce:{auth_id}")

    metadata = await get_oidc_metadata()
    client = AsyncOAuth2Client(
      client_id=OIDC_CLIENT_ID,
      client_secret=OIDC_CLIENT_SECRET,
      redirect_uri=url_for('auth_login.auth_callback', _external=True, _scheme=REDIRECT_URI_SCHEME),
      scope='openid profile email',
    )

    # Exchange code for token
    token = await client.fetch_token(
      metadata['token_endpoint'],
      code=code,
      grant_type='authorization_code'
    )
    await logger.debug(f'Auth Callback | token: {token}')

    # Decode ID token
    id_token = token.get('id_token')
    if not id_token:
      await logger.error('ID token missing in OAuth response')
      return await render_template('views/api/token.htm', error='ID token missing'), 400

    # Fetch the JWKs to verify signature
    async with httpx.AsyncClient() as http_client:
      jwks_resp = await http_client.get(metadata['jwks_uri'])
      jwks = jwks_resp.json()

    keys = JsonWebKey.import_key_set(jwks)
    claims = authlib_jwt.decode(id_token, key=keys)

    try:
      claims.validate_aud()
      claims.validate()
    except JoseError as e:
      await logger.error(f"JWT validation error: {e}")
      return await render_template('views/api/token.htm', error=str(e)), 400

    if claims.get('nonce') != nonce:
      await logger.error('Nonce mismatch in ID token')
      return await render_template('views/api/token.htm', error='Invalid nonce'), 400

    await logger.info(f'Auth Callback | user_info: {claims}')

    if API_GROUP not in claims.get('groups', []):
      await logger.error("You don't have Permissions to Access this API")
      return await render_template('views/api/token.htm', error="You don't have Permissions to Access this API"), 403

    session['user'] = claims
    response = await make_response(redirect(url_for('side_main.index')))
    response.set_cookie('auth_id', '', max_age=0, httponly=True, secure=True, samesite='Lax')
    return response

  except httpx.HTTPError as e:
    await logger.error(f"HTTP error during token exchange: {e}")
    return await render_template('views/api/token.htm', error="Token exchange failed"), 500
  except InvalidTokenError as e:
    await logger.error(f"Invalid ID Token: {e}")
    return await render_template('views/api/token.htm', error="Invalid ID Token"), 400
  except Exception as e:
    await logger.exception("Unknown error during OAuth callback")
    return await render_template('views/api/token.htm', error=str(e)), 500