import numpy as np
import math
import bisect
from PyQt6.QtOpenGLWidgets import QOpenGLWidget
from PyQt6.QtCore import Qt, QEvent, QThread, pyqtSignal
from PyQt6.QtGui import QTouchEvent, QSurfaceFormat, QMatrix4x4, QVector4D
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 = {'TRAVEL': False}
        segment_types = []          # 每条线段对应的类型名称（长度为线段数）

        x = y = z = e = 0.0
        vertex_idx = 0
        feature_type = 'OTHER'
        current_segment_type = 'OTHER'
        segment_start = 0
        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)]

        lod_steps = [2, 4]
        lod_levels = {1: None}
        for s in lod_steps:
            lod_levels[s] = {'points': [], 'colors': [], 'count': 0}

        grid_size = 10
        grid_segments = {}

        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

                        # 记录当前线段的类型
                        segment_types.append(seg_type)

                        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])

                        for step in lod_steps:
                            if vertex_idx % (step * 2) == 0:
                                lod_levels[step]['points'].extend([x, y, z, new_x, new_y, new_z])
                                lod_levels[step]['colors'].extend([*c, *c])
                                lod_levels[step]['count'] += 2

                        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

            if vertex_idx > 0 and max_x > min_x and max_y > min_y:
                dx = (max_x - min_x) / grid_size
                dy = (max_y - min_y) / grid_size
                for i in range(grid_size):
                    for j in range(grid_size):
                        grid_segments[(i, j)] = []
                # 填充网格，同时每个格子存储 (线段起始顶点索引, 类型)
                for seg_idx in range(0, vertex_idx, 2):
                    px = points[seg_idx * 3]
                    py = points[seg_idx * 3 + 1]
                    gx = min(int((px - min_x) / dx), grid_size - 1)
                    gy = min(int((py - min_y) / dy), grid_size - 1)
                    seg_type = segment_types[seg_idx // 2]
                    grid_segments[(gx, gy)].append((seg_idx, seg_type))

            lod_levels[1] = {'points': points, 'colors': colors, 'count': vertex_idx}

            result = {
                'lod_levels': lod_levels,
                '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,
                'grid_segments': grid_segments,
                'grid_size': grid_size,
                'bbox': (min_x, max_x, min_y, max_y, min_z, max_z)
            }
            self.finished.emit(result)
        except Exception as e:
            print("ParseGCode Error:", e)
            self.finished.emit({})


class GCodeViewerWidget(QOpenGLWidget):
    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),
        'WALL-INNER':        (0.25, 0.50, 0.81),
        'FILL':              (0.80, 0.75, 0.29),
        'SKIN':              (0.62, 0.38, 0.70),
        'SUPPORT':           (0.34, 0.70, 0.34),
        'SUPPORT-INTERFACE': (0.17, 0.42, 0.17),
        'SKIRT':             (0.00, 1.00, 1.00),
        'OTHER':             (0.67, 0.67, 0.67),
        'TRAVEL':            (0.25, 0.31, 0.38),
    }

    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;
        }
    """

    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):
        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.lod_levels = {}
        self.vbos = {}
        self.vertex_count = 0

        self.type_segments = {}
        self.type_visibility = {}

        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_vertices = 0
        self.layer_map = [(0, 0)]

        self.center_x = 110.0
        self.center_y = 110.0
        self.center_z = 0.0

        self.grid_segments = {}       # 格子 -> [(seg_idx, type_name)]
        self.grid_size = 10
        self.bbox = (0, 0, 0, 0, 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._rot_sensitivity = 0.1

        self.shader_program = None
        self.aPos_location = None
        self.aColor_location = None
        self.uMVP_location = None
        self.uDarken_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.lod_levels = result['lod_levels']
        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.grid_segments = result['grid_segments']
        self.grid_size = result['grid_size']
        self.bbox = result['bbox']
        self.progress_vertices = self.vertex_count
        self.vbos.clear()
        self.update()

    def update_by_filepos(self, filepos: int, is_printing: bool = True):
        if not self.layer_map:
            return
        if not is_printing:
            target_vertices = self.vertex_count
        else:
            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]
        if target_vertices != self.progress_vertices:
            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()

    def initializeGL(self):
        import OpenGL.GL as gl
        gl.glClearColor(0.15, 0.15, 0.15, 1.0)
        gl.glEnable(gl.GL_DEPTH_TEST)
        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()
        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")

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

    def _build_mvp(self):
        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 _is_tile_visible(self, gx, gy, mvp_mat):
        min_x, max_x, min_y, max_y, min_z, max_z = self.bbox
        if max_x <= min_x or max_y <= min_y:
            return True
        dx = (max_x - min_x) / self.grid_size
        dy = (max_y - min_y) / self.grid_size
        cx = min_x + dx * (gx + 0.5)
        cy = min_y + dy * (gy + 0.5)
        cz = (min_z + max_z) * 0.5
        v = mvp_mat * QVector4D(cx, cy, cz, 1.0)
        if v.w() == 0:
            return False
        ndc_x = v.x() / v.w()
        ndc_y = v.y() / v.w()
        ndc_z = v.z() / v.w()
        return -1.5 <= ndc_x <= 1.5 and -1.5 <= ndc_y <= 1.5 and -2.0 <= ndc_z <= 2.0

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

        if self.vertex_count == 0:
            return

        actual_display_verts = min(self.vertex_count, self.progress_vertices)

        if actual_display_verts > 500000:
            desired_step = 4
        elif actual_display_verts > 200000:
            desired_step = 2
        else:
            desired_step = 1

        if abs(self.view_zoom) > 600 and desired_step < 4:
            desired_step = 4
        elif abs(self.view_zoom) > 300 and desired_step < 2:
            desired_step = 2

        if desired_step not in self.vbos:
            self._create_vbo_for_step(desired_step)

        vbo_vertices, vbo_colors, vbo_count = self.vbos[desired_step]
        if vbo_count == 0:
            return

        scale = desired_step * 2
        lod_progress = actual_display_verts // scale
        if lod_progress > vbo_count:
            lod_progress = vbo_count

        self.shader_program.bind()
        mvp_mat = self._build_mvp()
        self.shader_program.setUniformValue(self.uMVP_location, mvp_mat)

        vbo_vertices.bind()
        self.shader_program.setAttributeBuffer(self.aPos_location, gl.GL_FLOAT, 0, 3, 0)
        self.shader_program.enableAttributeArray(self.aPos_location)

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

        gl.glDepthFunc(gl.GL_LEQUAL)

        use_culling = (desired_step == 1 and actual_display_verts > 100000 and self.grid_segments)

        if use_culling:
            # 空间剔除分支：现在每个格子条目是 (seg_idx, type_name)
            for pass_idx in range(2):
                if pass_idx == 0:
                    gl.glLineWidth(6.0)
                    self.shader_program.setUniformValue(self.uDarken_location, 0.8)
                else:
                    gl.glLineWidth(3.0)
                    self.shader_program.setUniformValue(self.uDarken_location, 1.0)

                for (gx, gy), seg_entries in self.grid_segments.items():
                    if not self._is_tile_visible(gx, gy, mvp_mat):
                        continue
                    for seg_idx, seg_type in seg_entries:
                        # 跳过不可见类型
                        if not self.type_visibility.get(seg_type, True):
                            continue
                        if seg_idx < actual_display_verts:
                            gl.glDrawArrays(gl.GL_LINES, seg_idx, 2)
        else:
            # 普通渲染分支
            for pass_idx in range(2):
                if pass_idx == 0:
                    gl.glLineWidth(6.0)
                    self.shader_program.setUniformValue(self.uDarken_location, 0.8)
                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:
                        lod_start = start // scale
                        lod_count = length // scale
                        if lod_start >= vbo_count:
                            continue
                        end = min(lod_start + lod_count, lod_progress)
                        actual_count = end - lod_start
                        if actual_count > 0:
                            gl.glDrawArrays(gl.GL_LINES, lod_start, actual_count)

        gl.glDepthFunc(gl.GL_LESS)
        self.shader_program.disableAttributeArray(self.aPos_location)
        self.shader_program.disableAttributeArray(self.aColor_location)
        vbo_vertices.release()
        vbo_colors.release()
        self.shader_program.release()

    def _create_vbo_for_step(self, step):
        import OpenGL.GL as gl
        lod_data = self.lod_levels.get(step)
        if not lod_data:
            lod_data = self.lod_levels[1]
            step = 1

        verts = np.array(lod_data['points'], dtype=np.float32)
        cols = np.array(lod_data['colors'], dtype=np.float32)

        vbo_v = QOpenGLBuffer(QOpenGLBuffer.Type.VertexBuffer)
        vbo_c = QOpenGLBuffer(QOpenGLBuffer.Type.VertexBuffer)
        vbo_v.create()
        vbo_v.bind()
        vbo_v.allocate(verts.tobytes(), verts.nbytes)
        vbo_v.release()
        vbo_c.create()
        vbo_c.bind()
        vbo_c.allocate(cols.tobytes(), cols.nbytes)
        vbo_c.release()

        self.vbos[step] = (vbo_v, vbo_c, lod_data['count'])

    # ── 交互部分保持不变 ──
    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 * self._rot_sensitivity
            self.view_rot_x = max(-90.0, min(0.0, self.view_rot_x))
            self.view_rot_z += dx * self._rot_sensitivity
        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 * self._rot_sensitivity
                self.view_rot_x = max(-90.0, min(0.0, self.view_rot_x))
                self.view_rot_z += dx * self._rot_sensitivity
                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 = math.hypot(dx_p, dy_p)
            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()