"""Mirror completed depth geometry and fill matching boundaries."""
import bpy
from mathutils import Matrix

from ..common.nodes import (
    interface_socket, group_input, float_input, bool_input, vector_input,
    evaluate_field, store_float_attribute, read_float_attribute, prepare_node_group,
)
from ..common.depth_surface import _math
from ..common.normal_map import ROTATION_ATTRIBUTE, FACE_ATTRIBUTE, AXIS_ATTRIBUTE
from ..common.direction import symmetry_legacy_direction_to_canonical
from ..common.smoothing import edge_boundary_field

NODE_GROUP_NAME = "O Image Cutout Symmetry"


def _align_front(group, geometry, controls):
    """Align existing geometry with the fixed XY symmetry plane."""
    nodes, links = group.nodes, group.links
    direction = symmetry_legacy_direction_to_canonical(group, controls['Direction'])
    align = nodes.new('FunctionNodeAlignEulerToVector')
    align.axis = 'Y'
    links.new(direction, align.inputs['Vector'])
    position = nodes.new('GeometryNodeInputPosition')
    rotate = nodes.new('ShaderNodeVectorRotate')
    rotate.rotation_type = 'EULER_XYZ'
    rotate.invert = True
    links.new(position.outputs[0], rotate.inputs['Vector'])
    links.new(align.outputs['Rotation'], rotate.inputs['Rotation'])
    offset = nodes.new('ShaderNodeCombineXYZ')
    offset_value = nodes.new('ShaderNodeMath')
    offset_value.operation = 'MULTIPLY'
    offset_value.inputs[1].default_value = -1.0
    links.new(controls['Offset'], offset_value.inputs[0])
    links.new(offset_value.outputs[0], offset.inputs['Y'])
    offset_position = nodes.new('ShaderNodeVectorMath')
    offset_position.operation = 'ADD'
    links.new(rotate.outputs[0], offset_position.inputs[0])
    links.new(offset.outputs[0], offset_position.inputs[1])
    scale_vector = nodes.new('ShaderNodeCombineXYZ')
    scale_vector.inputs['X'].default_value = 1
    scale_vector.inputs['Z'].default_value = 1
    links.new(controls['Scale'], scale_vector.inputs['Y'])
    scale = nodes.new('ShaderNodeVectorMath')
    scale.operation = 'MULTIPLY'
    links.new(offset_position.outputs[0], scale.inputs[0])
    links.new(scale_vector.outputs[0], scale.inputs[1])
    place = nodes.new('GeometryNodeSetPosition')
    links.new(geometry, place.inputs['Geometry'])
    links.new(scale.outputs[0], place.inputs['Position'])
    store = nodes.new('GeometryNodeStoreNamedAttribute')
    store.data_type, store.domain = 'FLOAT_VECTOR', 'FACE'
    store.inputs['Name'].default_value = ROTATION_ATTRIBUTE
    links.new(place.outputs[0], store.inputs['Geometry'])
    links.new(align.outputs['Rotation'], store.inputs['Value'])
    return store.outputs['Geometry']


def _boundary_edge_field(group):
    """Report one on edges that carry a single face."""
    neighbors = group.nodes.new('GeometryNodeInputMeshEdgeNeighbors')
    return _math(group, 'COMPARE', neighbors.outputs['Face Count'], 1)


def _fold_to_plane(group, geometry, tolerance):
    """Move over-plane points onto the plane and drop the faces that collapse."""
    nodes, links = group.nodes, group.links
    position = nodes.new('GeometryNodeInputPosition')
    components = nodes.new('ShaderNodeSeparateXYZ')
    links.new(position.outputs[0], components.inputs[0])
    beyond = _math(group, 'GREATER_THAN', components.outputs['Y'], tolerance)
    target = nodes.new('ShaderNodeCombineXYZ')
    links.new(components.outputs['X'], target.inputs['X'])
    links.new(components.outputs['Z'], target.inputs['Z'])
    place = nodes.new('GeometryNodeSetPosition')
    links.new(geometry, place.inputs['Geometry'])
    links.new(beyond, place.inputs['Selection'])
    links.new(target.outputs[0], place.inputs['Position'])
    # A face lying on the plane would be mirrored onto itself.
    center = evaluate_field(group, _math(group, 'ABSOLUTE', components.outputs['Y']), 'FLOAT', 'FACE')
    normal = nodes.new('GeometryNodeInputNormal')
    normal_axes = nodes.new('ShaderNodeSeparateXYZ')
    links.new(normal.outputs[0], normal_axes.inputs[0])
    flat = nodes.new('FunctionNodeBooleanMath')
    flat.operation = 'AND'
    links.new(_math(group, 'LESS_THAN', center, tolerance), flat.inputs[0])
    links.new(_math(group, 'GREATER_THAN', _math(group, 'ABSOLUTE', normal_axes.outputs['Y']), 1 - 1e-3), flat.inputs[1])
    delete = nodes.new('GeometryNodeDeleteGeometry')
    delete.domain, delete.mode = 'FACE', 'ALL'
    links.new(place.outputs[0], delete.inputs['Geometry'])
    links.new(flat.outputs[0], delete.inputs['Selection'])
    return _clean_crease(group, delete.outputs[0], tolerance)


