"""Reserve horizontal space for field calculations beside their geometry consumers."""

from collections import defaultdict

from .zones import node_zones, repeat_regions


def clear_repeat_regions(node_group, positions, sizes, bodies, *, gap):
    """Move unrelated calculation bodies and nodes outside reserved repeat areas."""
    regions = repeat_regions(node_group, positions, sizes)
    grouped = set().union(*(members for members, _ in bodies)) if bodies else set()
    groups = [members for members, _ in bodies] + [{node} for node in positions if node not in grouped]
    for members in groups:
        obstacles = [box for owned, box in regions if not members & owned]
        left = min(positions[node][0] for node in members)
        right = max(positions[node][0] + sizes[node].width for node in members)
        top = max(positions[node][1] for node in members)
        bottom = min(positions[node][1] - sizes[node].height for node in members)
        if not any(left < b and right > a and top > d and bottom < c for a, b, c, d in obstacles):
            continue
        occupied = obstacles + [(x, x + sizes[node].width, y, y - sizes[node].height)
                                for node, (x, y) in positions.items() if node not in members]
        occupied = [box for box in occupied if left < box[1] + gap and right > box[0] - gap]
        offsets = {edge for _, _, upper, lower in occupied
                   for edge in (upper + gap - bottom, lower - gap - top)}
        targets = {link.to_node for node in members for socket in node.outputs for link in socket.links
                   if link.to_node not in members and link.to_node in positions}
        above = (right - left < 1000 and len(targets) == 1
                 and bottom > positions[next(iter(targets))][1])
        offset = min((offset for offset in offsets
                      if all(top + offset <= lower - gap + 1e-5 or bottom + offset >= upper + gap - 1e-5
                             for _, _, upper, lower in occupied)),
                     key=lambda offset: (above and bottom + offset <= positions[next(iter(targets))][1], abs(offset)))
        for node in members:
            x, y = positions[node]
            positions[node] = (x, y + offset)
    for index, (members, _) in enumerate(bodies):
        bodies[index] = (members, (min(positions[node][0] for node in members),
                                  max(positions[node][0] + sizes[node].width for node in members),
                                  max(positions[node][1] for node in members),
                                  min(positions[node][1] - sizes[node].height for node in members)))


def enclose_repeat_branches(node_group, positions, sizes, bodies, *, gap):
    """Leave real horizontal space between repeat boundaries and their calculations."""
    for entry, end, members in node_zones(node_group):
        if entry.bl_idname != "GeometryNodeRepeatInput":
            continue
        interior = (members - {entry, end}) & positions.keys()
        if not interior:
            continue
        for boundary, direction in ((entry, -1), (end, 1)):
            start = positions[boundary][0]
            target = (min(positions[node][0] for node in interior) - gap - sizes[entry].width
                      if direction < 0 else
                      max(positions[node][0] + sizes[node].width for node in interior) + gap)
            distance = direction * (target - start)
            if distance <= 0:
                continue
            moving = {node for node, (x, _) in positions.items()
                      if node not in members and direction * (x - start) >= 0} | {boundary}
            for body, _ in bodies:
                if body & moving and not body & members:
                    moving.update(body)
            for node in moving:
                x, y = positions[node]
                positions[node] = (x + direction * distance, y)
    for index, (members, _) in enumerate(bodies):
        bodies[index] = (members, (min(positions[node][0] for node in members),
                                  max(positions[node][0] + sizes[node].width for node in members),
                                  max(positions[node][1] for node in members),
                                  min(positions[node][1] - sizes[node].height for node in members)))


def reserve_field_sections(node_group, layouts, positions, sizes, *, gap):
    """Extend execution stages to make room for substantial incoming calculations."""
    stages = {node: node for node in positions}
    for entry, _, members in node_zones(node_group):
        for node in members & positions.keys():
            stages[node] = entry
    sections = defaultdict(list)
    for region, owners, local, width, height in layouts:
        if width < 1000 or not owners:
            continue
        owner = min(owners, key=lambda node: (positions[node][0], node.name))
        sections[stages[owner]].append((region, owner, local, width, height))
    front = max((width for _, _, _, width, _ in layouts), default=0)
    bands = [min(x for x, _ in positions.values()) - front - gap] * 2
    for stage in sorted(sections, key=lambda node: (positions[node][0], node.name)):
        start = positions[stage][0]
        pending = []
        for region, owner, local, width, height in sorted(sections[stage], key=lambda item: -item[3]):
            band = min(range(2), key=lambda index: bands[index])
            left = bands[band] + gap
            for node in region:
                for socket in node.inputs:
                    for link in socket.links:
                        source = link.from_node
                        if source in positions and stages[source] != stage:
                            left = max(left, positions[source][0] + sizes[source].width + gap - local[node][0])
            offsets = [local[node][1] - sizes[node].sockets[socket]
                       + sizes[link.to_node].sockets[link.to_socket]
                       + positions[owner][1] - positions[link.to_node][1]
                       for node in region for socket in node.outputs for link in socket.links
                       if link.to_node in positions]
            above = max((abs(height + gap + offset) for offset in offsets), default=gap)
            below = max((abs(-sizes[owner].height - gap + offset) for offset in offsets), default=gap)
            exit_run = max(gap, min(above, below) * 0.8)
            finish = max(positions[owner][0] - exit_run, left + width)
            pending.append((owner, finish + exit_run))
            bands[band] = finish
        extra = max(finish - positions[owner][0] for owner, finish in pending)
        if extra > 0:
            for node, (x, y) in positions.items():
                if x >= start:
                    positions[node] = (x + extra, y)