from __future__ import annotations

import re
from typing import Dict

BASE64_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
BASE64_LOOKUP: Dict[str, int] = {ch: idx for idx, ch in enumerate(BASE64_ALPHABET)}


def utf8_encode(input_str: str) -> bytes:
    return input_str.encode("utf-8")


def utf8_decode(input_bytes: bytes) -> str:
    return input_bytes.decode("utf-8")


def encode_base64(data: bytes) -> str:
    out = []
    for i in range(0, len(data), 3):
        b0 = data[i] if i < len(data) else 0
        b1 = data[i + 1] if i + 1 < len(data) else 0
        b2 = data[i + 2] if i + 2 < len(data) else 0
        n = (b0 << 16) | (b1 << 8) | b2

        out.append(BASE64_ALPHABET[(n >> 18) & 63])
        out.append(BASE64_ALPHABET[(n >> 12) & 63])
        out.append(BASE64_ALPHABET[(n >> 6) & 63] if i + 1 < len(data) else "=")
        out.append(BASE64_ALPHABET[n & 63] if i + 2 < len(data) else "=")

    return "".join(out)


def decode_base64(base64: str) -> bytes:
    clean = re.sub(r"\s+", "", base64)
    if len(clean) % 4 != 0:
        raise ValueError("Invalid base64 length")

    out = bytearray()

    for i in range(0, len(clean), 4):
        c0, c1, c2, c3 = clean[i : i + 4]

        v0 = BASE64_LOOKUP.get(c0)
        v1 = BASE64_LOOKUP.get(c1)
        v2 = 0 if c2 == "=" else BASE64_LOOKUP.get(c2)
        v3 = 0 if c3 == "=" else BASE64_LOOKUP.get(c3)

        if v0 is None or v1 is None or (c2 != "=" and v2 is None) or (c3 != "=" and v3 is None):
            raise ValueError("Invalid base64 input")

        n = (v0 << 18) | (v1 << 12) | (v2 << 6) | v3
        out.append((n >> 16) & 255)
        if c2 != "=":
            out.append((n >> 8) & 255)
        if c3 != "=":
            out.append(n & 255)

    return bytes(out)


def encode_base64url(data: bytes) -> str:
    return encode_base64(data).replace("+", "-").replace("/", "_").rstrip("=")


def decode_base64url(input_str: str) -> bytes:
    base64 = input_str.replace("-", "+").replace("_", "/")
    pad = (4 - (len(base64) % 4)) % 4
    return decode_base64(base64 + ("=" * pad))


utf8Encode = utf8_encode
utf8Decode = utf8_decode
encodeBase64 = encode_base64
decodeBase64 = decode_base64
encodeBase64Url = encode_base64url
decodeBase64Url = decode_base64url

__all__ = [
    "decode_base64",
    "decode_base64url",
    "decodeBase64",
    "decodeBase64Url",
    "encode_base64",
    "encode_base64url",
    "encodeBase64",
    "encodeBase64Url",
    "utf8_decode",
    "utf8_encode",
    "utf8Decode",
    "utf8Encode",
]
