# collatz_encoder.py

from mathstream import (
    StreamNumber,
    add,
    mul,
    div,
    is_even,
    clear_logs,
)

from pathlib import Path
import string

LOG_DIR = Path("instance/log")

# --- CONFIG ---
CHUNK_ENCODING = "ascii"  # Only support ASCII for now
BLOCK_CHARS = 3           # Number of characters to encode per step

# --- Helpers ---
def encode_message_to_int_blocks(msg: str) -> list[int]:
    padded = msg + ("\0" * (BLOCK_CHARS - len(msg) % BLOCK_CHARS))  # pad with nulls
    blocks = [padded[i:i+BLOCK_CHARS] for i in range(0, len(padded), BLOCK_CHARS)]
    return [int.from_bytes(block.encode(CHUNK_ENCODING), "big") for block in blocks]

def decode_blocks_to_message(values: list[int]) -> str:
    chars = []
    for val in values:
        try:
            b = val.to_bytes(BLOCK_CHARS, "big")
            chars.append(b.decode(CHUNK_ENCODING, errors="ignore"))
        except Exception:
            continue
    return "".join(chars).rstrip("\0")

def collatz_step(n: StreamNumber, three: StreamNumber, two: StreamNumber, one: StreamNumber) -> StreamNumber:
    return div(n, two) if is_even(n) else add(mul(n, three), one)

# --- Collatz forward walker ---
def walk_collatz(n: StreamNumber, max_steps=128):
    """Run a forward Collatz walk and return integer values from each step."""
    one = StreamNumber(literal="1")
    two = StreamNumber(literal="2")
    three = StreamNumber(literal="3")

    steps = []
    for _ in range(max_steps):
        digits = "".join(n.stream())
        try:
            val = int(digits)
        except Exception as e:
            print("[!] Could not convert to int:", digits)
            break
        steps.append(val)
        if val == 1:
            break
        n = collatz_step(n, three, two, one)
    return steps

# --- Reverse constructor prototype ---
def reverse_engineer_seed(target_blocks: list[int]):
    # Just takes the first block as the seed for now
    seed = StreamNumber(literal=str(target_blocks[0]))
    return seed

# --- MAIN ---
def encode_message_to_collatz_seed(message: str):
    clear_logs()
    encoded_blocks = encode_message_to_int_blocks(message)
    print("[+] Encoded message blocks:", encoded_blocks)
    seed = reverse_engineer_seed(encoded_blocks)
    return seed

def decode_collatz_path_to_message(seed: StreamNumber, max_steps=128):
    steps = walk_collatz(seed, max_steps)
    filtered = [v for v in steps if v.bit_length() <= 8 * BLOCK_CHARS]
    return decode_blocks_to_message(filtered)

# --- CLI TEST ---
if __name__ == "__main__":
    msg = "hi there"
    print("[+] Original message:", msg)
    seed = encode_message_to_collatz_seed(msg)
    print("[+] Seed value:", "".join(seed.stream()))
    decoded = decode_collatz_path_to_message(seed)
    print("[+] Decoded:", decoded)