import heapq
import re

class Node:
    def __init__(self, node_id, node_type, name, arm_num, take_time, dependencies, location, target_obj=None, source_obj=None,delay_after=0):
        self.id = node_id
        self.type = node_type 
        self.name = name
        self.arm_num = arm_num 
        self.take_time = take_time
        self.dependencies = dependencies
        self.predecessors = len(dependencies)
        self.earliest_start = 0 
        self.successors = [] 
        self.location = location 
        self.target_obj = target_obj
        self.source_obj = source_obj
        self.delay_after = delay_after 
        
def choose_arm(arm_free, node, arm_locks=None):
    """Selects the appropriate arm(s) for a given node"""
    if arm_locks is None:
        arm_locks = {
            'left': {'locked': False, 'current_chain': None},
            'right': {'locked': False, 'current_chain': None}
        }
    left_locked = arm_locks['left']['locked']
    right_locked = arm_locks['right']['locked']
    
    if node.type == 'place':
        for arm in ['left', 'right']:
            if arm_locks[arm]['current_chain'] == node.source_obj:
                if arm_free[arm] <= node.earliest_start:
                    return arm
        return None

    # Case 0: both arms are locked
    if arm_locks["left"]['locked'] and arm_locks["right"]['locked']:
        current_chain = None
        if node.type == 'pick':
            if match := re.search(r'target="([^"]+)"', node.name):
                current_chain = f"{match.group(1)}"
        elif node.type == 'place' or node.type == 'stick_on' or node.type == 'cut' or node.type == 'pour_into':#这里是双臂动作添加的地方，添加技能需要修改参数
            if match := re.search(r'source="([^"]+)"', node.name):
                current_chain = f"{match.group(1)}"

        left_match = arm_locks['left']['current_chain'] == current_chain
        right_match = arm_locks['right']['current_chain'] == current_chain

        if node.arm_num == 2:
            if left_match and right_match and current_chain:
                return 'both'
            if left_match and current_chain and not arm_locks['right']['current_chain']:
                return 'both'
            if right_match and current_chain and not arm_locks['left']['current_chain']:
                return 'both'
        else:
            if left_match and not right_match:
                return 'left' if arm_free['left'] <= node.earliest_start else None
            if right_match and not left_match:
                return 'right' if arm_free['right'] <= node.earliest_start else None
            if left_match and right_match:
                return 'left' if arm_free['left'] < arm_free['right'] else 'right'
        return None

    # Case 1: the task requires both arms
    if node.arm_num == 2:
        return 'both'

    # Case 2: one arm is locked or both are free
    locked_arm = None
    if left_locked and not right_locked:
        locked_arm = 'left'
        free_arm = 'right'
    elif right_locked and not left_locked:
        locked_arm = 'right'
        free_arm = 'left'

    if locked_arm:
        if node.type == 'place':
            if arm_locks[locked_arm]['current_chain'] == node.source_obj:
                if arm_free[locked_arm] <= node.earliest_start:
                    return locked_arm
                else:
                    return None
            else:
                return None
        else:
            source_chain = node.source_obj
            if arm_locks[locked_arm]['current_chain'] == source_chain:
                if arm_free[locked_arm] <= node.earliest_start:
                    return locked_arm
                else:
                    return None
            else:
                if arm_free[free_arm] <= node.earliest_start:
                    return free_arm
                else:
                    return None

    left_arm_free_time = arm_free['left']
    right_arm_free_time = arm_free['right']

    if left_arm_free_time < right_arm_free_time:
        return 'left'
    elif right_arm_free_time < left_arm_free_time:
        return 'right'
    else:
        return 'left'


def print_arm(name, tasks):
    """Prints the task schedule for a single arm"""
    
    print(f"{name}")
    for task_info in tasks:
        start, end, task_name = task_info[:3]
        print(f"{start}-{end} {task_name}")
    print()

def get_successors(node_id, nodes):
    """Returns the list of successor node IDs for a given node ID"""
    
    if node_id in nodes:
        return nodes[node_id].successors
    return []

