"""
Z-Arch Session Service
-----------------------------------------------
Session lifecycle service for Z-Arch's dual-auth model.

This module:
  1. Mints encrypted `__session` cookies after Gateway JWT validation.
  2. Enforces strict JWT-to-cookie identity binding on protected flows.
  3. Exposes optional stateful verification hooks for revocation checks.
  4. Requires verified S2S tokens on inbound internal requests.

Runtime contract:
  - Gateway validates end-user JWTs and forwards claims in `X-Userinfo`.
  - `/session/login` creates an encrypted session cookie bound to user identity.
  - `/session/logout` verifies cookie and claims, invokes revocation hook, and clears cookie.
  - `/session/verify` allows gateway-side stateful checks when enabled.
  - `/session` returns minimal identity metadata for authenticated callers.

Security invariants:
  - Cookie integrity is enforced via Fernet encryption + authentication.
  - Expired, malformed, or tampered cookies are rejected.
  - Cookie UID/tenant must match JWT-derived identity for sensitive operations.
  - Internal session endpoints are protected by Z-Arch S2S token verification.

Copyright © 2026 RAM Cloud Code LLC
Licensed under the Apache License, Version 2.0.
"""

from flask import Flask, request, make_response, abort, g
from cryptography.fernet import Fernet, InvalidToken
import os, json, time, secrets
from zarch.auth.s2s import verify_s2s_token
from pathlib import Path
from typing import Callable, Optional, Dict, Any

# ---------------------------------------------
# Configuration
# ---------------------------------------------
COOKIE_NAME = "__session"
COOKIE_DOMAIN = os.getenv("COOKIE_DOMAIN", None) or None
COOKIE_SAMESITE = os.getenv("SESSION_SAMESITE", "Lax")
COOKIE_SECURE = True
COOKIE_HTTPONLY = True
COOKIE_TTL_SECS = int(os.getenv("SESSION_TTL_SECS", "1209600"))  # 14 days
MAX_IAT_SKEW = int(os.getenv("IAT_SKEW_SECS", "60"))  # seconds allowed between JWT iat and now

_fernet: Optional[Fernet] = None


def _get_fernet() -> Fernet:
    """Lazily load the Fernet key so imports don't fail when the secret isn't mounted."""
    global _fernet
    if _fernet is None:
        try:
            key = Path("/zarch/mounts/session_key/session_key").read_text().rstrip("\n")
        except FileNotFoundError as exc:
            raise RuntimeError("Missing session key at /zarch/mounts/session_key/session_key") from exc
        _fernet = Fernet(key)
    return _fernet

# ---------------------------------------------
# Session Hooks Registry
# ---------------------------------------------
_session_hooks: Dict[str, Optional[Callable]] = {
    "on_login": None,
    "on_logout": None,
    "on_verify": None,
}


def register_session_hook(hook_name: str, callback: Callable) -> None:
    """
    Register a user-defined hook for session lifecycle events.

    Supported hooks:
        - on_login: Called after minting a session cookie.
                    Signature: (sid: str, uid: str, tenant: str | None, iat: int, exp: int) -> None
        - on_logout: Called during logout to revoke a session.
                     Signature: (sid: str, uid: str, tenant: str | None) -> None
        - on_verify: Called to verify session state against a backend store.
                     Signature: (sid: str, uid: str, tenant: str | None, iat: int, exp: int) -> bool
                     Must return True if session is valid, False otherwise.

    Args:
        hook_name: One of 'on_login', 'on_logout', 'on_verify'.
        callback: The function to invoke for this hook.

    Raises:
        ValueError: If hook_name is not recognized.
    """
    if hook_name not in _session_hooks:
        raise ValueError(f"Unknown session hook: {hook_name}. Valid hooks: {list(_session_hooks.keys())}")
    _session_hooks[hook_name] = callback


def _invoke_hook(hook_name: str, *args, **kwargs) -> Any:
    """Invoke a registered hook if present. Returns None if hook not registered."""
    hook = _session_hooks.get(hook_name)
    if hook is not None:
        return hook(*args, **kwargs)
    return None


# ---------------------------------------------
# Helpers
# ---------------------------------------------
def _generate_sid() -> str:
    """Generate a cryptographically strong session ID."""
    return secrets.token_urlsafe(32)


