"""
Z-Arch Auth Module
-----------------------------------------------
Runtime entry point for all Z-Arch authentication services.

Provides two main namespaces:
    - ZArchAuth.session : session-cookie and login management
    - ZArchAuth.s2s     : service-to-service authentication

Usage Example:
    from zarch import ZArchAuth
    auth = ZArchAuth()

    # Start the session service Flask app
    app = auth.session.start()

    # Verify a user's encrypted session cookie
    claims = auth.session.verify(cookie_val, expect_uid="123")

    # Mint and verify internal service-to-service tokens
    token = auth.s2s.sign(req, target="backend")
    claims = auth.s2s.verify(req)

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

from .session import entry_point, verify_encrypted_cookie, register_session_hook
from .s2s import mint_s2s_token, verify_s2s_token
import os
import json
import urllib.request
import urllib.error
from typing import Optional, Tuple, Any
from google.auth.transport import requests as google_requests
from google.oauth2 import id_token
from flask import Request


class ZArchAuth:
    """
    Root controller exposing modular authentication subsystems for Z-Arch.

    Attributes:
        app (flask.Flask | None): Reference to the active Flask app if started.
        session (_Session): Interface for session cookie handling and login endpoints.
        s2s (_S2S): Interface for service-to-service JWT minting and verification.
    """

    def __init__(self):
        # Future: Verify Z-Arch license or entitlement key in Secret Manager
        self.app = None
        self.session = self._Session(self)
        self.s2s = self._S2S(self)

    # ==============================================================
    # Session namespace
    # ==============================================================
    class _Session:
        """Namespace for session-cookie management and verification."""

        def __init__(self, parent: "ZArchAuth"):
            self.parent = parent

        def start(self):
            """
            Initialize and return the Z-Arch session Flask app.

            Starts the internal session management service that exposes:
                - GET  /session          : Retrieve user identity metadata
                - POST /session/login    : Mint encrypted session cookie
                - POST /session/logout   : Revoke session cookie
                - POST /session/verify   : Internal stateful verification endpoint
                - GET  /session/health   : Health check endpoint

            Returns:
                flask.Flask: The initialized Flask application instance.
            """
            self.parent.app = entry_point()
            return self.parent.app

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

            Supported hooks:
                - on_login: Called after minting a session cookie.
                            Signature: (sid, uid, tenant, iat, exp) -> None
                - on_logout: Called during logout to revoke a session.
                             Signature: (sid, uid, tenant) -> None
                - on_verify: Called to verify session state against a backend store.
                             Signature: (sid, uid, tenant, iat, exp) -> bool

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

        def verify(
            self,
            cookie_val: str,
            expect_uid: Optional[str] = None,
            expect_tenant: Optional[str] = None,
            stateful_session: Optional[str] = None,
            session_service_url: Optional[str] = None,
        ) -> dict:
            """
            Decrypt and validate the encrypted Z-Arch session cookie.

            Performs two-phase verification:
                Phase 1: Cryptographic verification of the encrypted cookie (local, never bypassable).
                Phase 2: Optional authoritative state verification by calling the session service
                         when stateful sessions are enabled.

            Args:
                cookie_val (str): The raw encrypted cookie value.
                expect_uid (str | None): Expected user UID; if provided, must match cookie.
                expect_tenant (str | None): Expected tenant identifier; if provided, must match.
                stateful_session (str | None): Set to "TRUE" to enable stateful verification.
                                               When not "TRUE", verification is stateless after
                                               cryptographic checks.
                session_service_url (str | None): URL of the session service for stateful
                                                  verification. Required when stateful_session
                                                  is "TRUE".

            Returns:
                dict: Decoded cookie payload containing at least:
                      {
                        "sid": <session_id>,
                        "uid": <user_id>,
                        "tenant": <tenant_id>,
                        "iat": <issued_at>,
                        "exp": <expires_at>,
                        "v": <schema_version>
                      }

            Raises:
                ValueError: If cookie is missing, invalid, tampered, expired, or if stateful
                            verification fails (session revoked, service unreachable, or any
                            non-200 response).
            """
            # Phase 1: Cryptographic verification (always performed, never bypassable)
            payload = verify_encrypted_cookie(cookie_val, expect_uid, expect_tenant)

            # Phase 2: Optional stateful verification
            if stateful_session == "TRUE":
                if not session_service_url:
                    raise ValueError("SESSION_SERVICE_URL required for stateful session verification")

                self._verify_stateful(payload, session_service_url)

            return payload

        def _verify_stateful(self, payload: dict, session_service_url: str) -> None:
            """
            Perform stateful session verification via the session service.

            Makes a blocking S2S-authenticated HTTP call to the session service
            verification endpoint. Fails closed on any error.

            Args:
                payload: The cryptographically verified cookie payload.
                session_service_url: Base URL of the session service.

            Raises:
                ValueError: On any failure, timeout, or non-200 response.
            """
            sid = payload.get("sid")
            if not sid:
                raise ValueError("Session cookie missing sid for stateful verification")

            # Build verification request
            verify_url = session_service_url.rstrip("/") + "/session/verify"
            request_body = json.dumps({
                "sid": sid,
                "uid": payload.get("uid"),
                "tenant": payload.get("tenant"),
                "iat": payload.get("iat"),
                "exp": payload.get("exp"),
            }).encode("utf-8")

            req = urllib.request.Request(
                verify_url,
                data=request_body,
                method="POST",
                headers={
                    "Content-Type": "application/json",
                },
            )
            try:
                self.parent.s2s.sign(req, target="session", url=session_service_url)
            except PermissionError as e:
                raise ValueError(f"Failed to mint S2S token for session service: {e}")

            # Fail closed: any error, timeout, or non-200 response denies the request
            try:
                with urllib.request.urlopen(req, timeout=300) as resp:
                    if resp.status != 200:
                        raise ValueError("Session verification denied by session service")
                    response_data = json.loads(resp.read().decode("utf-8"))
                    if not response_data.get("valid"):
                        raise ValueError("Session revoked or invalid")
            except urllib.error.HTTPError as e:
                raise ValueError(f"Session verification failed: HTTP {e.code}")
            except urllib.error.URLError as e:
                raise ValueError(f"Session service unreachable: {e.reason}")
            except json.JSONDecodeError:
                raise ValueError("Invalid response from session service")
            except TimeoutError:
                raise ValueError("Session service timeout")

    # ==============================================================
    # S2S namespace
    # ==============================================================
    class _S2S:
        """Namespace for internal service-to-service trust and JWT operations."""

        def __init__(self, parent: "ZArchAuth"):
            self.parent = parent

        def sign(self, req: Any, target: str, url: str = ''):
            """
            Add a signed JWT to x-zarch-s2s-token header for authorizing this service to call another internal service.
            If GCP platform is detected and an upstream URL is provided, a Google OIDC token is added to the Authorization header.

            The Z-Arch S2S token represents an authenticated, short-lived identity assertion from
            the current service (iss) to its intended recipient (aud).

            Args:
                req (obj): The inbound HTTP request. Any object with a mutable `headers` attribute.
                target (str): The service ID of the intended audience, as declared
                              in build-time configuration.
                url (str): URL of the upstream service being called. Used to mint the Google OIDC token. 

            Raises:
                PermissionError: If the target service is not authorized per the
                                 embedded S2S_ALLOWED_TARGETS policy.
            """
            platform = os.getenv("ZARCH_PLATFORM")

            zarch_token = mint_s2s_token(target)
            req.headers["x-zarch-s2s-token"] = zarch_token

            if platform == "gcp":
                if not url:
                    print("[WARNING ZArchAuth] No URL provided to mint Google OIDC token")
                else:
                    req.headers["Authorization"] = f"Bearer {id_token.fetch_id_token(google_requests.Request(), url)}"


        def verify(self, req) -> dict:
            """
            Validate an inbound Z-Arch service-to-service JWT from the x-zarch-s2s-token header.

            Ensures the signature matches a trusted caller's public key, audience equals
            this service, and token freshness and policy constraints are respected.

            Args:
                req (flask.Request): The inbound HTTP request containing
                                     `x-zarch-s2s-token: <token>` header.

            Returns:
                dict: Verified JWT claims dictionary identical to the minted payload.

            Raises:
                PermissionError: If signature invalid, token expired, audience mismatch,
                                 or the caller is not authorized in the trust graph.
            """
            return verify_s2s_token(req)
