import numpy as np
from PyQt6.QtOpenGLWidgets import QOpenGLWidget
from PyQt6.QtCore import Qt, QEvent, QThread, pyqtSignal
from PyQt6.QtGui import QTouchEvent, QSurfaceFormat
from PyQt6.QtOpenGL import QOpenGLShaderProgram, QOpenGLShader, QOpenGLBuffer

class GCodeParseWorker(QThread):
    finished = pyqtSignal(dict)

    def __init__(self, filepath, type_map, default_colors, parent=None):
        super().__init__(parent)
        self.filepath = filepath
        self.TYPE_MAP = type_map
        self.DEFAULT_COLORS = default_colors

    def run(self):
        points = []
        colors = []
        type_segments = {}
        type_visibility = {}

        x = y = z = e = 0.0
        vertex_idx = 0
        feature_type = 'OTHER'
        current_segment_type = 'OTHER'
        segment_start = 0
        type_visibility['TRAVEL'] = False
        relative_e = False

        min_x = min_y = min_z = float('inf')
        max_x = max_y = max_z = float('-inf')
        
        current_offset = 0
        max_z_seen = -999.0
        layer_map = [(0, 0)]

        def add_segment(t_name, start, length):
            if length > 0:
                type_segments.setdefault(t_name, []).append((start, length))

        try:
            with open(self.filepath, 'rb') as f:
                for line_bytes in f:
                    current_offset += len(line_bytes)
                    line = line_bytes.decode('utf-8', errors='ignore').strip()
                    if not line:
                        continue
                    if line.startswith('M82'):
                        relative_e = False
                        continue
                    if line.startswith('M83'):
                        relative_e = True
                        continue
                    
                    if line.startswith(';'):
                        if line.startswith(';TYPE:'):
                            raw_type = line.split(':', 1)[1].strip()
                        else:
                            raw_type = line[1:].strip()
                            if raw_type not in self.TYPE_MAP and not any(k in raw_type.lower() for k in ('perimeter', 'infill', 'material', 'skirt/brim')):
                                continue
                            if 'Skirt/Brim' in raw_type:
                                raw_type = 'Skirt'
                        
                        new_type = self.TYPE_MAP.get(raw_type, 'OTHER')
                        if new_type != feature_type:
                            feature_type = new_type
                            if feature_type not in type_visibility:
                                type_visibility[feature_type] = True
                        continue

                    if line.startswith(('G0', 'G1')):
                        new_x, new_y, new_z = x, y, z
                        e_val = 0.0
                        has_e = False
                        parts = line.split(';')[0].split()
                        for p in parts[1:]:
                            try:
                                if p.startswith('X'): new_x = float(p[1:])
                                elif p.startswith('Y'): new_y = float(p[1:])
                                elif p.startswith('Z'): new_z = float(p[1:])
                                elif p.startswith('E'):
                                    e_val = float(p[1:])
                                    has_e = True
                            except ValueError:
                                pass
                        
                        if new_x == x and new_y == y and new_z == z:
                            if has_e:
                                if relative_e: e += e_val
                                else: e = e_val
                            continue

                        is_extrusion = False
                        if has_e:
                            if relative_e:
                                is_extrusion = e_val > 0
                                new_e = e + e_val
                            else:
                                is_extrusion = e_val > e
                                new_e = e_val
                        else:
                            new_e = e

                        if is_extrusion:
                            seg_type = feature_type
                            c = self.DEFAULT_COLORS.get(feature_type, self.DEFAULT_COLORS['OTHER'])
                        else:
                            seg_type = 'TRAVEL'
                            c = self.DEFAULT_COLORS['TRAVEL']
                            if 'TRAVEL' not in type_visibility:
                                type_visibility['TRAVEL'] = False

                        if seg_type != current_segment_type:
                            if vertex_idx > segment_start:
                                add_segment(current_segment_type, segment_start, vertex_idx - segment_start)
                            current_segment_type = seg_type
                            segment_start = vertex_idx

                        if new_x < min_x: min_x = new_x
                        if new_x > max_x: max_x = new_x
                        if new_y < min_y: min_y = new_y
                        if new_y > max_y: max_y = new_y
                        if new_z < min_z: min_z = new_z
                        if new_z > max_z: max_z = new_z

                        if new_z > max_z_seen and is_extrusion:
                            max_z_seen = new_z
                            layer_map.append((current_offset, vertex_idx))

                        points.extend([x, y, z, new_x, new_y, new_z])
                        colors.extend([*c, *c])
                        vertex_idx += 2
                        x, y, z, e = new_x, new_y, new_z, new_e
                    elif line.startswith('G92'):
                        parts = line.split(';')[0].split()
                        for p in parts[1:]:
                            try:
                                if p.startswith('E'): e = float(p[1:])
                                elif p.startswith('X'): x = float(p[1:])
                                elif p.startswith('Y'): y = float(p[1:])
                                elif p.startswith('Z'): z = float(p[1:])
                            except ValueError:
                                pass

            if vertex_idx > segment_start:
                add_segment(current_segment_type, segment_start, vertex_idx - segment_start)

            cx = (min_x + max_x) / 2.0 if max_x >= min_x else 110.0
            cy = (min_y + max_y) / 2.0 if max_y >= min_y else 110.0
            cz = (min_z + max_z) / 2.0 if max_z >= min_z else 0.0

            result = {
                'vertices': np.array(points, dtype=np.float32) if vertex_idx > 0 else np.zeros((0,), dtype=np.float32),
                'colors': np.array(colors, dtype=np.float32) if vertex_idx > 0 else np.zeros((0,), dtype=np.float32),
                'vertex_count': vertex_idx,
                'center_x': cx,
                'center_y': cy,
                'center_z': cz,
                'type_segments': type_segments,
                'type_visibility': type_visibility,
                'layer_map': layer_map
            }
            self.finished.emit(result)
        except Exception as e:
            print("ParseGCode Error:", e)
            self.finished.emit({})

