"""
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")
"""

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

__all__ = ["encode", "decode"]

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


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_salt() -> str:
    return secrets.token_hex(SALT_LENGTH)


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

    if num == 0:
        return CUSTOM_CHARSET[0]

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


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

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


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

    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))

    result = ".".join(encoded_parts)
    return f"{salt_value}.{result}" if salt_value else result


def decode(encoded_text: str, key: str) -> str:
    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 ""

    parts = encoded_text.split(".")
    
    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 = parts[0].lower()
            parts = parts[1:]

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

    for i, part in enumerate(parts):
        number = _from_custom_base(part)
        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))

    return "".join(decoded_chars)