def _clean_crease(group, geometry, tolerance):
    """Open interior in-plane edges left by folding a closed surface."""
    nodes, links = group.nodes, group.links
    heights = _edge_abs_y(group)
    edge_plane = nodes.new('FunctionNodeBooleanMath')
    edge_plane.operation = 'AND'
    links.new(_math(group, 'LESS_THAN', heights[0], tolerance), edge_plane.inputs[0])
    links.new(_math(group, 'LESS_THAN', heights[1], tolerance), edge_plane.inputs[1])
    neighbors = nodes.new('GeometryNodeInputMeshEdgeNeighbors')
    interior = _math(group, 'COMPARE', neighbors.outputs['Face Count'], 2)
    fold = nodes.new('FunctionNodeBooleanMath')
    fold.operation = 'AND'
    links.new(edge_plane.outputs[0], fold.inputs[0])
    links.new(interior, fold.inputs[1])

    face_index = nodes.new('GeometryNodeInputIndex')
    corners = nodes.new('GeometryNodeCornersOfFace')
    links.new(face_index.outputs[0], corners.inputs['Face Index'])
    face_fold = None
    for offset in range(4):
        corner = nodes.new('GeometryNodeOffsetCornerInFace')
        links.new(corners.outputs['Corner Index'], corner.inputs['Corner Index'])
        corner.inputs['Offset'].default_value = offset
        edge_of_corner = nodes.new('GeometryNodeEdgesOfCorner')
        links.new(corner.outputs['Corner Index'], edge_of_corner.inputs['Corner Index'])
        sample = nodes.new('GeometryNodeFieldAtIndex')
        sample.data_type, sample.domain = 'BOOLEAN', 'EDGE'
        links.new(fold.outputs[0], sample.inputs['Value'])
        links.new(edge_of_corner.outputs['Next Edge Index'], sample.inputs['Index'])
        if face_fold is None:
            face_fold = sample.outputs['Value']
        else:
            combined = nodes.new('FunctionNodeBooleanMath')
            combined.operation = 'OR'
            links.new(face_fold, combined.inputs[0])
            links.new(sample.outputs['Value'], combined.inputs[1])
            face_fold = combined.outputs['Boolean']
    cleanup = nodes.new('GeometryNodeDeleteGeometry')
    cleanup.domain, cleanup.mode = 'FACE', 'ALL'
    links.new(geometry, cleanup.inputs['Geometry'])
    links.new(face_fold, cleanup.inputs['Selection'])
    return cleanup.outputs[0]


def _off_plane_edge_field(group, tolerance):
    """Report one on edges that keep an endpoint off the symmetry plane."""
    heights = _edge_abs_y(group)
    return _math(group, 'GREATER_THAN', _math(group, 'MAXIMUM', heights[0], heights[1]), tolerance)


def _edge_abs_y(group):
    """Return both endpoint heights on the symmetry axis for an edge."""
    edge = group.nodes.new('GeometryNodeInputMeshEdgeVertices')
    heights = []
    for socket in ('Position 1', 'Position 2'):
        axis = group.nodes.new('ShaderNodeSeparateXYZ')
        group.links.new(edge.outputs[socket], axis.inputs[0])
        heights.append(_math(group, 'ABSOLUTE', axis.outputs['Y']))
    return heights


def _mirror(group, geometry):
    mirror = group.nodes.new('GeometryNodeTransform')
    mirror.inputs['Scale'].default_value = (1, -1, 1)
    group.links.new(geometry, mirror.inputs['Geometry'])
    flip = group.nodes.new('GeometryNodeFlipFaces')
    group.links.new(mirror.outputs[0], flip.inputs['Mesh'])
    return flip.outputs[0]


def _cut_fields(group, tolerance):
    """Return point-domain seam and outline fields for the cut."""
    nodes, links = group.nodes, group.links
    boundary = evaluate_field(group, edge_boundary_field(nodes, links), 'FLOAT', 'POINT')
    position = nodes.new('GeometryNodeInputPosition')
    components = nodes.new('ShaderNodeSeparateXYZ')
    links.new(position.outputs[0], components.inputs[0])
    absolute = _math(group, 'ABSOLUTE', components.outputs['Y'])

    def split(side):
        node = nodes.new('FunctionNodeBooleanMath')
        node.operation = 'AND'
        links.new(_math(group, 'GREATER_THAN', boundary, 0), node.inputs[0])
        links.new(side, node.inputs[1])
        return node.outputs[0]

    return split(_math(group, 'LESS_THAN', absolute, tolerance)), split(
        _math(group, 'GREATER_THAN', absolute, tolerance)
    )


