"""Face separation and connected thickness for Depth Surface assets."""

TEMPORARY_ATTRIBUTES = {
    "face_center": ("FLOAT_VECTOR", "FACE"),
    "face_camera": ("FLOAT_VECTOR", "FACE"),
    "corner_camera": ("FLOAT_VECTOR", "CORNER"),
    "cut_vertex": ("BOOLEAN", "CORNER"),
    "camera": ("FLOAT_VECTOR", "POINT"),
    "surface_position": ("FLOAT_VECTOR", "POINT"),
    "source_vertex": ("INT", "POINT"),
    "surface_normal": ("FLOAT_VECTOR", "POINT"),
}


def _store(group, geometry, name, value):
    node = group.nodes.new("GeometryNodeStoreNamedAttribute")
    node.data_type, node.domain = TEMPORARY_ATTRIBUTES[name]
    node.inputs["Name"].default_value = "o_depth_" + name
    group.links.new(geometry, node.inputs["Geometry"])
    group.links.new(value, node.inputs["Value"])
    return node.outputs["Geometry"]


def _read(group, name):
    node = group.nodes.new("GeometryNodeInputNamedAttribute")
    node.data_type = TEMPORARY_ATTRIBUTES[name][0]
    node.inputs["Name"].default_value = "o_depth_" + name
    return node.outputs["Attribute"]


def remove_surface_attributes(group, geometry):
    """Clear the reserved temporary namespace with one wildcard operation."""
    node = group.nodes.new("GeometryNodeRemoveAttribute")
    node.pattern_mode = "WILDCARD"
    node.inputs["Name"].default_value = "o_depth_*"
    group.links.new(geometry, node.inputs["Geometry"])
    return node.outputs["Geometry"]


def _math(group, operation, a, b=None):
    node = group.nodes.new("ShaderNodeMath")
    node.operation = operation
    for index, value in enumerate((a, b)):
        if value is None:
            continue
        if isinstance(value, (int, float)):
            node.inputs[index].default_value = value
        else:
            group.links.new(value, node.inputs[index])
    return node.outputs[0]


