from tea3.pretty_print import pretty_print
from tea3.cliutils import prompt_int, prompt_choice, prompt_list
from tea3.tea3model import Tea3Model
from tea3.variable_search import run_exhaustive
from tea3.sbox import run_sbox
from tea3.variable_xor import run_variable_xor, run_exhaustive_xor
from tea3.f31f32 import run_f31f32


def run_classic_cli():
    print("\nR registers are indexed 0–7; bits within each register are 0–7.")
    print("Enter -1 to print all registers, or all bits.")

    steps = prompt_int("How many steps do you want to run? (1–100): ", 1, 100)
    reg = prompt_int("Which R register do you want to inspect? (-1 or 0–7): ", -1, 7)
    bit = prompt_int("Which bit of that register?              (-1 or 0–7): ", -1, 7)

    print("-" * 50)

    model = Tea3Model()

    for i in range(steps):
        model.step()
        print(f"\n[Step {i + 1}]")

        regs = range(8) if reg == -1 else [reg]
        bits = range(8) if bit == -1 else [bit]

        for r in regs:
            for b in bits:
                poly = model.R_bits[r][b]
                print(f"R_bits[{r}][{b}] =")
                print(pretty_print(poly))
                print()

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

def run_exhaustive_cli():
    print("\nR registers are indexed 0–7; bits within each register are 0–7.")
    print("Enter -1 to print all bits in the chosen register.")

    steps = prompt_int("How many steps? (1–100): ", 1, 100)
    target_reg = prompt_int("Target register (0–7): ", 0, 7)
    target_bit = prompt_int("Target bit (-1 or 0–7): ", -1, 7)

    print("-" * 50)
    run_exhaustive(steps, target_reg, target_bit)

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


def run_variable_xor_cli():
    print("\nR registers are indexed 0–7; bits within each register are 0–7.")
    print("Enter the bit positions to XOR, e.g. `0 1`.")

    steps = prompt_int("How many steps? (1–100): ", 1, 100)
    target_reg = prompt_int("Target register (0–7): ", 0, 7)
    bits_to_xor = prompt_list(
        "Bits to XOR (0–7, separated by spaces, ranges with '-'): ",
        item_name="bit",
        min_value=0,
        max_value=7,
    )

    print("-" * 50)
    run_variable_xor(steps, target_reg, bits_to_xor)

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

def run_exhaustive_xor_cli():
    print("\nR registers are indexed 0–7; bits within each register are 0–7.")
    print("This mode searches all XOR combinations of a given size.")

    steps = prompt_int("How many steps? (1–100): ", 1, 100)
    target_reg = prompt_int("Target register (0–7): ", 0, 7)
    xor_size = prompt_int("Number of bits to XOR (2–8): ", 2, 8)

    print("-" * 50)

    run_exhaustive_xor(steps=steps, target_reg=target_reg, xor_size=xor_size)

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



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

    print("\nChoose a mode:")
    print("  1) Classic inspection")
    print("  2) Exhaustive variable-change search")
    print("  3) S box analysis")
    print("  4) variable XOR")
    print("  5) F31, F32 analysis")
    print("  6) Exhaustive XOR search")

    mode = prompt_choice("Your choice (1, 2, 3, 4, 5 or 6): ", {1, 2, 3, 4, 5, 6})

    if mode == 1:
        run_classic_cli()
    elif mode == 2:
        run_exhaustive_cli()
    elif mode == 3:
        run_sbox()
    elif mode == 4:
        run_variable_xor_cli()
    elif mode == 5:
        run_f31f32()
    elif mode == 6:
        run_exhaustive_xor_cli()


main()