from dataclasses import dataclass
import hashlib, base64, hmac
from typing import Any
import json, time

from starlette.requests import Request
from starlette.responses import JSONResponse

from app.config import JWT_SECRET
from app.policy import TokenPolicy, tuple_claim

@dataclass(frozen=True)
class AuthContext:
  policy:TokenPolicy|None

def error(status_code:int, detail:str) -> JSONResponse:
  return JSONResponse({"detail": detail}, status_code=status_code)

def b64url_decode(value:str) -> bytes:
  padding = "=" * (-len(value) % 4)
  return base64.urlsafe_b64decode((value + padding).encode())

def decode_hs256_jwt(token:str, secret:bytes) -> dict[str, Any]:
  parts = token.split(".")
  if len(parts) != 3:
    raise ValueError("JWT must have three parts")

  header_b64, payload_b64, signature_b64 = parts
  try:
    header = json.loads(b64url_decode(header_b64))
    payload = json.loads(b64url_decode(payload_b64))
  except (json.JSONDecodeError, ValueError) as exc:
    raise ValueError("JWT header or payload is invalid") from exc

  if not isinstance(header, dict) or not isinstance(payload, dict):
    raise ValueError("JWT header and payload must be JSON objects")
  if header.get("alg") != "HS256":
    raise ValueError("only HS256 JWTs are supported")
  if header.get("typ") not in (None, "JWT"):
    raise ValueError("JWT typ must be JWT")

  signing_input = f"{header_b64}.{payload_b64}".encode()
  expected = hmac.new(secret, signing_input, hashlib.sha256).digest()
  supplied = b64url_decode(signature_b64)
  if not hmac.compare_digest(supplied, expected):
    raise ValueError("JWT signature is invalid")

  now = int(time.time())
  exp = payload.get("exp")
  if exp is not None and (not isinstance(exp, int | float) or now >= int(exp)):
    raise ValueError("JWT has expired")
  nbf = payload.get("nbf")
  if nbf is not None and (not isinstance(nbf, int | float) or now < int(nbf)):
    raise ValueError("JWT is not active yet")
  iat = payload.get("iat")
  if iat is not None and (not isinstance(iat, int | float) or int(iat) > now + 300):
    raise ValueError("JWT iat is in the future")

  return payload

def policy_from_jwt_claims(token:str, claims:dict[str, Any]) -> TokenPolicy:
  name = str(claims.get("name") or claims.get("sub") or "jwt")
  return TokenPolicy(
    token=token,
    name=name,
    vaults=tuple_claim(claims.get("vaults"), ("*",)),
    paths=tuple_claim(claims.get("paths"), ("*",)),
    commands=tuple_claim(claims.get("commands"), ("*",)),
  )

def require_auth(request:Request) -> AuthContext|JSONResponse:
  if not JWT_SECRET:
    return error(500, "OBSIDIAN_JWT_SECRET is required")

  header = request.headers.get("authorization", "")
  if not header.startswith("Bearer "):
    return error(401, "Unauthorized")

  supplied = header.removeprefix("Bearer ")
  try:
    claims = decode_hs256_jwt(supplied, JWT_SECRET.encode())
    return AuthContext(policy=policy_from_jwt_claims(supplied, claims))
  except ValueError:
    return error(401, "Unauthorized")