"""Token-based authentication using HMAC-SHA256."""
from __future__ import annotations

import base64
import hashlib
import hmac
import json
import logging
import time
from typing import Callable, Dict, Optional, Sequence

from fastapi import Depends, HTTPException, Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse

logger = logging.getLogger(__name__)

# Paths that never require authentication
DEFAULT_EXEMPT_PATHS = ("/health", "/docs", "/openapi.json")


def generate_token(
    node_id: str,
    secret: str,
    ttl_hours: int = 24,
    role: str = "node",
    extra: Optional[Dict] = None,
) -> str:
    """Create an HMAC-SHA256 signed token with expiry.

    Format: base64(json_payload).base64(signature)
    """
    payload = {
        "node_id": node_id,
        "role": role,
        "iat": int(time.time()),
        "exp": int(time.time()) + ttl_hours * 3600,
    }
    if extra:
        payload.update(extra)

    payload_bytes = json.dumps(payload, separators=(",", ":"), sort_keys=True).encode()
    payload_b64 = base64.urlsafe_b64encode(payload_bytes).decode()

    sig = hmac.new(secret.encode(), payload_bytes, hashlib.sha256).digest()
    sig_b64 = base64.urlsafe_b64encode(sig).decode()

    return f"{payload_b64}.{sig_b64}"


def verify_token(token: str, secret: str) -> dict:
    """Verify an HMAC-SHA256 token and return its payload.

    Raises ValueError on invalid or expired tokens.
    """
    parts = token.split(".")
    if len(parts) != 2:
        raise ValueError("Malformed token")

    payload_b64, sig_b64 = parts

    try:
        payload_bytes = base64.urlsafe_b64decode(payload_b64)
        provided_sig = base64.urlsafe_b64decode(sig_b64)
    except Exception:
        raise ValueError("Malformed token encoding")

    expected_sig = hmac.new(secret.encode(), payload_bytes, hashlib.sha256).digest()
    if not hmac.compare_digest(provided_sig, expected_sig):
        raise ValueError("Invalid token signature")

    try:
        payload = json.loads(payload_bytes)
    except json.JSONDecodeError:
        raise ValueError("Corrupt token payload")

    exp = payload.get("exp", 0)
    if exp and time.time() > exp:
        raise ValueError("Token expired")

    return payload


class AuthMiddleware(BaseHTTPMiddleware):
    """Starlette middleware that enforces Bearer token authentication."""

    def __init__(
        self,
        app,
        secret: str,
        exempt_paths: Sequence[str] = DEFAULT_EXEMPT_PATHS,
        token_store=None,
    ):
        super().__init__(app)
        self.secret = secret
        self.exempt_paths = tuple(exempt_paths)
        self.token_store = token_store  # Optional TokenStore for revocation checks

    async def dispatch(self, request: Request, call_next):
        path = request.url.path.rstrip("/")

        # Skip auth for exempt paths
        if path in self.exempt_paths or path.startswith("/docs") or path.startswith("/redoc"):
            return await call_next(request)

        auth_header = request.headers.get("authorization", "")
        if not auth_header.startswith("Bearer "):
            return JSONResponse(
                status_code=401,
                content={"detail": "Missing or invalid Authorization header"},
            )

        token = auth_header[7:]  # Strip "Bearer "
        try:
            payload = verify_token(token, self.secret)
        except ValueError as e:
            return JSONResponse(
                status_code=401,
                content={"detail": str(e)},
            )

        # Check token revocation if a token store is configured
        if self.token_store is not None:
            if self.token_store.is_revoked(token):
                logger.warning("Revoked token used: node=%s", payload.get("node_id", "unknown"))
                return JSONResponse(
                    status_code=401,
                    content={"detail": "Token has been revoked"},
                )

        # Attach payload to request state for downstream use
        request.state.token_payload = payload
        return await call_next(request)


def require_role(role: str) -> Callable:
    """FastAPI dependency that checks the token payload for a specific role."""

    async def _check(request: Request):
        payload = getattr(request.state, "token_payload", None)
        if payload is None:
            raise HTTPException(401, "Authentication required")
        token_role = payload.get("role", "")
        if token_role != role and token_role != "admin":
            raise HTTPException(403, f"Role '{role}' required, got '{token_role}'")
        return payload

    return Depends(_check)
