from qgis.core import *
from qgis.gui import *


@qgsfunction(args="auto", group="Custom", referenced_columns=[])
def get_projected_length_of_3D_line(geom, length):
    vertices = geom.constGet().points()
    proj_length = 0
    total_3Dlength = 0
    for point in range(len(vertices) - 1):
        segment_3Dlength = QgsLineString(
            [vertices[point], vertices[point + 1]]
        ).length3D()
        if total_3Dlength + segment_3Dlength > length:
            proj_length += (
                (length - total_3Dlength) / segment_3Dlength
            ) * QgsLineString([vertices[point], vertices[point + 1]]).length()
            break
        total_3Dlength += segment_3Dlength
        proj_length += QgsLineString([vertices[point], vertices[point + 1]]).length()
    return proj_length