def verify_encrypted_cookie(cookie_val: str,
                            expect_uid: str | None = None,
                            expect_tenant: str | None = None) -> dict:
    """
    Decrypt and verify the encrypted session cookie.

    Args:
        cookie_val (str): The raw encrypted cookie value.
        expect_uid (str, optional): If provided, ensure cookie uid matches.
        expect_tenant (str, optional): If provided, ensure cookie tenant matches.

    Returns:
        dict: The decrypted cookie payload if valid.

    Raises:
        ValueError: If cookie is missing, invalid, tampered, or expired.
    """
    if not cookie_val:
        raise ValueError("Missing session cookie")

    try:
        raw = _get_fernet().decrypt(cookie_val.encode())
        data = json.loads(raw.decode())
    except (InvalidToken, json.JSONDecodeError):
        raise ValueError("Invalid or tampered session cookie")

    exp = int(data.get("exp", 0))
    now = int(time.time())
    if exp < now:
        raise ValueError("Session cookie expired")

    if expect_uid and data.get("uid") != expect_uid:
        raise ValueError("Cookie UID mismatch")
    if expect_tenant and data.get("tenant") != expect_tenant:
        raise ValueError("Cookie tenant mismatch")

    return data


def _get_userinfo() -> dict:
    """Decode and validate the user info header from Gateway."""
    hdr = request.headers.get("X-Userinfo")
    if not hdr:
        abort(401, description="Missing user info header")
    try:
        data = json.loads(hdr)
    except json.JSONDecodeError:
        abort(401, description="Malformed user info header")
    return data


def _verify_iat_freshness(claims: dict):
    """Reject tokens whose iat is too old or in the future."""
    iat = int(claims.get("iat", 0))
    now = int(time.time())
    if abs(now - iat) > MAX_IAT_SKEW:
        abort(401, description="Token iat too old or invalid")


def _mint_cookie_payload(claims: dict) -> tuple[bytes, str]:
    """
    Build and encrypt session cookie payload.

    Returns:
        Tuple of (encrypted_cookie_bytes, sid)
    """
    sid = _generate_sid()
    now = int(time.time())
    exp = now + COOKIE_TTL_SECS
    payload = {
        "sid": sid,
        "uid": claims.get("sub") or claims.get("uid") or claims.get("user_id"),
        "tenant": claims.get("firebase", {}).get("tenant") or claims.get("tenant") or claims.get("tid"),
        "iat": now,
        "exp": exp,
        "v": 1,
    }
    return _get_fernet().encrypt(json.dumps(payload).encode()), sid


def _decrypt_cookie(cookie_val: str) -> dict:
    """Decrypt and validate session cookie payload."""
    try:
        raw = _get_fernet().decrypt(cookie_val.encode())
        data = json.loads(raw.decode())
    except (InvalidToken, json.JSONDecodeError):
        abort(401, description="Invalid or tampered cookie")
    if int(data.get("exp", 0)) < time.time():
        abort(401, description="Expired cookie")
    return data

