import bpy
import math
import bmesh
import random
from time import sleep
import bpy_extras

# This makes it easy to keep revising and running the script
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete(use_global=False)

# ccmp = Color Component
def random_ccmp():
	return random.randrange(0, 1000) / 1000.0

def random_color():
	return (random_ccmp(), random_ccmp(), random_ccmp(), 1.0)

mat = bpy.data.materials.new(name="VertexColorMaterial")
mat.use_nodes = True
bsdf = mat.node_tree.nodes.get("Principled BSDF")
color_node = mat.node_tree.nodes.new(type='ShaderNodeVertexColor')
color_node.layer_name = "Col"
mat.node_tree.links.new(color_node.outputs['Color'], bsdf.inputs['Base Color'])

def point_average(A, B):
	return [(A[0] + B[0]) / 2, ((A[1] + B[1]) / 2), ((A[2] + B[2]) / 2)]

def line_length(A, B): # Works ONLY for axis-aligned lines with only one component different
	return sum([math.fabs(i - j) for i,j in zip(A, B)])

verts = []
faces = []
def add_rp(A, B, C, D, E, F, G, H):
	voff = len(verts)
	verts.extend([A, B, C, D, E, F, G, H])
	faces_z = [(0, 1, 3, 2), (0, 1, 5, 4), (0, 2, 6, 4), (2, 3, 7, 6), (3, 1, 5, 7), (4, 5, 7, 6)]
	faces.extend([[n + voff for n in f] for f in faces_z])

thickness = 0.0

def add_cap(P):
	cap_size = thickness * 5
	A = (P[0] + cap_size, P[1] + cap_size, +cap_size)
	B = (P[0] - cap_size, P[1] + cap_size, +cap_size)
	C = (P[0] + cap_size, P[1] + cap_size, -cap_size)
	D = (P[0] - cap_size, P[1] + cap_size, -cap_size)
	E = (P[0] + cap_size, P[1] - cap_size, +cap_size)
	F = (P[0] - cap_size, P[1] - cap_size, +cap_size)
	G = (P[0] + cap_size, P[1] - cap_size, -cap_size)
	H = (P[0] - cap_size, P[1] - cap_size, -cap_size)
	add_rp(A, B, C, D, E, F, G, H)

def drawline(A, B, short_top = False):
	global thickness
	if thickness == 0.0:
		thickness = line_length(A, B) / 10
	if math.fabs(A[0] - B[0]) < math.fabs(A[1] - B[1]): # Along X axis
		expansion = -thickness if A[1] > B[1] else +thickness
		a = (A[0] + thickness, A[1] - expansion, thickness)
		b = (A[0] - thickness, A[1] - expansion, thickness)
		c = (A[0] + thickness, A[1] - expansion, -thickness)
		d = (A[0] - thickness, A[1] - expansion, -thickness)
		if short_top:
			expansion -= thickness * 2
		e = (B[0] + thickness, B[1] + expansion, thickness)
		f = (B[0] - thickness, B[1] + expansion, thickness)
		g = (B[0] + thickness, B[1] + expansion, -thickness)
		h = (B[0] - thickness, B[1] + expansion, -thickness)
		add_rp(a, b, c, d, e, f, g, h)	
	else: # Along Y axis
		expansion = +thickness if A[0] > B[0] else -thickness
		a = (A[0] - expansion, A[1] + thickness, thickness)
		b = (A[0] - expansion, A[1] - thickness, thickness)
		c = (A[0] - expansion, A[1] + thickness, -thickness)
		d = (A[0] - expansion, A[1] - thickness, -thickness)
		if short_top:
			expansion += thickness * 2
		e = (B[0] + expansion, B[1] + thickness, thickness)
		f = (B[0] + expansion, B[1] - thickness, thickness)
		g = (B[0] + expansion, B[1] + thickness, -thickness)
		h = (B[0] + expansion, B[1] - thickness, -thickness)
		add_rp(a, b, c, d, e, f, g, h)	

def replace_line(A, B, subdivisions_left=1): 
	I = point_average(A, B)
	H = point_average(A, I)
	J = point_average(I, B)
	direction = 0
	if math.fabs(A[0] - B[0]) > math.fabs(A[1] - B[1]):
		direction = 1
	length = line_length(A, H)

	K = H.copy()
	K[direction] -= length
	L = I.copy()
	L[direction] -= length
	M = I.copy()
	M[direction] += length
	N = J.copy()
	N[direction] += length

	if subdivisions_left > 0:
		replace_line(A, H, subdivisions_left - 1)
		replace_line(H, K, subdivisions_left - 1)
		replace_line(K, L, subdivisions_left - 1)
		replace_line(I, L, subdivisions_left - 1)
		replace_line(M, I, subdivisions_left - 1)
		replace_line(M, N, subdivisions_left - 1)
		replace_line(N, J, subdivisions_left - 1)
		replace_line(J, B, subdivisions_left - 1)
	else:					  
		drawline(A, H)
		drawline(H, K)		 
		drawline(K, L)		 
		drawline(L, I, short_top = True)		 
		drawline(I, M)		 
		drawline(M, N)	
		drawline(N, J) 
		drawline(J, B)
		[add_cap(P) for P in [H, K, L, I, M, N, J, B]]

replace_line((-10, 0, 0), (10, 0, 0), 3)

worm_mesh = bpy.data.meshes.new(name="Worm Mesh")
worm_mesh.from_pydata(vertices=verts, edges=[], faces=faces)
worm_mesh.validate(verbose=True)
new_object = bpy_extras.object_utils.object_data_add(bpy.context, worm_mesh)
new_object.data.materials.append(mat)

bpy.ops.object.mode_set(mode='EDIT')
worm_bmesh = bmesh.from_edit_mesh(new_object.data)
color_layer = worm_bmesh.loops.layers.color.new("Col")
for i, face in enumerate(worm_bmesh.faces):
	if i % 6 == 0:
		face_color = random_color()
	for loop in face.loops:
		loop[color_layer] = face_color
bmesh.update_edit_mesh(new_object.data)
bpy.ops.object.mode_set(mode='OBJECT')