from tea3.pretty_print import pretty_print
from tea3.tea3model import Tea3Model


def prompt_int(message: str, lo: int, hi: int) -> int:
    while True:
        raw = input(message).strip()
        try:
            value = int(raw)
        except ValueError:
            print("Please enter a whole number.")
            continue
        if lo <= value <= hi:
            return value
        print(f"Value must be between {lo} and {hi} (inclusive).")


def main() -> None:
    print("=" * 50)
    print(" Tea3 Model ")
    print("=" * 50)

    steps = prompt_int("\nHow many steps do you want to run? (1–100): ", 1, 100)

    print("\nR registers are indexed 0–7; bits within each register are 0–7.")
    reg = prompt_int("Which R register do you want to inspect? (0–7): ", 0, 7)
    bit = prompt_int("Which bit of that register?           (0–7): ", 0, 7)

    print(f"\nRunning {steps} step(s), watching R[{reg}][{bit}]")
    print("-" * 50)

    model = Tea3Model(max_steps=steps)

    for i in range(steps):
        model.step()
        poly = model.R_bits[reg][bit]
        print(f"\n[Step {i+1}]  R_bits[{reg}][{bit}] =")
        print(pretty_print(poly))

    print("\n" + "=" * 50)
    print("Done.")


main()