
import bpy
import bmesh
import numpy as np
import bgl
from mathutils import Vector, Matrix
import gpu
from gpu_extras.batch import batch_for_shader


def bounds_2d_co(coords, width, height):
    from bpy_extras.view3d_utils import location_3d_to_region_2d as space_2d

    region = bpy.context.region
    rv3d = bpy.context.space_data.region_3d

    co_2d = [space_2d(region, rv3d, co) for co in coords]

    print("co_2d", co_2d)

    x_value = sorted(co_2d, key=lambda x:x[0])
    y_value = sorted(co_2d, key=lambda x:x[1])

    min_x = x_value[0][0]
    min_y = y_value[0][1]

    box_width = x_value[-1][0] - min_x
    box_height = y_value[-1][1] - min_y

    x_margin = 1 / width # one pixel margin
    y_margin = 1 / height

    scale_fac = 2 / max(box_width, box_height)

    if box_width < box_height :
        center_offset = (box_height - box_width) /2
        min_x -= center_offset
    else :
        center_offset = (box_width - box_height) /2
        min_y -= center_offset

    co_2d_bounds = []
    for co in co_2d:
        x = (co[0] - min_x) *scale_fac * (1 - 2*x_margin) -1 + x_margin#+ x_offset
        y = (co[1] - min_y) *scale_fac * (1 - 2*y_margin) -1 + y_margin#-height/2 #* (1-y_offset) #+ y_offset

        co_2d_bounds.append(Vector((x, y)))
    

    return co_2d_bounds


def render_widget(shape, filepath, width=32, height=32) :
    dg = bpy.context.evaluated_depsgraph_get()

    bm = bmesh.new()
    bm.from_object(shape, dg)

    width *= 2
    height *= 2

    coords = [shape.matrix_world @ v.co for v in bm.verts]
    coords = bounds_2d_co(coords, width, height)
    
    indices = [(e.verts[0].index, e.verts[1].index) for e in bm.edges]

    bm.free()

    offscreen = gpu.types.GPUOffScreen(width, height)
    shader = gpu.shader.from_builtin('2D_UNIFORM_COLOR')
    batch = batch_for_shader(shader, 'LINES', {"pos": coords}, indices=indices)

    with offscreen.bind():
        fb = gpu.state.active_framebuffer_get()
        fb.clear(color=(0.0, 0.0, 0.0, 0.0))
        with gpu.matrix.push_pop():
            # reset matrices -> use normalized device coordinates [-1, 1]
            gpu.matrix.load_matrix(Matrix.Identity(4))
            gpu.matrix.load_projection_matrix(Matrix.Identity(4))


            shader.bind()
            gpu.state.line_width_set(4*2)
            #bgl.glEnable(bgl.GL_BLEND)
            #bgl.glEnable( bgl.GL_LINE_SMOOTH )
            shader.uniform_float("color", (0, 0, 0, 0.2))
            batch.draw(shader)
            
            gpu.state.line_width_set(2*2)
            shader.uniform_float("color", (0.85, 0.85, 0.85, 1))
            batch.draw(shader)

        buffer = fb.read_color(0, 0, width, height, 4, 0, 'FLOAT')
        buffer.dimensions = width * height * 4
        #buffer = bgl.Buffer(bgl.GL_BYTE, width * height * 4)
        #bgl.glReadBuffer(bgl.GL_BACK)
        #bgl.glReadPixels(0, 0, width, height, bgl.GL_RGBA, bgl.GL_UNSIGNED_BYTE, buffer)

    offscreen.free()

    #icon_name = '.' + shape.name + '_icon.png'
    icon_name = f'.{shape.name}_icon.png'
    image = bpy.data.images.get(icon_name)

    if image:
        bpy.data.images.remove(image)

    image = bpy.data.images.new(icon_name, width, height, alpha=True, float_buffer=True)

    #image_data = np.asarray(buffer, dtype=np.uint8)
    #image.pixels.foreach_set(image_data.flatten()/255)

    #image_data = np.asarray(buffer, dtype=np.uint8)
    image.pixels.foreach_set(buffer)

    image.scale(int(width/2), int(height/2))

    image.filepath_raw = str(filepath)

    image.save()

    bpy.data.images.remove(image)