def split_depth_surface(group, geometry, image, strength, uniform):
    """Keep each face's camera sample on its own corners before splitting."""
    nodes, links = group.nodes, group.links
    uv = nodes.new("GeometryNodeInputNamedAttribute")
    uv.data_type = "FLOAT_VECTOR"
    uv.inputs["Name"].default_value = "UVMap"
    face_uv = nodes.new("GeometryNodeFieldOnDomain")
    face_uv.data_type = "FLOAT_VECTOR"
    face_uv.domain = "FACE"
    links.new(uv.outputs["Attribute"], face_uv.inputs["Value"])
    texture = nodes.new("GeometryNodeImageTexture")
    texture.interpolation = "Linear"
    texture.extension = "EXTEND"
    links.new(image, texture.inputs["Image"])
    links.new(face_uv.outputs["Value"], texture.inputs["Vector"])
    position = nodes.new("GeometryNodeInputPosition")
    faces = _store(group, geometry, "face_center", position.outputs[0])
    faces = _store(group, faces, "face_camera", texture.outputs["Color"])
    axes = nodes.new("ShaderNodeSeparateXYZ")
    links.new(_read(group, "face_camera"), axes.inputs[0])
    center_field = _read(group, "face_center")
    centers, depths = [], []
    for side in (0, 1):
        corner = nodes.new("GeometryNodeCornersOfEdge")
        corner.inputs["Sort Index"].default_value = side
        face = nodes.new("GeometryNodeFaceOfCorner")
        links.new(corner.outputs["Corner Index"], face.inputs["Corner Index"])
        for data_type, field, results in (
            ("FLOAT_VECTOR", center_field, centers),
            ("FLOAT", axes.outputs["Z"], depths),
        ):
            sample = nodes.new("GeometryNodeFieldAtIndex")
            sample.data_type = data_type
            sample.domain = "FACE"
            links.new(face.outputs["Face Index"], sample.inputs["Index"])
            links.new(field, sample.inputs["Value"])
            results.append(sample.outputs["Value"])
    distance = nodes.new("ShaderNodeVectorMath")
    distance.operation = "DISTANCE"
    for index, center in enumerate(centers):
        links.new(center, distance.inputs[index])
    delta = _math(group, "ABSOLUTE", _math(group, "SUBTRACT", *depths))
    scaled = _math(group, "MULTIPLY", delta, uniform)
    # Cross multiplication avoids division and stays invariant under uniform scaling.
    selected = _math(
        group,
        "GREATER_THAN",
        _math(group, "MULTIPLY", scaled, strength),
        _math(group, "MULTIPLY", distance.outputs["Value"], 2.5),
    )
    selected = _math(
        group,
        "MULTIPLY",
        selected,
        _math(group, "GREATER_THAN", distance.outputs["Value"], 0.0),
    )
    neighbors = nodes.new("GeometryNodeInputMeshEdgeNeighbors")
    selected = _math(
        group,
        "MULTIPLY",
        selected,
        _math(group, "COMPARE", neighbors.outputs["Face Count"], 2.0),
    )
    cut_edges = nodes.new("GeometryNodeFieldOnDomain")
    cut_edges.data_type = "FLOAT"
    cut_edges.domain = "EDGE"
    links.new(selected, cut_edges.inputs["Value"])
    cut_vertices = nodes.new("GeometryNodeFieldOnDomain")
    cut_vertices.data_type = "FLOAT"
    cut_vertices.domain = "POINT"
    links.new(cut_edges.outputs["Value"], cut_vertices.inputs["Value"])
    corners = _store(group, faces, "corner_camera", _read(group, "face_camera"))
    corners = _store(
        group,
        corners,
        "cut_vertex",
        _math(group, "GREATER_THAN", cut_vertices.outputs["Value"], 0.0),
    )
    split = nodes.new("GeometryNodeSplitEdges")
    links.new(corners, split.inputs["Mesh"])
    links.new(selected, split.inputs["Selection"])
    enabled = _math(group, "GREATER_THAN", strength, 0.0)
    choose = nodes.new("GeometryNodeSwitch")
    choose.input_type = "GEOMETRY"
    links.new(enabled, choose.inputs["Switch"])
    links.new(geometry, choose.inputs["False"])
    links.new(split.outputs["Mesh"], choose.inputs["True"])
    return choose.outputs["Output"]


def store_surface_camera(group, geometry, original, strength):
    """Change cut depth along the original ray, preserving every other vertex."""
    separated = _read(group, "corner_camera")
    enabled = _math(
        group,
        "MULTIPLY",
        _math(group, "GREATER_THAN", strength, 0.0),
        _read(group, "cut_vertex"),
    )
    original_axes = group.nodes.new("ShaderNodeSeparateXYZ")
    separated_axes = group.nodes.new("ShaderNodeSeparateXYZ")
    group.links.new(original, original_axes.inputs[0])
    group.links.new(separated, separated_axes.inputs[0])
    ratio = _math(
        group, "DIVIDE", separated_axes.outputs["Z"], original_axes.outputs["Z"]
    )
    projected = group.nodes.new("ShaderNodeVectorMath")
    projected.operation = "SCALE"
    group.links.new(original, projected.inputs[0])
    group.links.new(ratio, projected.inputs[3])
    choose = group.nodes.new("GeometryNodeSwitch")
    choose.input_type = "VECTOR"
    group.links.new(enabled, choose.inputs["Switch"])
    group.links.new(original, choose.inputs["False"])
    group.links.new(projected.outputs[0], choose.inputs["True"])
    geometry = _store(group, geometry, "camera", choose.outputs["Output"])
    return geometry, _read(group, "camera")


