import re
import json

def extract_last_code_block(text: str):
    code_blocks = re.findall(r'```.*?\n(.*?)```', text, re.DOTALL)
    if not code_blocks:
        code_blocks = re.findall(r'```(.*?)```', text, re.DOTALL)
    return code_blocks[-1].strip() if code_blocks else None

def calculate_switches(car_order, colors):
    """
    Calculate the number of color switches for a given car order
    
    Args:
    - car_order: List of car IDs
    - colors: List of colors corresponding to each car
    
    Returns:
    - Number of color switches
    """
    if not car_order:
        return 0
    
    switches = 0
    for i in range(1, len(car_order)):
        current_car = car_order[i]
        prev_car = car_order[i-1]
        # Car IDs start from 1, but list indices start from 0, so subtract 1
        if colors[current_car-1] != colors[prev_car-1]:
            switches += 1
    
    return switches

def is_valid_car_order(car_order, original_car_ids, K):
    """
    Check if the car order is valid
    
    Args:
    - car_order: Model's output car order
    - original_car_ids: Original car order
    - K: Maximum allowed movement range
    
    Returns:
    - Boolean indicating validity and optional error message
    """
    # Check if all cars are included and no duplicates
    if sorted(car_order) != sorted(original_car_ids):
        return False, "Car list contains duplicates or missing cars"
    
    # Check if each car's position is within allowed range
    for new_pos, car_id in enumerate(car_order):
        # Car IDs start from 1, positions also start from 1
        new_pos += 1
        # Get car's position in original order
        original_pos = original_car_ids.index(car_id) + 1
        # Check if movement is within allowed range
        if abs(new_pos - original_pos) > K:
            return False, f"Car {car_id} moved from position {original_pos} to {new_pos}, which exceeds limit K={K}"
    
    return True, ""

def verify(pred, answer, meta):
    """
    Verify if the model's prediction is correct
    
    Args:
    - pred: Model's prediction (text)
    - answer: Correct answer (car order)
    - meta: Problem metadata
    
    Returns:
    - True/False indicating verification result, or None if prediction cannot be parsed
    """
    # Parse metadata
    if isinstance(meta, str):
        try:
            meta = json.loads(meta)
        except json.JSONDecodeError:
            return 0
    elif isinstance(meta, dict):
        pass
    else:
        raise ValueError('meta should be dict or str')
    
    # Get original car order, colors and constraints
    original_car_ids = meta.get("car_ids", [])
    colors = meta.get("colors", [])
    K = meta.get("K", 0)
    min_switches = meta.get("min_switches", 0)
    
    # Extract JSON code block from prediction
    car_order = None
    final_answer = extract_last_code_block(pred)
    if not final_answer:
        return 0

    try:
        # Look for pattern like [1, 2, 3, 4, 5]
        pattern = r'\[([\d\s,]+)\]'
        matches = re.findall(pattern, final_answer)
        if matches:
            car_order = [int(x.strip()) for x in matches[0].split(',') if x.strip()]
        else:
            return 0
    except Exception:
        return 0
    
    if not car_order:
        return 0
    
    # Check validity of car order
    valid, error_msg = is_valid_car_order(car_order, original_car_ids, K)
    if not valid:
        return 0
    
    # Calculate color switches and compare with optimal solution
    switches = calculate_switches(car_order, colors)
    if switches == min_switches:
        return 1
    return 0