def _falloff(group, center, rings=2):
    """Blur a 0/1 center into a falloff over the requested rings."""
    blur = group.nodes.new('GeometryNodeBlurAttribute')
    blur.data_type = 'FLOAT'
    blur.inputs['Iterations'].default_value = rings
    group.links.new(center, blur.inputs['Value'])
    return _math(group, 'MAXIMUM', center, blur.outputs[0])


def _apply_planar_relax(group, geometry, influence, smooth):
    """Move points inside the plane by their blurred position, gated by influence."""
    nodes, links = group.nodes, group.links
    position = nodes.new('GeometryNodeInputPosition')
    blur = nodes.new('GeometryNodeBlurAttribute')
    blur.data_type = 'FLOAT_VECTOR'
    links.new(position.outputs[0], blur.inputs['Value'])
    links.new(smooth, blur.inputs['Iterations'])
    difference = nodes.new('ShaderNodeVectorMath')
    difference.operation = 'SUBTRACT'
    links.new(blur.outputs[0], difference.inputs[0])
    links.new(position.outputs[0], difference.inputs[1])
    in_plane = nodes.new('ShaderNodeVectorMath')
    in_plane.operation = 'MULTIPLY'
    in_plane.inputs[1].default_value = (.5, 0, .5)
    links.new(difference.outputs[0], in_plane.inputs[0])
    scaled = nodes.new('ShaderNodeVectorMath')
    scaled.operation = 'SCALE'
    links.new(in_plane.outputs[0], scaled.inputs[0])
    links.new(influence, scaled.inputs[3])
    place = nodes.new('GeometryNodeSetPosition')
    links.new(geometry, place.inputs['Geometry'])
    links.new(scaled.outputs[0], place.inputs['Offset'])
    return place.outputs[0]


def _relax_seam(group, geometry, seam_field, smooth):
    """Relax the cut seam in plane through a two-ring falloff."""
    return _apply_planar_relax(group, geometry, _falloff(group, seam_field), smooth)


def _extrude_boundary(group, geometry, selection, y_factor):
    """Extrude selected edges to a fraction of their current height."""
    nodes, links = group.nodes, group.links
    extrude = nodes.new('GeometryNodeExtrudeMesh')
    extrude.mode = 'EDGES'
    extrude.inputs['Offset Scale'].default_value = 0
    links.new(geometry, extrude.inputs['Mesh'])
    links.new(selection, extrude.inputs['Selection'])
    position = nodes.new('GeometryNodeInputPosition')
    scale = nodes.new('ShaderNodeVectorMath')
    scale.operation = 'MULTIPLY'
    scale.inputs[1].default_value = (1, y_factor, 1)
    links.new(position.outputs[0], scale.inputs[0])
    place = nodes.new('GeometryNodeSetPosition')
    links.new(extrude.outputs['Mesh'], place.inputs['Geometry'])
    links.new(extrude.outputs['Top'], place.inputs['Selection'])
    links.new(scale.outputs[0], place.inputs['Position'])
    store = nodes.new('GeometryNodeStoreNamedAttribute')
    store.data_type, store.domain = 'FLOAT', 'FACE'
    store.inputs['Name'].default_value = FACE_ATTRIBUTE
    store.inputs['Value'].default_value = 3
    links.new(place.outputs[0], store.inputs['Geometry'])
    links.new(extrude.outputs['Side'], store.inputs['Selection'])
    return store.outputs[0], extrude.outputs['Top']


def _build_wall(group, front, boundary):
    """Attach the curved half wall to the front and mark its faces."""
    geometry, top = _extrude_boundary(group, front, boundary, .5)
    geometry, _ = _extrude_boundary(group, geometry, top, 0)
    return geometry


