import base64
import os

import cryptography.hazmat.primitives.kdf.scrypt
import cryptography.exceptions

from typing import (
	Literal,
	overload,
	Optional,
)


class PasswordUtils:
	@overload
	@classmethod
	def secret_hash(
		cls,
		secret: str | bytes,
		mode: Literal['base64'],
		salt: Optional[bytes] = None,
	) -> tuple[str, str]: ...

	@overload
	@classmethod
	def secret_hash(
		cls,
		secret: str | bytes,
		mode: Literal['bytes'],
		salt: Optional[bytes] = None,
	) -> tuple[bytes, bytes]: ...

	@classmethod
	def secret_hash(
		cls,
		secret: str | bytes,
		mode: Literal['bytes', 'base64'],
		salt: Optional[bytes] = None,
	) -> tuple[str, str] | tuple[bytes, bytes]:
		if salt is None:
			salt = os.urandom(16)

		if isinstance(secret, str):
			secret = secret.encode('utf-8')
		# derive
		kdf = cls._scrypt_init(salt=salt)

		hashed_secret = kdf.derive(secret)

		if mode == 'bytes':
			return (salt, hashed_secret)
		elif mode == 'base64':
			res_tuple = tuple(
				(
					base64.b64encode(o).decode('utf-8')
					for o in (
						salt,
						hashed_secret,
					)
				)
			)
			return (res_tuple[0], res_tuple[1])
		else:
			raise NotImplementedError

	@classmethod
	def _scrypt_init(cls, salt: bytes) -> cryptography.hazmat.primitives.kdf.scrypt.Scrypt:
		return cryptography.hazmat.primitives.kdf.scrypt.Scrypt(
			salt=salt,
			length=32,
			n=2**14,
			r=8,
			p=1,
		)

	@classmethod
	def secret_check(
		cls,
		secret: str | bytes,
		salt: str | bytes,
		hashed_secret: str | bytes,
	) -> bool:
		if isinstance(salt, str):
			salt = base64.b64decode(salt)

		if isinstance(secret, str):
			secret = secret.encode('utf-8')

		if isinstance(hashed_secret, str):
			hashed_secret = base64.b64decode(hashed_secret)

		kdf = cls._scrypt_init(salt=salt)

		try:
			kdf.verify(secret, hashed_secret)
			return True
		except cryptography.exceptions.InvalidKey:
			return False