def thicken_depth_surface(group, geometry, thickness, normal_smooth=None):
    """Connect caps with single quads by source vertex and cap identity."""
    nodes, links = group.nodes, group.links
    front_geometry = geometry
    position = nodes.new("GeometryNodeInputPosition")
    geometry = _store(group, geometry, "surface_position", position.outputs[0])
    surface_store = geometry.node
    flattened = nodes.new("ShaderNodeVectorMath")
    flattened.operation = "MULTIPLY"
    flattened.inputs[1].default_value = (1.0, 1.0, 0.0)
    links.new(position.outputs[0], flattened.inputs[0])
    sheet = nodes.new("GeometryNodeSetPosition")
    links.new(geometry, sheet.inputs["Geometry"])
    links.new(flattened.outputs[0], sheet.inputs["Position"])
    index = nodes.new("GeometryNodeInputIndex")
    geometry = _store(
        group, sheet.outputs["Geometry"], "source_vertex", index.outputs[0]
    )
    saved = geometry

    # Weld in an integer identity grid, never in scene coordinates. Nearby cuts
    # retain distinct IDs; copies of one rim vertex share the same ID and layer.
    identity = nodes.new("ShaderNodeCombineXYZ")
    for operation, axis in (("MODULO", "X"), ("DIVIDE", "Y")):
        part = nodes.new("FunctionNodeIntegerMath")
        part.operation = operation
        part.inputs[1].default_value = 4096
        links.new(_read(group, "source_vertex"), part.inputs[0])
        links.new(part.outputs[0], identity.inputs[axis])

    extrude = nodes.new("GeometryNodeExtrudeMesh")
    extrude.mode = "FACES"
    extrude.inputs["Individual"].default_value = False
    extrusion_offset = nodes.new("ShaderNodeCombineXYZ")
    extrusion_offset.inputs["Z"].default_value = -1.0
    links.new(extrusion_offset.outputs[0], extrude.inputs["Offset"])
    links.new(saved, extrude.inputs["Mesh"])
    sides = nodes.new("GeometryNodeSeparateGeometry")
    sides.domain = "FACE"
    links.new(extrude.outputs["Mesh"], sides.inputs["Geometry"])
    links.new(extrude.outputs["Side"], sides.inputs["Selection"])
    offset = nodes.new("ShaderNodeCombineXYZ")
    offset.inputs["Z"].default_value = -1.0
    back = nodes.new("GeometryNodeSetPosition")
    links.new(saved, back.inputs["Geometry"])
    links.new(offset.outputs[0], back.inputs["Offset"])
    flip_back = nodes.new("GeometryNodeFlipFaces")
    links.new(back.outputs["Geometry"], flip_back.inputs["Mesh"])
    flip_sides = nodes.new("GeometryNodeFlipFaces")
    links.new(sides.outputs["Selection"], flip_sides.inputs["Mesh"])
    join = nodes.new("GeometryNodeJoinGeometry")
    for value in (
        saved,
        flip_back.outputs["Mesh"],
        flip_sides.outputs["Mesh"],
    ):
        links.new(value, join.inputs["Geometry"])
    encoded = nodes.new("GeometryNodeSetPosition")
    links.new(join.outputs["Geometry"], encoded.inputs["Geometry"])
    links.new(identity.outputs[0], encoded.inputs["Position"])
    merge = nodes.new("GeometryNodeMergeByDistance")
    merge.mode = "ALL"
    merge.inputs["Distance"].default_value = 0.1
    links.new(encoded.outputs["Geometry"], merge.inputs["Geometry"])
    displacement = nodes.new("ShaderNodeCombineXYZ")
    layer_position = nodes.new("GeometryNodeInputPosition")
    layer_axes = nodes.new("ShaderNodeSeparateXYZ")
    links.new(layer_position.outputs[0], layer_axes.inputs[0])
    links.new(
        _math(group, "MULTIPLY", layer_axes.outputs["Z"], thickness),
        displacement.inputs["Z"],
    )
    restored = nodes.new("ShaderNodeVectorMath")
    restored.operation = "ADD"
    links.new(_read(group, "surface_position"), restored.inputs[0])
    links.new(displacement.outputs[0], restored.inputs[1])
    result = nodes.new("GeometryNodeSetPosition")
    links.new(merge.outputs["Geometry"], result.inputs["Geometry"])
    links.new(restored.outputs[0], result.inputs["Position"])
    local_position = nodes.new("GeometryNodeInputPosition")
    local_axes = nodes.new("ShaderNodeSeparateXYZ")
    links.new(local_position.outputs[0], local_axes.inputs[0])
    links.new(local_axes.outputs["Z"], identity.inputs["Z"])
    if normal_smooth is not None:
        normal = nodes.new("GeometryNodeInputNormal")
        blur = nodes.new("GeometryNodeBlurAttribute")
        blur.data_type = "FLOAT_VECTOR"
        links.new(normal_smooth, blur.inputs["Iterations"])
        links.new(normal.outputs[0], blur.inputs["Value"])
        normalize = nodes.new("ShaderNodeVectorMath")
        normalize.operation = "NORMALIZE"
        links.new(blur.outputs[0], normalize.inputs[0])
        front_geometry = _store(group, front_geometry, "surface_normal", normalize.outputs[0])
        links.new(front_geometry, surface_store.inputs["Geometry"])
        normal_offset = nodes.new("ShaderNodeVectorMath")
        normal_offset.operation = "SCALE"
        links.new(_read(group, "surface_normal"), normal_offset.inputs[0])
        links.new(displacement.inputs["Z"].links[0].from_socket, normal_offset.inputs[3])
        links.new(normal_offset.outputs[0], restored.inputs[1])
        nodes.remove(displacement)
    return result.outputs["Geometry"]