class GCodeViewerWidget(QOpenGLWidget):
    """
    3D G-code 预览控件（OpenGL ES / eglfs 兼容版）
    使用可编程管线替代固定管线，适配树莓派 4B
    """

    # PrusaSlicer 类型映射、颜色定义等保持不变
    TYPE_MAP = {
        'External perimeter': 'WALL-OUTER',
        'Perimeter': 'WALL-INNER',
        'Overhang perimeter': 'WALL-OUTER',
        'Solid infill': 'SKIN',
        'Top solid infill': 'SKIN',
        'Bridge infill': 'SKIN',
        'Internal infill': 'FILL',
        'Support material': 'SUPPORT',
        'Support material interface': 'SUPPORT-INTERFACE',
        'Skirt': 'SKIRT',
        'Brim': 'SKIRT',
        'Custom': 'OTHER',
    }

    DEFAULT_COLORS = {
        'WALL-OUTER':        (0.92, 0.55, 0.22),  # 0xeb8b38
        'WALL-INNER':        (0.25, 0.50, 0.81),  # 0x4080cf
        'FILL':              (0.80, 0.75, 0.29),  # 0xccc04b
        'SKIN':              (0.62, 0.38, 0.70),  # 0x9e60b3
        'SUPPORT':           (0.34, 0.70, 0.34),  # 0x57b357
        'SUPPORT-INTERFACE': (0.17, 0.42, 0.17),  # 0x2b6b2b
        'SKIRT':             (0.00, 1.00, 1.00),  # 0x00ffff
        'OTHER':             (0.67, 0.67, 0.67),  # 0xaaaaaa
        'TRAVEL':            (0.25, 0.31, 0.38),  # 0x405060
    }

    # 顶点着色器（GLSL ES 1.00）
    VERTEX_SHADER = """
        attribute vec3 aPos;
        attribute vec3 aColor;
        varying vec3 vColor;
        uniform mat4 uMVP;
        void main() {
            gl_Position = uMVP * vec4(aPos, 1.0);
            vColor = aColor;
        }
    """

    # 片段着色器（添加可调节颜色深度的 uniform 以实现边界加深）
    FRAGMENT_SHADER = """
        precision mediump float;
        varying vec3 vColor;
        uniform float uDarken;
        void main() {
            gl_FragColor = vec4(vColor * uDarken, 1.0);
        }
    """

    def __init__(self, parent=None):
        
        
        
        # 请求 OpenGL ES 2.0 上下文
        fmt = QSurfaceFormat()
        fmt.setRenderableType(QSurfaceFormat.RenderableType.OpenGLES)
        fmt.setVersion(2, 0)
        
        super().__init__(parent)
        self.setMinimumSize(400, 300)
        self.setAttribute(Qt.WidgetAttribute.WA_AcceptTouchEvents, True)
        self.setFormat(fmt)
        # 数据缓冲区
        self.vertices = None
        self.colors = None
        self.vertex_count = 0
        self.vbo_vertices = None
        self.vbo_colors = None
        self.vbo_ready = False

        # 类型分段
        self.type_segments = {}
        self.type_visibility = {}
        self.current_type = 'OTHER'

        # 视角参数
        self.view_rot_x = -60.0
        self.view_rot_z = 45.0
        self.view_zoom = -250.0  # 稍微退后一点以便看到整个打印平台
        self.view_trans_x = 0.0
        self.view_trans_y = 0.0

        # ...
        self.progress_ratio = 1.0
        self.progress_vertices = 0

        # 按需渲染：缓存当前的 filepos 对应的顶点数
        self.layer_map = [(0, 0)]  # (offset_bytes, vertex_idx)
        self.last_reported_filepos = -1

        # 模型中心点
        self.center_x = 110.0
        self.center_y = 110.0
        self.center_z = 0.0

        # 触摸状态
        self.last_mouse_pos = None
        self._touch_points = {}
        self._pinch_start_dist = 0.0
        self._pinch_start_zoom = 0.0
        self._pinch_start_center = None
        self._pinch_start_trans = (0.0, 0.0)
        self._ignore_wheel = False

        # 着色器程序
        self.shader_program = None
        self.aPos_location = None
        self.aColor_location = None
        self.uMVP_location = None

    # ── 公开接口 ──
    def load_gcode(self, filepath: str):
        if hasattr(self, '_worker') and self._worker.isRunning():
            self._worker.terminate()
            self._worker.wait()
            
        self._worker = GCodeParseWorker(filepath, self.TYPE_MAP, self.DEFAULT_COLORS)
        self._worker.finished.connect(self._on_parse_finished)
        self._worker.start()

    def _on_parse_finished(self, result: dict):
        if not result:
            return
        
        self.vertices = result['vertices']
        self.colors = result['colors']
        self.vertex_count = result['vertex_count']
        self.center_x = result['center_x']
        self.center_y = result['center_y']
        self.center_z = result['center_z']
        self.type_segments = result['type_segments']
        self.type_visibility = result['type_visibility']
        self.layer_map = result['layer_map']
        
        self.vbo_ready = False
        # 初始时先不显示，让 update_by_filepos 决定，或者如果是0则自动更新
        # 这里把 progress_vertices 赋予当前的 target
        target = 0
        if self.layer_map:
            target = self.layer_map[-1][1] # 全显
        
        self.progress_vertices = target
        self.last_reported_filepos = -1
        self.update()

    def update_processes(self, progress: float):
        pass

    def update_by_filepos(self, filepos: int, is_printing: bool = True):
        import bisect
        if not hasattr(self, 'layer_map') or not self.layer_map:
            return
        
        if not is_printing:
            # 不在打印途中时，渲染包含所有顶点（完整模型）
            target_vertices = self.vertex_count
        else:
            # 使用二分查找快速定位当前 filepos 对应的最多展示顶点数
            keys = [item[0] for item in self.layer_map]
            idx = bisect.bisect_right(keys, filepos)
            target_vertices = 0 if idx == 0 else self.layer_map[idx-1][1]
            
        # 按需渲染：只有当层级（计算出的可见顶点数）真正发生跳变时才调用 update()
        if target_vertices != getattr(self, 'progress_vertices', -1):
            self.progress_vertices = target_vertices
            self.update()

    def update_switch(self, type_name: str, visible: bool):
        if type_name in self.type_visibility:
            self.type_visibility[type_name] = visible
            self.update()

    def set_view_angles(self, rot_x: float, rot_z: float, zoom: float = None):
        self.view_rot_x = max(-90.0, min(0.0, rot_x))
        self.view_rot_z = rot_z
        if zoom is not None:
            self.view_zoom = zoom
        self.update()

        if length == 0:
            return
        self.type_segments.setdefault(type_name, []).append((start, length))

    # ── OpenGL 可编程管线初始化 ──
    def initializeGL(self):
        import OpenGL.GL as gl
        gl.glClearColor(0.15, 0.15, 0.15, 1.0)
        gl.glEnable(gl.GL_DEPTH_TEST)
        gl.glLineWidth(1.0)

        # 编译着色器
        self.shader_program = QOpenGLShaderProgram()
        self.shader_program.addShaderFromSourceCode(QOpenGLShader.ShaderTypeBit.Vertex, self.VERTEX_SHADER)
        self.shader_program.addShaderFromSourceCode(QOpenGLShader.ShaderTypeBit.Fragment, self.FRAGMENT_SHADER)
        self.shader_program.link()

        # 获取属性和 uniform 位置
        self.aPos_location = self.shader_program.attributeLocation("aPos")
        self.aColor_location = self.shader_program.attributeLocation("aColor")
        self.uMVP_location = self.shader_program.uniformLocation("uMVP")
        self.uDarken_location = self.shader_program.uniformLocation("uDarken")

        # 创建缓冲对象
        self.vbo_vertices = QOpenGLBuffer(QOpenGLBuffer.Type.VertexBuffer)
        self.vbo_colors = QOpenGLBuffer(QOpenGLBuffer.Type.VertexBuffer)

    def resizeGL(self, w, h):
        import OpenGL.GL as gl
        gl.glViewport(0, 0, w, h)

    # ── 构建 MVP 矩阵（替代 glOrtho/glRotate） ──
    def _build_mvp(self):
        from PyQt6.QtGui import QMatrix4x4
        mat = QMatrix4x4()
        aspect = self.width() / self.height() if self.height() else 1
        mat.perspective(45.0, aspect, 1.0, 1000.0)
        
        mat.translate(self.view_trans_x, self.view_trans_y, self.view_zoom)
        mat.rotate(self.view_rot_x, 1.0, 0.0, 0.0)
        mat.rotate(self.view_rot_z, 0.0, 0.0, 1.0)
        mat.translate(-self.center_x, -self.center_y, -self.center_z)
        
        return mat

    # ── 渲染 ──
    def paintGL(self):
        import OpenGL.GL as gl
        gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)

        if self.vertices is None or self.vertex_count == 0:
            return

        if not self.vbo_ready:
            self._create_vbos()
            self.vbo_ready = True

        # 使用着色器程序
        self.shader_program.bind()

        # 设置 MVP 矩阵
        import math
        from PyQt6.QtGui import QMatrix4x4
        mvp_mat = self._build_mvp()
        self.shader_program.setUniformValue(self.uMVP_location, mvp_mat)

        # 绑定 VBO
        self.vbo_vertices.bind()
        self.shader_program.setAttributeBuffer(self.aPos_location, gl.GL_FLOAT, 0, 3, 0)
        self.shader_program.enableAttributeArray(self.aPos_location)

        self.vbo_colors.bind()
        self.shader_program.setAttributeBuffer(self.aColor_location, gl.GL_FLOAT, 0, 3, 0)
        self.shader_program.enableAttributeArray(self.aColor_location)

        # 允许 z-fighting 覆盖，用于同一位置多次渲染线条
        gl.glDepthFunc(gl.GL_LEQUAL)

        # 渲染两次：一次绘制加粗加深的边界底线，一次绘制正常宽度的原色骨架线
        # （在树莓派等性能有限的平台上，使用真实的3D圆柱/方块代替线条会导致顶点数暴增十几倍直接卡顿，
        # 因此通过动态加粗像素级线宽来性能无损地模拟出“体积感”）
        for pass_idx in range(2):
            if pass_idx == 0:
                gl.glLineWidth(6.0)  # 底线宽度（加大以模拟体积轮廓）
                self.shader_program.setUniformValue(self.uDarken_location, 0.8) # 加深颜色至 40% 亮度
            else:
                gl.glLineWidth(3.0)  # 主体宽度（加大以模拟线条厚度）
                self.shader_program.setUniformValue(self.uDarken_location, 1.0) # 保持原色

            # 按类型分段绘制
            for type_name, segments in self.type_segments.items():
                if not self.type_visibility.get(type_name, True):
                    continue
                for start, length in segments:
                    if start >= self.progress_vertices:
                        continue
                    end = start + length
                    visible_start = start
                    visible_count = length
                    if end > self.progress_vertices:
                        visible_count = self.progress_vertices - start
                    if visible_count > 0:
                        gl.glDrawArrays(gl.GL_LINES, visible_start, visible_count)

        # 恢复默认深度测试模式
        gl.glDepthFunc(gl.GL_LESS)

        self.shader_program.disableAttributeArray(self.aPos_location)
        self.shader_program.disableAttributeArray(self.aColor_location)
        self.vbo_vertices.release()
        self.vbo_colors.release()
        self.shader_program.release()

    def _create_vbos(self):
        if self.vbo_vertices.isCreated():
            self.vbo_vertices.destroy()
        if self.vbo_colors.isCreated():
            self.vbo_colors.destroy()

        self.vbo_vertices.create()
        self.vbo_vertices.bind()
        self.vbo_vertices.allocate(self.vertices.tobytes(), self.vertices.nbytes)
        self.vbo_vertices.release()

        self.vbo_colors.create()
        self.vbo_colors.bind()
        self.vbo_colors.allocate(self.colors.tobytes(), self.colors.nbytes)
        self.vbo_colors.release()

    # ── 触摸/鼠标交互（完全不变） ──
    def mousePressEvent(self, event):
        self.last_mouse_pos = event.position()

    def mouseMoveEvent(self, event):
        if self.last_mouse_pos is None:
            return
        dx = event.position().x() - self.last_mouse_pos.x()
        dy = event.position().y() - self.last_mouse_pos.y()
        if event.buttons() & Qt.MouseButton.LeftButton:
            self.view_rot_x += dy * 0.5
            self.view_rot_x = max(-90.0, min(0.0, self.view_rot_x))  # 限制垂直视角的翻转
            self.view_rot_z += dx * 0.5
        self.last_mouse_pos = event.position()
        self.update()

    def wheelEvent(self, event):
        delta = event.angleDelta().y() / 120
        self.view_zoom += delta * 10
        self.update()

    def event(self, e):
        if e.type() in (QEvent.Type.TouchBegin, QEvent.Type.TouchUpdate, QEvent.Type.TouchEnd):
            self._ignore_wheel = True
            self.touchEvent(e)
            return True
        elif e.type() == QEvent.Type.Wheel:
            if self._ignore_wheel:
                self._ignore_wheel = False
                return True
        return super().event(e)

    def touchEvent(self, event: QTouchEvent):
        points = event.points()
        if not points:
            return

        if event.type() == QEvent.Type.TouchEnd:
            self._touch_points.clear()
            self._pinch_start_center = None
            self._pinch_start_dist = 0.0
            event.accept()
            return

        if len(points) == 1:
            p = points[0]
            # 如果是刚检测到单指，或者从双指变回单指
            if p.id() not in self._touch_points or self._pinch_start_center is not None:
                self._touch_points.clear()
                self._touch_points[p.id()] = p.position()
                self._pinch_start_center = None
            else:
                last = self._touch_points[p.id()]
                dx = p.position().x() - last.x()
                dy = p.position().y() - last.y()
                self.view_rot_x += dy * 0.5
                self.view_rot_x = max(-90.0, min(0.0, self.view_rot_x))  # 限制垂直视角的翻转
                self.view_rot_z += dx * 0.5
                self._touch_points[p.id()] = p.position()
                self.update()

        elif len(points) == 2:
            p1, p2 = points[0], points[1]
            dx_p = p1.position().x() - p2.position().x()
            dy_p = p1.position().y() - p2.position().y()
            dist = (dx_p**2 + dy_p**2)**0.5
            center_x = (p1.position().x() + p2.position().x()) / 2.0
            center_y = (p1.position().y() + p2.position().y()) / 2.0

            # 初始化双指状态
            if self._pinch_start_center is None:
                self._pinch_start_dist = dist
                self._pinch_start_zoom = self.view_zoom
                self._pinch_start_center = (center_x, center_y)
                self._pinch_start_trans = (self.view_trans_x, self.view_trans_y)
                self._touch_points.clear() # 清除单指记录
            else:
                # 缩放 (双指捏合)
                if self._pinch_start_dist > 0:
                    scale = dist / self._pinch_start_dist
                    self.view_zoom = self._pinch_start_zoom * (1 / scale)
                
                # 平移 (双指并行移动)
                dcx = center_x - self._pinch_start_center[0]
                dcy = center_y - self._pinch_start_center[1]
                pan_speed = abs(self.view_zoom) * 0.002
                self.view_trans_x = self._pinch_start_trans[0] + dcx * pan_speed
                self.view_trans_y = self._pinch_start_trans[1] - dcy * pan_speed
                
                self.update()
        event.accept()