import base64
import hmac
import os
from hashlib import sha256

from django.conf import settings
from django.utils.crypto import pbkdf2


_VERSION = b"v1"
_SALT_SIZE = 16
_TAG_SIZE = 32
_PREFIX = "enc1:"


def _derive_key(salt: bytes) -> bytes:
    return pbkdf2(settings.SECRET_KEY, salt, 260000, digest=sha256)


def _keystream(key: bytes, length: int) -> bytes:
    out = bytearray()
    counter = 0
    while len(out) < length:
        block = sha256(key + counter.to_bytes(4, "big")).digest()
        out.extend(block)
        counter += 1
    return bytes(out[:length])


def encrypt_text(plaintext: str) -> str:
    if not plaintext:
        return ""

    raw = plaintext.encode("utf-8")
    salt = os.urandom(_SALT_SIZE)
    key = _derive_key(salt)
    stream = _keystream(key, len(raw))
    ciphertext = bytes(a ^ b for a, b in zip(raw, stream))
    tag = hmac.new(key, ciphertext, sha256).digest()
    payload = _VERSION + salt + tag + ciphertext
    return _PREFIX + base64.urlsafe_b64encode(payload).decode("ascii")


def decrypt_text(token: str) -> str:
    if not token:
        return ""
    if not token.startswith(_PREFIX):
        return token

    payload = base64.urlsafe_b64decode(token[len(_PREFIX) :].encode("ascii"))
    version = payload[:2]
    if version != _VERSION:
        raise ValueError("Unsupported encrypted payload version")

    salt_start = 2
    salt_end = salt_start + _SALT_SIZE
    salt = payload[salt_start:salt_end]

    tag_start = salt_end
    tag_end = tag_start + _TAG_SIZE
    tag = payload[tag_start:tag_end]

    ciphertext = payload[tag_end:]
    key = _derive_key(salt)
    expected = hmac.new(key, ciphertext, sha256).digest()
    if not hmac.compare_digest(tag, expected):
        raise ValueError("Encrypted payload validation failed")

    stream = _keystream(key, len(ciphertext))
    raw = bytes(a ^ b for a, b in zip(ciphertext, stream))
    return raw.decode("utf-8")
