from __future__ import annotations

import base64
import hashlib
import os
from dataclasses import dataclass
from functools import lru_cache

from fastapi import HTTPException
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa


@dataclass(frozen=True)
class LoginPublicKey:
    key_id: str
    public_key_pem: str
    algorithm: str = "RSA-OAEP-SHA256"


def _normalize_env_pem(value: str) -> bytes:
    return value.replace("\\n", "\n").encode("utf-8")


@lru_cache(maxsize=1)
def _load_private_key():
    private_key_path = os.getenv("AUTH_LOGIN_PRIVATE_KEY_PATH")
    private_key_pem = os.getenv("AUTH_LOGIN_PRIVATE_KEY")

    if private_key_path:
        with open(private_key_path, "rb") as key_file:
            key_data = key_file.read()
        return serialization.load_pem_private_key(key_data, password=None)

    if private_key_pem:
        return serialization.load_pem_private_key(_normalize_env_pem(private_key_pem), password=None)

    return rsa.generate_private_key(public_exponent=65537, key_size=2048)


@lru_cache(maxsize=1)
def get_login_public_key() -> LoginPublicKey:
    private_key = _load_private_key()
    public_key = private_key.public_key()
    public_der = public_key.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    public_pem = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    ).decode("utf-8")
    key_id = hashlib.sha256(public_der).hexdigest()[:16]
    return LoginPublicKey(key_id=key_id, public_key_pem=public_pem)


def decrypt_login_password(password_ciphertext: str, login_key_id: str) -> str:
    key_info = get_login_public_key()
    if login_key_id != key_info.key_id:
        raise HTTPException(status_code=400, detail="登录密钥已过期，请刷新后重试")

    try:
        ciphertext = base64.b64decode(password_ciphertext.encode("utf-8"), validate=True)
        plaintext = _load_private_key().decrypt(
            ciphertext,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None,
            ),
        )
        return plaintext.decode("utf-8")
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=400, detail="登录密码密文无效，请刷新后重试") from exc