def schedule(nodes):
    """Main scheduling algorithm using a priority queue to simulate task execution over time"""
    
    event_queue = []
    arm_free = {'left': 0, 'right': 0}
    schedule_log = {'left': [], 'right': []}
    completed = {}

    arm_locks = {
        'left': {'locked': False, 'current_chain': None},
        'right': {'locked': False, 'current_chain': None}
    }

    for node in nodes.values():
        if node.predecessors == 0:
            heapq.heappush(event_queue, (node.earliest_start, 0, 'available', node.id))

    while event_queue:
        current_time, _, event_type, *payload = heapq.heappop(event_queue)
        if event_type == 'available':
            node_id = payload[0]
            node = nodes[node_id]
            chosen_arm = choose_arm(arm_free, node, arm_locks)
            len_left = len(schedule_log['left'])
            len_right = len(schedule_log['right'])
            pick_left = 0
            pick_right = 0
            other_arm = None
            now_arm = None
            other_arm_successors = []
            if  len_left > 0 and schedule_log['left'][-1][2].startswith("pick"):
                pick_left = 1
            if  len_right > 0 and schedule_log['right'][-1][2].startswith("pick"):
                pick_right = 1
            if not chosen_arm and pick_left and pick_right and node.arm_num == 2:
                if node.source_obj==arm_locks['left']['current_chain']:
                    now_arm ='left'
                    other_arm ='right'
                if node.source_obj==arm_locks['right']['current_chain']:
                    other_arm ='left'
                    now_arm = 'right'
                if schedule_log[other_arm]:
                    pick_id=schedule_log[other_arm][-1][3]
                    del schedule_log['right'][-1]
                    other_arm_successors=get_successors(pick_id,nodes)
                    other_arm_successors.append(pick_id)
                    event_queue = [item for item in event_queue if item[3] not in other_arm_successors]
                    heapq.heapify(event_queue)
                    start = arm_free[now_arm]
                    end = start + node.take_time
                    arm_free.update({'left': end, 'right': end})
                    schedule_log['left'].append((start, end, node.name, node.id))
                    schedule_log['right'].append((start, end, node.name, node.id))
                    if (now_arm=='left'):
                        arm_locks['left'].update({'locked': True, 'current_chain': node.source_obj})
                        arm_locks['right'].update({'locked': True, 'current_chain': None})
                    if (now_arm=='right'):
                        arm_locks['left'].update({'locked': True, 'current_chain': None})
                        arm_locks['right'].update({'locked': True, 'current_chain': node.source_obj})

                    heapq.heappush(event_queue, (end, 1, 'completed', node_id, 'both'))
                    heapq.heappush(event_queue, (end, 1, 'available', pick_id))
                    continue

            if not chosen_arm:
                heapq.heappush(event_queue, (current_time + 1, 1, 'available', node.id))
                continue

            if node.type == 'place':
                if chosen_arm != 'both' and arm_locks[chosen_arm]['current_chain'] != node.source_obj:
                    heapq.heappush(event_queue, (current_time + 1, 1, 'available', node.id))
                    continue

            if node.type == 'pick':
                if chosen_arm != 'both' and arm_locks[chosen_arm]['locked']:
                    heapq.heappush(event_queue, (current_time + 1, 1, 'available', node.id))
                    continue
                if chosen_arm != 'both':
                    arm_locks[chosen_arm] = {
                        'locked': True,
                        'current_chain': node.target_obj
                    }

            if chosen_arm == 'both':
                source_obj = node.source_obj
                left_match = arm_locks['left']['current_chain'] == source_obj
                right_match = arm_locks['right']['current_chain'] == source_obj
                both_unlocked = not arm_locks['left']['locked'] and not arm_locks['right']['locked']

                if (left_match or right_match) or both_unlocked:
                    start = max(arm_free.values())
                    end = start + node.take_time
                    arm_free.update({'left': end, 'right': end})
                    schedule_log['left'].append((start, end, node.name, node.id))
                    schedule_log['right'].append((start, end, node.name, node.id))
                    
                    if(left_match):
                        arm_locks['left'].update({'locked': True, 'current_chain': source_obj})
                        arm_locks['right'].update({'locked': True, 'current_chain': None})
                    if (right_match):
                        arm_locks['left'].update({'locked': True, 'current_chain': None})
                        arm_locks['right'].update({'locked': True, 'current_chain': source_obj})
                    if both_unlocked:
                        arm_locks['left'].update({'locked': True, 'current_chain': source_obj})
                        arm_locks['right'].update({'locked': True, 'current_chain': source_obj})
                    heapq.heappush(event_queue, (end, 1, 'completed', node_id, chosen_arm))
                else:
                    heapq.heappush(event_queue, (current_time + 1, 1, 'available', node.id))

            else:
                start = max(arm_free[chosen_arm], node.earliest_start)
                end = start + node.take_time
                arm_free[chosen_arm] = end
                schedule_log[chosen_arm].append((start, end, node.name, node.id))
                heapq.heappush(event_queue, (end, 1, 'completed', node_id, chosen_arm))

            # print_arm("left:", schedule_log['left'])
            # print_arm("right:", schedule_log['right'])
            # print(" ")
            
        elif event_type == 'completed':
            node_id, used_arm = payload
            node = nodes[node_id]
            completed[node_id] = current_time

            if used_arm == 'both':
                if arm_locks['left']['current_chain']==node.source_obj:
                    arm_locks['right']['locked'] = False
                    arm_locks['right']['current_chain'] = None
                if arm_locks['right']['current_chain']==node.source_obj:
                    arm_locks['left']['locked'] = False
                    arm_locks['left']['current_chain'] = None

            elif used_arm in ['left', 'right']:
                if node.type == 'place'and arm_locks[used_arm]['current_chain'] == node.source_obj:
                    arm_locks[used_arm]['locked'] = False
                    arm_locks[used_arm]['current_chain'] = None
                elif arm_locks[used_arm]['current_chain'] is None:
                    arm_locks[used_arm]['locked'] = False
            for succ_id in node.successors:
                succ = nodes[succ_id]
                new_earliest = current_time + succ.delay_after
                if new_earliest > succ.earliest_start:
                    succ.earliest_start = new_earliest

                succ.predecessors -= 1
                if succ.predecessors == 0:
                    heapq.heappush(event_queue, (succ.earliest_start, 0, 'available', succ.id))

    def format_log(log):
        return sorted([(s, e, n) for s, e, n, _ in log], key=lambda x: x[0])

    return {
        'left': format_log(schedule_log['left']),
        'right': format_log(schedule_log['right']),
        'total': max(completed.values(), default=0)
    }