def _projection_math(nodes, operation):
    node = nodes.new("ShaderNodeMath")
    node.operation = operation
    return node


def project_depth_surface(
    group, geometry, camera_position, uniform_scale, reference_depth, depth_scale
):
    """Interpolate the rectangular plane into the model camera coordinates."""
    nodes, links = group.nodes, group.links
    scaled = nodes.new("ShaderNodeVectorMath")
    scaled.operation = "SCALE"
    links.new(camera_position, scaled.inputs[0])
    links.new(uniform_scale, scaled.inputs[3])
    separate = nodes.new("ShaderNodeSeparateXYZ")
    links.new(scaled.outputs["Vector"], separate.inputs["Vector"])
    relative_depth = _projection_math(nodes, "SUBTRACT")
    links.new(reference_depth, relative_depth.inputs[0])
    links.new(separate.outputs["Z"], relative_depth.inputs[1])
    front = _projection_math(nodes, "MULTIPLY")
    links.new(relative_depth.outputs[0], front.inputs[0])
    links.new(depth_scale, front.inputs[1])

    negative_y = _projection_math(nodes, "MULTIPLY")
    negative_y.inputs[1].default_value = -1.0
    links.new(separate.outputs["Y"], negative_y.inputs[0])
    projected = nodes.new("ShaderNodeCombineXYZ")
    links.new(separate.outputs["X"], projected.inputs["X"])
    links.new(negative_y.outputs[0], projected.inputs["Y"])
    current_position = nodes.new("GeometryNodeInputPosition")
    blend = nodes.new("ShaderNodeMix")
    blend.data_type = "VECTOR"
    blend.clamp_factor = False
    links.new(depth_scale, blend.inputs[0])
    links.new(current_position.outputs["Position"], blend.inputs[4])
    links.new(projected.outputs["Vector"], blend.inputs[5])
    blended_axes = nodes.new("ShaderNodeSeparateXYZ")
    links.new(blend.outputs["Result"], blended_axes.inputs["Vector"])
    surface_position = nodes.new("ShaderNodeCombineXYZ")
    links.new(blended_axes.outputs["X"], surface_position.inputs["X"])
    links.new(blended_axes.outputs["Y"], surface_position.inputs["Y"])
    links.new(front.outputs[0], surface_position.inputs["Z"])
    project_geometry = nodes.new("GeometryNodeSetPosition")
    links.new(geometry, project_geometry.inputs["Geometry"])
    links.new(surface_position.outputs["Vector"], project_geometry.inputs["Position"])

    return project_geometry.outputs["Geometry"]