import base64
import hashlib
import json
import secrets
from functools import partial
from typing import Any, TypedDict

json_dump = partial(json.dumps, separators=(",", ":"))


def b64encode(b: bytes) -> str:
    return base64.b64encode(b).decode().rstrip("=")


def b64decode(s: str) -> bytes:
    return base64.b64decode(b64padded(s))


def b64padded(s: str) -> str:
    return s + "=" * (4 - len(s) % 4)


def _encode_params(params: dict[str, Any]) -> str:
    return ",".join(f"{k}={v}" for k, v in params.items())


class _PhcScryptParams(TypedDict, total=False):
    n: int
    r: int
    p: int


def phc_scrypt(
    secret: bytes,
    *,
    salt: bytes | None = None,
    params: _PhcScryptParams = {},  # noqa: B006
) -> str:
    """Return the scrypt expanded secret in PHC string format.

    Uses somewhat sane defaults.

    For more information on the PHC string format, see
    https://github.com/P-H-C/phc-string-format/blob/master/phc-sf-spec.md
    """

    if salt is None:
        salt = secrets.token_bytes(16)

    n = params.get("n", 2**14)  # CPU/Memory cost factor
    r = params.get("r", 8)  # block size
    p = params.get("p", 1)  # parallelization factor
    # maxmem = 2 * 128 * n * r * p
    hashed_secret = hashlib.scrypt(secret, salt=salt, n=n, r=r, p=p)

    phc = "".join(
        f"${x}"
        for x in [
            "scrypt",
            _encode_params({"n": n, "r": r, "p": p}),
            b64encode(salt),
            b64encode(hashed_secret),
        ]
    )

    return phc


def phc_compare(*, secret: str, phc_string: str) -> bool:
    args = parse_phc(phc_string)

    if args["id"] != "scrypt":
        raise ValueError(f"Algorithm not supported: {args['id']}")

    encoded = phc_scrypt(b64decode(secret), salt=args["salt"], params=args["params"])

    return secrets.compare_digest(encoded, phc_string)


class _PhcParts(TypedDict):
    # $<id>[$v=<version>][$<param>=<value>(,<param>=<value>)*][$<salt>[$<hash>]]
    id: str  # the symbolic name for the function
    version: int | None  # the algorithm version
    params: dict[str, int]
    salt: bytes | None
    hash: bytes | None


def parse_phc(s: str) -> _PhcParts:
    parts = _PhcParts(id="", version=None, params={}, salt=None, hash=None)

    _, parts["id"], *rest = s.split("$")

    if rest and rest[0].startswith("v="):
        parts["version"] = int(rest.pop(0))
    if rest and "=" in rest[0]:
        parts["params"] = {
            kv[0]: int(kv[1])
            for p in rest.pop(0).split(",")
            if len(kv := p.split("=", 2)) == 2
        }
    if rest:
        parts["salt"] = b64decode(rest.pop(0))
    if rest:
        parts["hash"] = b64decode(rest.pop(0))

    return parts
