# #region AuthJwtModule [C:5] [TYPE Module] [SEMANTICS auth, validate, jwt, token]
#
# @BRIEF JWT token generation and validation logic with server-side revocation.
# @LAYER Core
# @RELATION DEPENDS_ON -> [auth_config]
# @RELATION DEPENDS_ON -> [TokenBlacklist]
#
# @INVARIANT  Tokens must include aud, iss, iat, jti, exp, and sub claims.
#            Revoked tokens are rejected even if not expired.
# @PRE  JWT secret configured in environment
# @POST  Token encode/decode/revoke functions exported
# @SIDE_EFFECT  None
# @DATA_CONTRACT  TokenPayload -> JWT string

from datetime import UTC, datetime, timedelta
import hashlib
import uuid

from jose import jwt
from sqlalchemy.orm import Session

from ...models.auth import TokenBlacklist
from ..logger import belief_scope
from .config import auth_config


# #region create_access_token [TYPE Function]
# @BRIEF Generates a new JWT access token with aud, iss, iat, jti, exp claims.
# @PRE      data dict contains 'sub' (user_id) and optional 'scopes' (roles).
# @POST     Returns a signed JWT string with unique jti for revocation.
# @RELATION DEPENDS_ON -> [auth_config]
#
def create_access_token(data: dict, expires_delta: timedelta | None = None) -> str:
    with belief_scope("create_access_token"):
        to_encode = data.copy()
        now = datetime.now(UTC)
        if expires_delta:
            expire = now + expires_delta
        else:
            expire = now + timedelta(
                minutes=auth_config.ACCESS_TOKEN_EXPIRE_MINUTES
            )

        to_encode.update({
            "exp": expire,
            "iat": now,
            "jti": str(uuid.uuid4()),
            "aud": auth_config.JWT_AUDIENCE,
            "iss": auth_config.JWT_ISSUER,
        })
        encoded_jwt = jwt.encode(
            to_encode, auth_config.SECRET_KEY, algorithm=auth_config.ALGORITHM
        )
        return encoded_jwt


# #endregion create_access_token


# #region decode_token [TYPE Function]
# @BRIEF Decodes and validates a JWT token — verifies aud, iss, exp, and blacklist.
# @PRE      token is a signed JWT string.
# @POST     Returns the decoded payload if valid. Raises JWTError if invalid.
# @RELATION DEPENDS_ON -> [auth_config]
#
def decode_token(token: str) -> dict:
    with belief_scope("decode_token"):
        payload = jwt.decode(
            token,
            auth_config.SECRET_KEY,
            algorithms=[auth_config.ALGORITHM],
            audience=auth_config.JWT_AUDIENCE,
            issuer=auth_config.JWT_ISSUER,
            options={
                "verify_aud": True,
                "verify_exp": True,
                "verify_iat": True,
                "require": ["exp", "sub", "aud", "iss", "iat", "jti"],
            },
        )
        return payload


# #endregion decode_token


# #region _hash_token [TYPE Function]
# @BRIEF SHA-256 hash of a token string for blacklist lookup.
# @PRE  token is a non-empty JWT string.
# @POST  Returns 64-character hex digest.
def _hash_token(token: str) -> str:
    return hashlib.sha256(token.encode()).hexdigest()


# #endregion _hash_token


# #region _prune_blacklist [TYPE Function]
# @BRIEF Delete expired entries from the token blacklist.
# @PRE  db is a valid SQLAlchemy session.
# @POST  Expired blacklist rows are removed.
def _prune_blacklist(db: Session) -> None:
    now = datetime.now(UTC)
    db.query(TokenBlacklist).filter(TokenBlacklist.expires_at < now).delete()
    db.commit()


# #endregion _prune_blacklist


# #region blacklist_token [TYPE Function]
# @BRIEF Revoke a JWT token by adding its hash to the blacklist.
# @PRE  token was issued by this service and hasn't expired.
# @POST  Token is blacklisted and subsequent decode_token calls will reject it.
# @SIDE_EFFECT  Writes to token_blacklist table; prunes expired entries.
# @RELATION DEPENDS_ON -> [TokenBlacklist]
def blacklist_token(token: str, db: Session) -> None:
    with belief_scope("blacklist_token"):
        _prune_blacklist(db)
        try:
            payload = decode_token(token)
        except Exception:
            return  # Already invalid — no need to blacklist

        jti = payload.get("jti")
        exp_raw = payload.get("exp")
        if not jti or not exp_raw:
            return

        if isinstance(exp_raw, datetime):
            expires_at = exp_raw
        else:
            expires_at = datetime.fromtimestamp(exp_raw, tz=UTC)

        existing = db.query(TokenBlacklist).filter(TokenBlacklist.jti == jti).first()
        if existing:
            return

        entry = TokenBlacklist(
            jti=jti,
            token_hash=_hash_token(token),
            expires_at=expires_at,
        )
        db.add(entry)
        db.commit()


# #endregion blacklist_token


# #region is_token_blacklisted [TYPE Function]
# @BRIEF Check if a JWT token has been revoked.
# @PRE  db is a valid SQLAlchemy session.
# @POST  Returns True if token is blacklisted, False otherwise.
# @RELATION DEPENDS_ON -> [TokenBlacklist]
def is_token_blacklisted(token: str, db: Session) -> bool:
    with belief_scope("is_token_blacklisted"):
        try:
            payload = decode_token(token)
        except Exception:
            return True  # Invalid tokens are treated as blacklisted

        jti = payload.get("jti")
        if not jti:
            return False

        return db.query(TokenBlacklist).filter(TokenBlacklist.jti == jti).first() is not None


# #endregion is_token_blacklisted

# #endregion AuthJwtModule
