from __future__ import annotations

import base64
import hashlib
import secrets
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any
from uuid import uuid4

import jwt
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa

from app.models import BackendCredential


def hash_secret(secret: str) -> str:
    return hashlib.sha256(secret.encode("utf-8")).hexdigest()


def verify_secret(secret: str, digest: str) -> bool:
    return secrets.compare_digest(hash_secret(secret), digest)


def generate_client_secret() -> str:
    return secrets.token_urlsafe(32)


class JwtSigner:
    def __init__(self, key_path: Path, issuer: str, ttl_seconds: int):
        self.key_path = key_path
        self.issuer = issuer.rstrip("/")
        self.ttl_seconds = ttl_seconds
        self._private_key = self._load_or_create_private_key()
        self._public_key = self._private_key.public_key()
        self._kid = self._compute_kid()

    def _load_or_create_private_key(self):
        self.key_path.parent.mkdir(parents=True, exist_ok=True)
        if self.key_path.exists():
            return serialization.load_pem_private_key(
                self.key_path.read_bytes(),
                password=None,
            )
        private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
        self.key_path.write_bytes(pem)
        return private_key

    def _compute_kid(self) -> str:
        der = self._public_key.public_bytes(
            encoding=serialization.Encoding.DER,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        return hashlib.sha256(der).hexdigest()[:16]

    def build_jwks(self) -> dict[str, Any]:
        public_numbers = self._public_key.public_numbers()

        def _b64url(value: int) -> str:
            raw = value.to_bytes((value.bit_length() + 7) // 8, "big")
            return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii")

        return {
            "keys": [
                {
                    "kty": "RSA",
                    "use": "sig",
                    "alg": "RS256",
                    "kid": self._kid,
                    "n": _b64url(public_numbers.n),
                    "e": _b64url(public_numbers.e),
                }
            ]
        }

    def issue_access_token(
        self,
        *,
        credential: BackendCredential,
        audience: str,
        scopes: list[str],
    ) -> tuple[str, int]:
        now = datetime.now(timezone.utc)
        exp = now + timedelta(seconds=self.ttl_seconds)
        payload = {
            "iss": self.issuer,
            "sub": f"backend:{credential.backend_id}",
            "client_id": credential.client_id,
            "backend_id": credential.backend_id,
            "aud": audience,
            "scp": scopes,
            "iat": int(now.timestamp()),
            "nbf": int(now.timestamp()),
            "exp": int(exp.timestamp()),
            "jti": str(uuid4()),
        }
        token = jwt.encode(
            payload,
            self._private_key,
            algorithm="RS256",
            headers={"kid": self._kid},
        )
        return token, self.ttl_seconds

    def decode(self, token: str) -> dict[str, Any]:
        return jwt.decode(
            token,
            self._public_key,
            algorithms=["RS256"],
            issuer=self.issuer,
            options={
                "require": ["iss", "sub", "aud", "exp", "iat", "jti"],
                "verify_aud": False,
            },
        )
