#!/usr/bin/env python3
"""
Base64 tokenizer for zero-copy ingestion of structured data.

Use when "prompts" are control tables, morph maps, receipts, or seeds
(already base64-encoded) - avoids UTF-8 decode/re-escape overhead.

Alphabet: A-Za-z0-9+/ (64 symbols) + = (pad)
"""
import base64
from typing import List, Tuple

class Base64Tokenizer:
    """
    Tokenizer for base64-encoded data.

    Treats bytes as 64-symbol alphabet + pad, chunks at N symbols per token.
    Reports bytes/token and enables zero-copy ingress for structured guidance.
    """

    # Base64 alphabet (RFC 4648)
    ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
    PAD = "="

    def __init__(self, symbols_per_token: int = 4):
        """
        Initialize base64 tokenizer.

        Args:
            symbols_per_token: Number of base64 symbols per token (default: 4)
                - 4 symbols = 3 bytes per token (standard base64 quantum)
                - 8 symbols = 6 bytes per token
        """
        self.symbols_per_token = symbols_per_token
        self.bytes_per_token = (symbols_per_token * 3) // 4  # Base64: 4 chars = 3 bytes

        # Build vocab: 64 symbols + pad + special tokens
        self.vocab = {char: idx for idx, char in enumerate(self.ALPHABET + self.PAD)}
        self.vocab_size = len(self.vocab)

    def encode(self, data: bytes) -> List[int]:
        """
        Encode bytes to token IDs.

        Args:
            data: Raw bytes

        Returns:
            List of token IDs
        """
        # Convert bytes to base64 string
        b64_str = base64.b64encode(data).decode('ascii')

        # Chunk into tokens
        tokens = []
        for i in range(0, len(b64_str), self.symbols_per_token):
            chunk = b64_str[i:i+self.symbols_per_token]
            # Convert chunk to token ID (simple: use first char's ID as token)
            # In real implementation, this would be a learned embedding
            token_id = self.vocab.get(chunk[0], 0)
            tokens.append(token_id)

        return tokens

    def decode(self, token_ids: List[int]) -> bytes:
        """
        Decode token IDs to bytes.

        Args:
            token_ids: List of token IDs

        Returns:
            Decoded bytes
        """
        # In real implementation, this would use learned embeddings
        # For now, placeholder that returns empty
        return b""

    def get_stats(self, data: bytes) -> dict:
        """
        Get tokenization statistics.

        Args:
            data: Input bytes

        Returns:
            Stats dict with symbols_per_token, bytes_per_token, etc.
        """
        b64_str = base64.b64encode(data).decode('ascii')
        num_tokens = (len(b64_str) + self.symbols_per_token - 1) // self.symbols_per_token

        return {
            "input_bytes": len(data),
            "b64_symbols": len(b64_str),
            "num_tokens": num_tokens,
            "symbols_per_token": self.symbols_per_token,
            "bytes_per_token": len(data) / num_tokens if num_tokens > 0 else 0,
            "vocab_size": self.vocab_size,
            "zero_copy": True  # No UTF-8 decode/re-escape
        }

def test_b64_tokenizer():
    """Test base64 tokenizer."""
    tokenizer = Base64Tokenizer(symbols_per_token=4)

    # Test data
    data = b"Hello, World! This is a test of the base64 tokenizer."

    # Encode
    tokens = tokenizer.encode(data)
    stats = tokenizer.get_stats(data)

    print(f"Input: {len(data)} bytes")
    print(f"Base64: {stats['b64_symbols']} symbols")
    print(f"Tokens: {stats['num_tokens']}")
    print(f"Bytes/token: {stats['bytes_per_token']:.2f}")
    print(f"Zero-copy: {stats['zero_copy']}")

if __name__ == "__main__":
    test_b64_tokenizer()