# ---------------------------------------------
# Flask App entry
# ---------------------------------------------
def entry_point():
    app = Flask(__name__)

    # ---------------------------------------------
    # s2s hop verification
    # ---------------------------------------------
    @app.before_request
    def verify_internal_caller():
        """
        Flask pre-request hook that validates inbound Z-Arch S2S tokens.

        - Runs before every route handler.
        - Aborts with 401 if the token is missing or invalid.
        - On success, attaches verified claims to request context.
        """
        try:
            claims = verify_s2s_token(request)
            # attach claims to flask.g for downstream use
            g.s2s_claims = claims
        except PermissionError as e:
            abort(401, description=f"S2S token invalid: {e}")

    # ---------------------------------------------
    # /session/login
    # ---------------------------------------------
    @app.post("/session/login")
    def session_login():
        """Create encrypted session cookie from Gateway-validated claims."""
        # Gateway already verified the JWT. Just parse the headers.
        claims = _get_userinfo()
        _verify_iat_freshness(claims)

        cookie_val, sid = _mint_cookie_payload(claims)

        # Extract identity for hook invocation
        uid = claims.get("sub") or claims.get("uid") or claims.get("user_id")
        tenant = claims.get("firebase", {}).get("tenant") or claims.get("tenant") or claims.get("tid")
        now = int(time.time())
        exp = now + COOKIE_TTL_SECS

        # Invoke user-defined on_login hook to persist session state
        _invoke_hook("on_login", sid=sid, uid=uid, tenant=tenant, iat=now, exp=exp)

        resp = make_response({"ok": True, "uid": uid})
        resp.set_cookie(
            key=COOKIE_NAME,
            value=cookie_val.decode(),
            max_age=COOKIE_TTL_SECS,
            httponly=COOKIE_HTTPONLY,
            secure=COOKIE_SECURE,
            samesite=COOKIE_SAMESITE,
            domain=COOKIE_DOMAIN,
            path="/",
        )
        resp.headers["Cache-Control"] = "no-store"
        return resp

    # ---------------------------------------------
    # /session/logout
    # ---------------------------------------------
    @app.post("/session/logout")
    def session_logout():
        """Validate existing cookie and clear it."""
        # JWT is validated by Gateway openid-connect plugin upstream.
        # We verify the header format here as a sanity check.
        hdr = request.headers.get("Authorization", "")
        if not hdr.startswith("Bearer "):
            abort(401)

        claims = _get_userinfo()
        cookie_val = request.cookies.get(COOKIE_NAME)
        if not cookie_val:
            abort(401, description="Missing session cookie")

        payload = _decrypt_cookie(cookie_val)

        # Strict identity binding
        jwt_uid = claims.get("sub") or claims.get("uid") or claims.get("user_id")
        jwt_tenant = claims.get("firebase", {}).get("tenant") or claims.get("tenant") or claims.get("tid")
        if payload.get("uid") != jwt_uid:
            abort(401, description="UID mismatch")
        if payload.get("tenant") != jwt_tenant:
            abort(401, description="Tenant mismatch")

        # Invoke user-defined on_logout hook to revoke session (idempotent)
        sid = payload.get("sid")
        if sid:
            _invoke_hook("on_logout", sid=sid, uid=jwt_uid, tenant=jwt_tenant)

        # Invalidate cookie client-side
        resp = make_response({"ok": True})
        resp.set_cookie(
            key=COOKIE_NAME,
            value="",
            max_age=0,
            httponly=COOKIE_HTTPONLY,
            secure=COOKIE_SECURE,
            samesite=COOKIE_SAMESITE,
            domain=COOKIE_DOMAIN,
            path="/",
        )
        resp.headers["Cache-Control"] = "no-store"
        return resp

    # ---------------------------------------------
    # /session/verify  –  Internal verification endpoint
    # ---------------------------------------------
    @app.post("/session/verify")
    def session_verify():
        """
        Internal endpoint for stateful session verification.

        Called by the gateway when stateful sessions are enabled.
        Receives session metadata (never the raw cookie) and invokes
        the on_verify hook to check session state against a backend store.

        Request JSON body:
            {
                "sid": <session_id>,
                "uid": <user_id>,
                "tenant": <tenant_id | null>,
                "iat": <issued_at_timestamp>,
                "exp": <expiration_timestamp>
            }

        Returns:
            200 OK with {"valid": true} if session is valid.
            403 Forbidden if session is revoked or not found.
            400 Bad Request if request body is malformed.
        """
        try:
            data = request.get_json(force=True)
        except Exception:
            abort(400, description="Invalid JSON body")

        sid = data.get("sid")
        uid = data.get("uid")
        tenant = data.get("tenant")
        iat = data.get("iat")
        exp = data.get("exp")

        # Validate required fields
        if not sid or not uid:
            abort(400, description="Missing required fields: sid, uid")
        if iat is None or exp is None:
            abort(400, description="Missing required fields: iat, exp")

        # Check expiration server-side as a safeguard
        now = int(time.time())
        if int(exp) < now:
            resp = make_response({"valid": False, "reason": "expired"}, 403)
            resp.headers["Cache-Control"] = "no-store"
            return resp

        # Invoke user-defined on_verify hook
        # If no hook registered, default to valid (stateless fallback)
        hook_result = _invoke_hook("on_verify", sid=sid, uid=uid, tenant=tenant, iat=iat, exp=exp)

        # Fail closed: if hook returns anything other than True, deny
        if hook_result is None:
            # No hook registered - stateless mode, pass through
            valid = True
        else:
            valid = hook_result is True

        if not valid:
            resp = make_response({"valid": False, "reason": "revoked_or_invalid"}, 403)
            resp.headers["Cache-Control"] = "no-store"
            return resp

        resp = make_response({"valid": True}, 200)
        resp.headers["Cache-Control"] = "no-store"
        return resp

    # ---------------------------------------------
    # /session  –  Safe identity info endpoint
    # ---------------------------------------------
    @app.get("/session")
    def session_info():
        """
        Return minimal, non-sensitive identity info from the Gateway OIDC header.
        Works for any OIDC IdP (Firebase, Auth0, Entra, etc.).
        """
        hdr = request.headers.get("X-Userinfo")
        if not hdr:
            abort(401, description="Missing X-Userinfo header")

        try:
            claims = json.loads(hdr)
        except json.JSONDecodeError:
            abort(401, description="Malformed X-Userinfo header")

        uid = claims.get("sub") or claims.get("uid") or claims.get("user_id")
        email = (
        claims.get("email")
        or claims.get("preferred_username")
        or claims.get("upn")
        or claims.get("email_verified")
        )

        tenant = claims.get("firebase", {}).get("tenant") or claims.get("tenant") or claims.get("tid")

        if not uid:
            abort(401, description="Missing sub claim")

        data = {"uid": uid}
        if email:
            data["email"] = email
        if tenant:
            data["tenant"] = tenant

        resp = make_response(data, 200)
        resp.headers["Cache-Control"] = "no-store"
        return resp

    # ---------------------------------------------
    # /session/health
    # ---------------------------------------------
    @app.get("/session/health")
    def health():
        return {"ok": True}

    return app
