import ast
import re

MIN_DISK_ID = 1
MAX_DISK_ID = 5
DISK_IDS = tuple(list(range(MIN_DISK_ID, MAX_DISK_ID + 1)))
NUM_POSTS = 3

CRE_NEXT_STATE = re.compile(r"(?i)\bnext_state\b\s*=\s*(\[\s*\[)", re.DOTALL)
CRE_MOVE = re.compile(r"(?i)\bmove\b\s*=\s*(\[[^\[\]]*\])")


SYSTEM_PROMPT = """
You are a helpful assistant. Solve this puzzle for me.

There are three pegs and n disks of different sizes stacked on the first peg. The disks are
numbered from 1 (smallest) to n (largest). Disk moves in this puzzle should follow:
1. Only one disk can be moved at a time.
2. Each move consists of taking the upper disk from one stack and placing it on top of
another stack.
3. A larger disk may not be placed on top of a smaller disk.
The goal is to move the entire stack to the third peg.

Example: With 3 disks numbered 1 (smallest), 2, and 3 (largest), the initial state is [[3, 2,
1], [], []], and a solution might be:
moves = [[1, 0, 2], [2, 0, 1], [1, 2, 1], [3, 0, 2], [1, 1, 0], [2, 1, 2], [1, 0, 2]]
This means: Move disk 1 from peg 0 to peg 2, then move disk 2 from peg 0 to peg 1, and so on.

Requirements:
- The positions are 0-indexed (the leftmost peg is 0).
- Ensure your answer includes a single next move in this EXACT FORMAT:
'''move = [disk id, from peg, to peg]'''
- Ensure your answer includes the next state resulting from applying the move to the current
state in this EXACT FORMAT:
'''next_state = [[...], [...], [...]]'''
"""

USER_TEMPLATE = """
Rules:
- Only one disk can be moved at a time.
- Only the top disk from any stack can be moved.
- A larger disk may not be placed on top of a smaller disk.
For all moves, follow the standard Tower of Hanoi procedure:
If the previous move did not move disk 1, move disk 1 clockwise one peg (0 -> 1 -> 2 -> 0).
If the previous move did move disk 1, make the only legal move that does not involve moving
disk1.
Use these clear steps to find the next move given the previous move and current state
"""


def extract_balanced_brackets(text, start_idx):
    """Extract a substring with balanced brackets [[...]] starting at start_idx"""
    bracket_stack = []
    i = start_idx
    while i < len(text):
        if text[i] == '[':
            bracket_stack.append('[')
        elif text[i] == ']':
            if not bracket_stack:
                break
            bracket_stack.pop()
        if not bracket_stack:
            return text[start_idx:i + 1]
        i += 1
    return text[start_idx:i] + ']'


def parse_move_state_repair(response_text):
    try:
        move_matches = CRE_MOVE.findall(response_text)
        if not move_matches:
            raise ValueError("No 'move = [...]' found in response.")
        # last mention of 'move = [...]'
        move = ast.literal_eval(move_matches[-1].strip())
    except Exception as e:
        raise ValueError(f"Could find 'move = [...]' in response ") from e
    try:
        # normalize whitespace and look for 'next_state = [ ['
        matches = list(CRE_NEXT_STATE.finditer(response_text))
        if not matches:
            raise ValueError("'next_state = [ [' not found in response.")
        # take last match
        start_idx = matches[-1].start(1)
        next_state_str = extract_balanced_brackets(response_text, start_idx).strip()
        next_state = ast.literal_eval(next_state_str)
    except Exception as e:
        raise ValueError(f"No 'next_state = [[' in response {e}.") from e
    return move, next_state


def _validate_move(move, disk_ids=DISK_IDS):
    """ move variable must be list of 3 integers """
    assert isinstance(move, list), f"isinstance({repr(move)}, list)"
    assert len(move) == 3, f"len({move}) == 3"
    assert all(isinstance(x, int) for x in move), f"all(isinstance(x, int) for x in {move})"
    assert all((x in disk_ids) for x in move), f"all((x in disk_ids) for x in {move})"
    return move
    # if (not isinstance(move, list)
    #         or len(move) != 3
    #         or not all(isinstance(x, int) for x in move)
    #     ):
    #     raise ValueError("'move = [...]' must be a list of exactly 3 integers.")