def build_image_cutout_symmetry_group():
    existing = bpy.data.node_groups.get(NODE_GROUP_NAME)
    if existing is not None:
        return existing
    group = bpy.data.node_groups.new(NODE_GROUP_NAME, 'GeometryNodeTree')
    group.is_modifier, group.use_fake_user = True, True
    group.color_tag = 'GEOMETRY'
    interface_socket(group, 'Geometry', 'INPUT', 'NodeSocketGeometry')
    vector_input(group, 'Direction', (0, 0, 1), subtype='DIRECTION').description = 'Align this direction with the local symmetry axis.'
    float_input(group, 'Offset', 0, minimum=-10, maximum=10, subtype='DISTANCE').description = 'Offset geometry along the aligned symmetry axis before scaling.'
    float_input(group, 'Scale', 1, maximum=2).description = 'Scale the offset geometry along the aligned symmetry axis.'
    options = group.interface.new_panel(name='Options', default_closed=True)
    bool_input(group, 'Fill Sides', True, parent=options).description = 'Connect matching front and mirrored boundaries.'
    smooth = interface_socket(group, 'Fill Smooth', 'INPUT', 'NodeSocketInt', options)
    smooth.default_value, smooth.min_value, smooth.max_value = 2, 0, 16
    smooth.description = 'Smooth the outline and the wall it builds through a two-ring falloff. Zero keeps the raw outline.'
    seam = interface_socket(group, 'Seam Smooth', 'INPUT', 'NodeSocketInt', options)
    seam.default_value, seam.min_value, seam.max_value = 2, 0, 16
    seam.description = 'Smooth the cut seam inside the symmetry plane. Zero keeps the raw seam.'
    interface_socket(group, 'Geometry', 'OUTPUT', 'NodeSocketGeometry')
    nodes, links = group.nodes, group.links
    inputs = group_input(nodes, set()).outputs
    aligned = _align_front(group, inputs['Geometry'], inputs)
    bounds = nodes.new('GeometryNodeBoundBox')
    links.new(aligned, bounds.inputs['Geometry'])
    size = nodes.new('ShaderNodeVectorMath')
    size.operation = 'DISTANCE'
    links.new(bounds.outputs['Min'], size.inputs[0])
    links.new(bounds.outputs['Max'], size.inputs[1])
    tolerance = _math(group, 'MULTIPLY', _math(group, 'MAXIMUM', size.outputs['Value'], 1e-8), 1e-7)
    is_boundary = _boundary_edge_field(group)
    seam_field, outline_field = _cut_fields(group, tolerance)
    front = _fold_to_plane(group, aligned, tolerance)
    shade = nodes.new('GeometryNodeSetShadeSmooth')
    shade.domain = 'FACE'
    links.new(front, shade.inputs['Geometry'])
    front = store_float_attribute(group, shade.outputs[0], FACE_ATTRIBUTE, 1., domain='FACE')
    front = _relax_seam(group, front, seam_field, inputs['Seam Smooth'])
    boundary = nodes.new('FunctionNodeBooleanMath')
    boundary.operation = 'AND'
    links.new(is_boundary, boundary.inputs[0])
    # The seam sits on the plane and welds through the merge, so only off-plane edges need walls.
    links.new(_off_plane_edge_field(group, tolerance), boundary.inputs[1])
    fill_smooth = _math(group, 'MULTIPLY', inputs['Fill Smooth'], inputs['Fill Sides'])
    front = _apply_planar_relax(group, front, _falloff(group, outline_field), fill_smooth)
    walled = _build_wall(group, front, boundary.outputs[0])
    wall_shade = nodes.new('GeometryNodeSetShadeSmooth')
    wall_shade.domain = 'FACE'
    links.new(_math(group, 'COMPARE', read_float_attribute(group, FACE_ATTRIBUTE), 3), wall_shade.inputs['Selection'])
    links.new(walled, wall_shade.inputs['Geometry'])
    body = nodes.new('GeometryNodeSwitch')
    body.input_type = 'GEOMETRY'
    links.new(inputs['Fill Sides'], body.inputs['Switch'])
    links.new(front, body.inputs['False'])
    links.new(wall_shade.outputs[0], body.inputs['True'])
    mark_back = nodes.new('GeometryNodeStoreNamedAttribute')
    mark_back.data_type, mark_back.domain = 'FLOAT', 'FACE'
    mark_back.inputs['Name'].default_value = FACE_ATTRIBUTE
    mark_back.inputs['Value'].default_value = 2
    links.new(_mirror(group, body.outputs[0]), mark_back.inputs['Geometry'])
    links.new(_math(group, 'COMPARE', read_float_attribute(group, FACE_ATTRIBUTE), 1), mark_back.inputs['Selection'])
    join = nodes.new('GeometryNodeJoinGeometry')
    for geometry in (body.outputs[0], mark_back.outputs['Geometry']):
        links.new(geometry, join.inputs[0])
    merge = nodes.new('GeometryNodeMergeByDistance')
    links.new(tolerance, merge.inputs['Distance'])
    links.new(join.outputs[0], merge.inputs['Geometry'])
    geometry = store_float_attribute(group, merge.outputs[0], AXIS_ATTRIBUTE, 1., domain='FACE')
    orient = nodes.new('GeometryNodeTransform')
    orient.inputs['Rotation'].default_value = Matrix.Rotation(1.5707963267948966, 4, 'Z').to_euler()
    links.new(geometry, orient.inputs['Geometry'])
    links.new(orient.outputs[0], nodes.new('NodeGroupOutput').inputs['Geometry'])
    prepare_node_group(group)
    return group