"""
custom_encoder
--------------
A reusable custom encoding module.

Usage:
    from wc32 import encode, decode

    # Basic usage (deterministic - same key + text = same output)
    encrypted = encode("hello", "key123")
    decrypted = decode(encrypted, "key123")

    # With random salt (recommended - non-deterministic)
    encrypted = encode("hello", "key123", salt=True)
    decrypted = decode(encrypted, "key123")

    # With integrity check
    encrypted = encode("hello", "key123", integrity=True)
    decrypted = decode(encrypted, "key123")  # Raises error if wrong key

    # Binary data
    encrypted = encode(b"\\x00\\x01\\x02", "key123", binary=True)
    decrypted = decode(encrypted, "key123", binary=True)

    # Custom charset
    encrypted = encode("hello", "key123", charset="ABC123")
    decrypted = decode(encrypted, "key123", charset="ABC123")
"""

import base64
import hashlib
import hmac
import secrets
import string
from typing import List, Optional, Union

__all__ = ["encode", "decode", "encode_batch", "decode_batch"]

DEFAULT_CHARSET = "QwErTyUiOpAsDfGhJkLzXcVbNm0123456789@#$%^&*-_=+"
SALT_LENGTH = 8
INTEGRITY_LENGTH = 8
HEXDIGITS = string.hexdigits.lower() + string.hexdigits.upper()


def _get_charset(charset: Optional[str] = None) -> tuple:
    if charset is None or len(charset) < 10:
        return DEFAULT_CHARSET, len(DEFAULT_CHARSET)
    return charset, len(charset)


def _generate_keystream(key: str, length: int, salt: Optional[str] = None) -> List[int]:
    if not key:
        raise ValueError("Key must not be empty")

    key_material = f"{key}:{salt}" if salt else key
    seed = hashlib.sha256(key_material.encode()).hexdigest()
    keystream = []
    for i in range(length):
        piece = seed[i % len(seed)]
        keystream.append(ord(piece))
    return keystream


def _generate_integrity(text: str, key: str, charset: str, base: int) -> str:
    signature = hmac.new(key.encode(), text.encode(), hashlib.sha256).digest()
    integrity_chars = []
    for b in signature[:INTEGRITY_LENGTH]:
        integrity_chars.append(charset[b % base])
    return "".join(integrity_chars)


def _verify_integrity(text: str, key: str, expected: str, charset: str, base: int) -> bool:
    actual = _generate_integrity(text, key, charset, base)
    return actual[:INTEGRITY_LENGTH] == expected


def _generate_salt() -> str:
    return secrets.token_hex(SALT_LENGTH)


def _to_custom_base(num: int, charset: str, base: int) -> str:
    if num < 0:
        raise ValueError("Cannot encode negative numbers")

    if num == 0:
        return charset[0]

    result = ""
    while num > 0:
        num, rem = divmod(num, base)
        result = charset[rem] + result
    return result


def _from_custom_base(encoded: str, charset: str, base: int) -> int:
    if not encoded:
        raise ValueError("Encoded value cannot be empty")

    num = 0
    for char in encoded:
        if char not in charset:
            raise ValueError(f"Invalid character '{char}' in encoded string")
        num = num * base + charset.index(char)
    return num


def encode(
    text: Union[str, bytes],
    key: str,
    salt: Union[bool, str] = False,
    integrity: bool = False,
    binary: bool = False,
    charset: Optional[str] = None
) -> str:
    if not isinstance(key, str):
        raise TypeError("Key must be a string")
    if not isinstance(text, (str, bytes)):
        raise TypeError("Text must be a string or bytes")
    if not text:
        return ""

    cs, base = _get_charset(charset)

    if isinstance(text, bytes):
        if not binary:
            raise TypeError("Binary data requires binary=True")
        text = base64.b64encode(text).decode('ascii')
    elif binary:
        raise TypeError("String data cannot be used with binary=True, use bytes")

    salt_value: Optional[str] = None
    if isinstance(salt, bool):
        salt_value = _generate_salt() if salt else None
    elif isinstance(salt, str):
        salt_value = salt
    elif salt is not False:
        raise TypeError("Salt must be a boolean or string")

    keystream = _generate_keystream(key, len(text), salt_value)
    encoded_parts = []

    for i, char in enumerate(text):
        char_val = ord(char)
        if char_val > 1114111:
            raise ValueError(f"Invalid Unicode character at position {i}")
        transformed = char_val + keystream[i] + (i * 7)
        encoded_parts.append(_to_custom_base(transformed, cs, base))

    result = ".".join(encoded_parts)

    if integrity:
        integrity_tag = _generate_integrity(text, key, cs, base)
        result = f"i{integrity_tag}.{result}"

    if salt_value:
        result = f"{salt_value}.{result}"

    return result


def decode(
    encoded_text: str,
    key: str,
    binary: bool = False,
    charset: Optional[str] = None,
    verify_integrity: bool = True
) -> Union[str, bytes]:
    if not isinstance(encoded_text, str):
        raise TypeError("Encoded text must be a string")
    if not isinstance(key, str):
        raise TypeError("Key must be a string")
    if not encoded_text:
        return b"" if binary else ""

    cs, base = _get_charset(charset)

    parts = encoded_text.split(".")
    
    integrity_tag = None
    if len(parts) > 1 and parts[0].startswith('i') and len(parts[0]) > 1:
        integrity_tag = parts[0][1:1 + INTEGRITY_LENGTH]
        parts = parts[1:]

    salt = None
    if len(parts) > 1:
        first_part = parts[0]
        if all(c in HEXDIGITS for c in first_part) and 8 <= len(first_part) <= 32:
            salt = first_part.lower()
            parts = parts[1:]

    keystream = _generate_keystream(key, len(parts), salt)
    decoded_chars = []

    for i, part in enumerate(parts):
        number = _from_custom_base(part, cs, base)
        original = number - keystream[i] - (i * 7)
        if original < 0 or original > 1114111:
            raise ValueError(f"Invalid decoded character at position {i}: wrong key or corrupted data")
        decoded_chars.append(chr(original))

    result = "".join(decoded_chars)

    if integrity_tag and verify_integrity:
        if not _verify_integrity(result, key, integrity_tag, cs, base):
            raise ValueError("Integrity check failed: wrong key or corrupted data")

    if binary:
        return base64.b64decode(result)

    return result


def encode_batch(
    texts: List[Union[str, bytes]],
    key: str,
    salt: Union[bool, str] = False,
    integrity: bool = False,
    binary: bool = False,
    charset: Optional[str] = None
) -> List[str]:
    return [
        encode(text, key, salt=salt, integrity=integrity, binary=binary, charset=charset)
        for text in texts
    ]


def decode_batch(
    encoded_texts: List[str],
    key: str,
    binary: bool = False,
    charset: Optional[str] = None,
    verify_integrity: bool = True
) -> List[Union[str, bytes]]:
    return [
        decode(text, key, binary=binary, charset=charset, verify_integrity=verify_integrity)
        for text in encoded_texts
    ]
