
import sys
import os
import struct
import time
from panda3d.core import (
    Filename, loadPrcFileData, GeomVertexFormat,
    GeomVertexWriter, InternalName, Shader, Texture, SamplerState,
    Vec3, Vec4, Point2, Point3, LMatrix4f, ShaderBuffer, GeomEnums, OmniBoundingVolume, Quat,
    TransparencyAttrib, BoundingSphere, NodePath,
    GraphicsEngine, WindowProperties, FrameBufferProperties,
    GraphicsPipe, GraphicsOutput, Camera, DisplayRegion, OrthographicLens,
    BoundingBox
)

import p3dimgui.backend as p3dimgui_backend
import p3dimgui.shaders as p3dimgui_shaders
from imgui_bundle import imgui
from rpcore.effect import Effect

# Work around p3dimgui import-order issue where backend may import an unrelated
# top-level "shaders" module and miss these globals.
if not hasattr(p3dimgui_backend, "VERT_SHADER"):
    p3dimgui_backend.VERT_SHADER = p3dimgui_shaders.VERT_SHADER
if not hasattr(p3dimgui_backend, "FRAG_SHADER"):
    p3dimgui_backend.FRAG_SHADER = p3dimgui_shaders.FRAG_SHADER

ImGuiBackend = p3dimgui_backend.ImGuiBackend

from .ssbo_controller import ObjectController

class SSBOEditor:
    """
    SSBO Editor Component
    ====================
    Encapsulates the SSBO rendering, ImGui editor, and interaction logic.
    Can be integrated into any ShowBase application using RenderPipeline.
    """
    
    def __init__(self, base_app, render_pipeline, model_path=None, font_path=None):
        self.base = base_app
        self.rp = render_pipeline
        self.controller = None
        self.model = None
        self.ssbo = None
        self.font_path = font_path
        # Picking resources may be created later when a model is loaded.
        self.pick_buffer = None
        self.pick_texture = None
        self.pick_cam = None
        self.pick_cam_np = None
        self.pick_lens = None
        
        # Internal State
        self.selected_name = None
        self.selected_ids = []
        self.search_text = ""
        self.last_search_text = None
        self.filtered_nodes = []
        self.debug_mode = False
        self.keys = {}
        self._ssbo_transform_active = False
        self._ssbo_selected_local_indices = []
        self._ssbo_transform_snapshot = None
        self._ssbo_gizmo_proxy = None
        self._ssbo_proxy_start = {"pos": None, "quat": None, "scale": None}
        self._bound_transform_gizmo = None
        
        # Initialize ImGui Backend if not already present
        if not hasattr(self.base, 'imgui_backend'):
            print("[SSBOEditor] Initializing ImGui Backend...")
            self.base.imgui_backend = ImGuiBackend()
        
        self.load_font()
        
        # Register Events
        self.base.accept("imgui-new-frame", self.draw_imgui)
        self.base.accept("f", self.focus_on_selected)
        self.base.accept("d", self.toggle_debug)
        self.base.accept("mouse1", self.on_mouse_click)
        
        # Register Input Tasks
        for key in ['arrow_up', 'arrow_down', 'arrow_left', 'arrow_right', 'z', 'x']:
            self.base.accept(key, self.keys.__setitem__, [key, True])
            self.base.accept(f"{key}-up", self.keys.__setitem__, [key, False])
            
        # Add Tasks
        self.base.taskMgr.add(self.update_task, "update_task")
        
        # Load Model if provided
        if model_path:
            self.load_model(model_path)

    def load_font(self):
        """Load custom font for ImGui"""
        io = imgui.get_io()
        
        # Load Chinese Glyph Ranges
        glyph_ranges = None
        try:
            if hasattr(io.fonts, 'get_glyph_ranges_chinese_full'):
                glyph_ranges = io.fonts.get_glyph_ranges_chinese_full()
            elif hasattr(io.fonts, 'get_glyph_ranges_chinese_simplified_common'):
                glyph_ranges = io.fonts.get_glyph_ranges_chinese_simplified_common()
        except Exception as e:
            print(f"[SSBOEditor] Warning: Could not get Chinese glyph ranges: {e}")

        try:
            if self.font_path and os.path.exists(self.font_path):
                io.fonts.clear()
                # If glyph_ranges is None, it uses default (Basic Latin)
                if glyph_ranges:
                    io.fonts.add_font_from_file_ttf(self.font_path, 18.0, glyph_ranges=glyph_ranges)
                else:
                    io.fonts.add_font_from_file_ttf(self.font_path, 18.0)
            else:
                # Fallback to default or common font
                default_font = os.path.join(os.path.dirname(os.path.dirname(__file__)), "font", "msyh.ttc")
                if os.path.exists(default_font):
                    io.fonts.clear()
                    io.fonts.add_font_from_file_ttf(default_font, 18.0, glyph_ranges=glyph_ranges)
                else:
                    io.fonts.clear()
                    io.fonts.add_font_default()
        except Exception as e:
            print(f"[SSBOEditor] Font load error: {e}")
            io.fonts.clear()
            io.fonts.add_font_default()

    def load_model(self, model_path):
        """Load and process a model — NO custom shader, uses RP default rendering."""
        print(f"[SSBOEditor] Loading model: {model_path}")
        fn = Filename.fromOsSpecific(model_path)
        self.model = self.base.loader.loadModel(fn)
        
        self.controller = ObjectController()
        count = self.controller.bake_ids_and_collect(self.model)
        self._ssbo_transform_active = False
        self._ssbo_selected_local_indices = []
        self._ssbo_transform_snapshot = None
        self._cleanup_ssbo_proxy()
        
        self.model.reparent_to(self.base.render)
        
        # NO rp.set_effect() — use RP default rendering for max FPS
        # NO SSBO creation — vertex positions are baked
        
        # Setup GPU Picking (uses simple vertex-color shader)
        self.setup_gpu_picking()
        
        print(f"[SSBOEditor] Model loaded. Total objects: {count}")

    # No custom effect needed — RP default rendering for maximum FPS

    def _inject_ssbo_into_shadow_state(self, effect_path):
        """Inject SSBO inputs into RP shadow tag state"""
        try:
            if not hasattr(self.rp.tag_mgr, 'containers'): return

            shadow_container = self.rp.tag_mgr.containers.get("shadow")
            if not shadow_container: return

            tag_value = self.model.get_tag(shadow_container.tag_name)
            if not tag_value: return
            
            effect = Effect.load(effect_path, {})
            if effect is None: return
            
            shadow_shader = effect.get_shader_obj("shadow")
            if shadow_shader is None: return
            
            # Since inputs are now on Nodes (Chunks), we just need to ensure the shader is applied.
            # extra_inputs is no longer needed if the inputs are on the nodes themselves?
            # Wait, RP might override state.
            # But specific shader inputs on NodePath have priority over State inputs usually?
            # Let's try applying without extra inputs first.
            
            self.rp.tag_mgr.apply_state(
                "shadow", self.model, shadow_shader,
                tag_value, 65)
            
            print(f"[SSBO Shadow] Re-applied shadow state (tag='{tag_value}')")
        except Exception as e:
            print(f"[SSBO Shadow] Error injecting shadow state: {e}")

    def setup_gpu_picking(self):
        """Setup GPU Picking (Basic implementation)"""
        # ... (Buffer setup code remains same) ...
        win_props = WindowProperties()
        win_props.set_size(1, 1)
        fb_props = FrameBufferProperties()
        fb_props.set_rgba_bits(8, 8, 8, 8)
        fb_props.set_depth_bits(16)
        
        self.pick_buffer = self.base.graphicsEngine.make_output(
            self.base.pipe, "pick_buffer", -100,
            fb_props, win_props,
            GraphicsPipe.BF_refuse_window,
            self.base.win.get_gsg(), self.base.win
        )
        
        if not self.pick_buffer:
            print("[GPU Picking] Failed to create buffer!")
            return
            
        self.pick_texture = Texture()
        self.pick_texture.set_minfilter(Texture.FT_nearest)
        self.pick_texture.set_magfilter(Texture.FT_nearest)
        self.pick_buffer.add_render_texture(self.pick_texture, GraphicsOutput.RTM_copy_ram)
        
        self.pick_cam = Camera("pick_camera")
        self.pick_cam_np = self.base.cam.attach_new_node(self.pick_cam)
        self.pick_lens = self.base.camLens.make_copy()
        self.pick_cam.set_lens(self.pick_lens)
        
        dr = self.pick_buffer.make_display_region()
        dr.set_camera(self.pick_cam_np)
        
        # Load pick shader
        current_dir = os.path.dirname(os.path.abspath(__file__))
        pick_vert = os.path.join(current_dir, "shaders", "pick_id.vert")
        pick_frag = os.path.join(current_dir, "shaders", "pick_id.frag")
        
        pick_vert = Filename.fromOsSpecific(pick_vert).getFullpath()
        pick_frag = Filename.fromOsSpecific(pick_frag).getFullpath()

        try:
            pick_shader = Shader.load(
                Shader.SL_GLSL,
                pick_vert,
                pick_frag
            )
            pick_scene = getattr(self.controller, "pick_model", None) if self.controller else None
            if pick_scene is None:
                pick_scene = self.model
            self.pick_cam.set_scene(pick_scene)
            initial_state = NodePath("initial")
            initial_state.set_shader(pick_shader, 100)
            # Remove global SSBO input, Chunks have their own inputs
            # initial_state.set_shader_input("transforms", ssbo) 
            self.pick_cam.set_initial_state(initial_state.get_state())
        except Exception as e:
            print(f"[GPU Picking] Warning: pick shaders failed to load: {e}")
            print("Picking disabled.")
            return
        
        self.pick_buffer.set_active(False)
        self.pick_buffer.set_clear_color(Vec4(0, 0, 0, 0))
        self.pick_buffer.set_clear_color_active(True)

    def pick_object(self, mx, my):
        if (not self.pick_buffer or not self.pick_texture or not self.pick_lens or
                not self.controller or not self.model):
            return False
        
        self.pick_lens.set_fov(0.1)
        self.pick_lens.set_film_offset(0, 0)
        self.pick_cam.set_lens(self.pick_lens)
        
        near_point = Point3()
        far_point = Point3()
        self.base.camLens.extrude(Point2(mx, my), near_point, far_point)
        
        self.pick_cam_np.set_pos(0, 0, 0)
        self.pick_cam_np.look_at(far_point)
        
        self.pick_buffer.set_active(True)
        self.base.graphicsEngine.render_frame()
        self.pick_buffer.set_active(False)
        self.base.graphicsEngine.extract_texture_data(
            self.pick_texture, self.base.win.get_gsg()
        )
        
        ram_image = self.pick_texture.get_ram_image_as("RGBA")
        if ram_image:
            data = memoryview(ram_image)
            if len(data) >= 4:
                r, g, b, a = data[0], data[1], data[2], data[3]
                if a > 0:
                    hit_id = r + (g << 8)
                    node_key = self.controller.id_to_name.get(hit_id)
                    if node_key:
                        print(f"[Pick] Hit: ID={hit_id} -> {node_key}")
                        self.select_node(node_key)
                        return True
        
        self.selected_name = None
        self.selected_ids = []
        return False


    def on_mouse_click(self):
        io = imgui.get_io()
        if io.want_capture_mouse:
            return
        if self.base.mouseWatcherNode.has_mouse():
            mpos = self.base.mouseWatcherNode.get_mouse()
            # If clicking gizmo, skip SSBO pick.
            if self._try_start_gizmo_drag(mpos.x, mpos.y):
                return
            prev_selected = self.selected_name
            hit = self.pick_object(mpos.x, mpos.y)
            # SSBO miss must clear current selection.
            if not hit:
                self._sync_selection_none()
                # Always fallback to legacy ray pick when SSBO misses.
                # This keeps scene selection usable if SSBO ID mapping is incomplete.
                self._fallback_legacy_pick(mpos.x, mpos.y)
            elif prev_selected != self.selected_name:
                # Ensure selection visuals refresh when SSBO selection changes.
                self._sync_selection_from_key(self.selected_name)

    def toggle_debug(self):
        self.debug_mode = not self.debug_mode

    def clear_selection(self):
        pass  # No selection mask texture needed without custom shader

    def update_selection_mask(self):
        pass  # No selection mask texture needed without custom shader

    def select_node(self, key):
        if not self.controller or key not in self.controller.name_to_ids:
            return
        self.selected_name = key
        self.selected_ids = self.controller.name_to_ids.get(key, [])
        self._sync_selection_from_key(key)

    def _sync_selection_from_key(self, key):
        """Sync SSBO picked key to legacy SelectionSystem."""
        try:
            if hasattr(self.base, "selection") and self.base.selection:
                kind, target = self._resolve_ssbo_selection_target(key)
                if kind == "proxy":
                    target_np = target
                else:
                    target_np = target if target is not None else self.model
                if target_np is None or target_np.isEmpty():
                    target_np = self.model
                self.base.selection.updateSelection(target_np)
        except Exception as e:
            print(f"[SSBOEditor] selection sync failed: {e}")

    def _sync_selection_none(self):
        """Clear legacy SelectionSystem selection."""
        try:
            self._ssbo_transform_active = False
            self._ssbo_selected_local_indices = []
            self._ssbo_transform_snapshot = None
            self._cleanup_ssbo_proxy()
            if hasattr(self.base, "selection") and self.base.selection:
                self.base.selection.updateSelection(None)
        except Exception as e:
            print(f"[SSBOEditor] clear selection sync failed: {e}")

    def bind_transform_gizmo(self, transform_gizmo):
        """Bind TransformGizmo drag hooks so SSBO sub-object transforms can follow gizmo."""
        self._bound_transform_gizmo = transform_gizmo
        if not transform_gizmo:
            return
        hooks = {
            "move": {
                "drag_start": [self._on_ssbo_gizmo_drag_start],
                "drag_move": [self._on_ssbo_gizmo_drag_move],
                "drag_end": [self._on_ssbo_gizmo_drag_end],
            },
            "rotate": {
                "drag_start": [self._on_ssbo_gizmo_drag_start],
                "drag_move": [self._on_ssbo_gizmo_drag_move],
                "drag_end": [self._on_ssbo_gizmo_drag_end],
            },
            "scale": {
                "drag_start": [self._on_ssbo_gizmo_drag_start],
                "drag_move": [self._on_ssbo_gizmo_drag_move],
                "drag_end": [self._on_ssbo_gizmo_drag_end],
            },
        }
        try:
            if hasattr(transform_gizmo, "set_event_hooks"):
                transform_gizmo.set_event_hooks(hooks, replace=False)
            print("[SSBOEditor] TransformGizmo hooks bound")
        except Exception as e:
            print(f"[SSBOEditor] bind transform gizmo failed: {e}")

    def _resolve_ssbo_selection_target(self, key):
        """Resolve selected SSBO key to proxy node (preferred) or regular node."""
        self._ssbo_transform_active = False
        self._ssbo_transform_snapshot = None
        self._ssbo_selected_local_indices = []

        if not self.controller or not key:
            return "node", self.model
        global_ids = self.controller.name_to_ids.get(key, [])
        local_indices = self.controller.get_local_indices_from_global_ids(global_ids)
        self._ssbo_selected_local_indices = local_indices
        if local_indices:
            print(f"[SSBOEditor] selection locals={len(local_indices)} key={key}")
            center = self.controller.get_selection_center(local_indices)
            proxy = self._ensure_ssbo_proxy(center)
            return "proxy", proxy
        target_np = self.controller.key_to_node.get(key)
        if target_np is None or target_np.isEmpty():
            target_np = self.model
        return "node", target_np

    def _ensure_ssbo_proxy(self, center):
        if self._ssbo_gizmo_proxy is None or self._ssbo_gizmo_proxy.isEmpty():
            self._ssbo_gizmo_proxy = self.base.render.attach_new_node("ssbo_transform_proxy")
            self._ssbo_gizmo_proxy.setTag("is_ssbo_proxy", "1")
        self._ssbo_gizmo_proxy.set_pos(center)
        self._ssbo_gizmo_proxy.set_hpr(0, 0, 0)
        self._ssbo_gizmo_proxy.set_scale(1, 1, 1)
        return self._ssbo_gizmo_proxy

    def _cleanup_ssbo_proxy(self):
        if self._ssbo_gizmo_proxy and not self._ssbo_gizmo_proxy.isEmpty():
            self._ssbo_gizmo_proxy.removeNode()
        self._ssbo_gizmo_proxy = None

    def _on_ssbo_gizmo_drag_start(self, payload):
        try:
            target = payload.get("target") if payload else None
            if not target or target != self._ssbo_gizmo_proxy:
                self._ssbo_transform_active = False
                return
            if not self.controller or not self._ssbo_selected_local_indices:
                self._ssbo_transform_active = False
                return
            self._ssbo_transform_snapshot = self.controller.begin_transform_session(
                self._ssbo_selected_local_indices
            )
            self._ssbo_proxy_start = {
                "pos": Vec3(target.getPos(self.base.render)),
                "quat": Quat(target.getQuat(self.base.render)),
                "scale": Vec3(target.getScale()),
            }
            self._ssbo_transform_active = True
            print(f"[SSBOEditor] drag_start locals={len(self._ssbo_selected_local_indices)}")
        except Exception as e:
            self._ssbo_transform_active = False
            print(f"[SSBOEditor] drag_start bridge failed: {e}")

    def _on_ssbo_gizmo_drag_move(self, payload):
        try:
            if not self._ssbo_transform_active:
                return
            target = payload.get("target") if payload else None
            if not target or target != self._ssbo_gizmo_proxy:
                return
            start_pos = self._ssbo_proxy_start.get("pos")
            start_quat = self._ssbo_proxy_start.get("quat")
            start_scale = self._ssbo_proxy_start.get("scale")
            if start_pos is None or start_quat is None or start_scale is None:
                return

            curr_pos = Vec3(target.getPos(self.base.render))
            curr_quat = Quat(target.getQuat(self.base.render))
            curr_scale = Vec3(target.getScale())

            delta_pos = curr_pos - start_pos
            inv_start_quat = Quat(start_quat)
            inv_start_quat.invertInPlace()
            delta_quat = curr_quat * inv_start_quat
            delta_scale = Vec3(
                curr_scale.x / start_scale.x if abs(start_scale.x) > 1e-8 else 1.0,
                curr_scale.y / start_scale.y if abs(start_scale.y) > 1e-8 else 1.0,
                curr_scale.z / start_scale.z if abs(start_scale.z) > 1e-8 else 1.0,
            )

            self.controller.apply_transform_session(
                self._ssbo_transform_snapshot,
                delta_pos,
                delta_quat,
                delta_scale,
            )
        except Exception as e:
            print(f"[SSBOEditor] drag_move bridge failed: {e}")

    def _on_ssbo_gizmo_drag_end(self, payload):
        try:
            if self._ssbo_transform_active:
                print(f"[SSBOEditor] drag_end locals={len(self._ssbo_selected_local_indices)}")
            self._ssbo_transform_active = False
            self._ssbo_transform_snapshot = None
        except Exception as e:
            print(f"[SSBOEditor] drag_end bridge failed: {e}")

    def _fallback_legacy_pick(self, mx, my):
        """Fallback to legacy ray picking when SSBO misses."""
        try:
            if not hasattr(self.base, "event_handler") or not self.base.event_handler:
                return
            win_w, win_h = self.base.win.getSize()
            x = (mx + 1.0) * 0.5 * win_w
            y = (1.0 - my) * 0.5 * win_h
            self.base.event_handler.mousePressEventLeft({"x": x, "y": y})
        except Exception as e:
            print(f"[SSBOEditor] legacy fallback pick failed: {e}")

    def _try_start_gizmo_drag(self, mouse_x=None, mouse_y=None):
        """Try to start gizmo drag using the existing SelectionSystem pipeline."""
        try:
            new_transform = getattr(self.base, "newTransform", None)
            if (
                new_transform is not None and
                mouse_x is not None and
                mouse_y is not None and
                self._is_mouse_on_new_gizmo(new_transform, mouse_x, mouse_y)
            ):
                return True
            selection = getattr(self.base, "selection", None)
            if not selection or not selection.gizmo:
                return False
            win_w, win_h = self.base.win.getSize()
            mpos = self.base.mouseWatcherNode.get_mouse()
            x = (mpos.x + 1.0) * 0.5 * win_w
            y = (1.0 - mpos.y) * 0.5 * win_h

            axis = selection.gizmoHighlightAxis or selection.checkGizmoClick(x, y)
            if axis:
                selection.startGizmoDrag(axis, x, y)
                return True
        except Exception as e:
            print(f"[SSBOEditor] gizmo drag start failed: {e}")
        return False

    def _is_mouse_on_new_gizmo(self, new_transform, mouse_x, mouse_y):
        """Refresh and query hover state for TransformGizmo on current click position."""
        try:
            mouse_pos = Point3(mouse_x, mouse_y, 0.0)
            for gizmo_name in ("move_gizmo", "rotate_gizmo", "scale_gizmo"):
                gizmo = getattr(new_transform, gizmo_name, None)
                if not gizmo or not getattr(gizmo, "attached", False):
                    continue
                hover_updater = getattr(gizmo, "_update_hover_highlight", None)
                if callable(hover_updater):
                    hover_updater(mouse_pos)
            return bool(getattr(new_transform, "is_hovering", False))
        except Exception as e:
            print(f"[SSBOEditor] new gizmo hover check failed: {e}")
            return False

    def focus_on_selected(self):
        if self.selected_name and self.selected_ids:
            first_id = self.selected_ids[0]
            pos = self.controller.get_world_pos(first_id)
            dist = 100
            self.base.camera.set_pos(pos.x, pos.y - dist, pos.z + dist * 0.5)
            self.base.camera.look_at(pos)

    def draw_imgui(self):
        if not self.controller: return

        imgui.set_next_window_pos((10, 10), imgui.Cond_.first_use_ever)
        imgui.set_next_window_size((350, 600), imgui.Cond_.first_use_ever)
        
        expanded, opened = imgui.begin("Scene Tree (Component)")
        if expanded:
            imgui.text(f"FPS: {globalClock.getAverageFrameRate():.1f}")
            imgui.separator()
            
            changed, self.search_text = imgui.input_text("Search", self.search_text, 256)
            
            if imgui.begin_child("ObjectList", (0, 380), child_flags=imgui.ChildFlags_.borders):
                if self.search_text != self.last_search_text:
                    self.last_search_text = self.search_text
                    search_lower = self.search_text.lower()
                    self.filtered_nodes = []
                    for key in self.controller.node_list:
                        display = self.controller.display_names.get(key, key.split('/')[-1])
                        if not search_lower or (search_lower in display.lower() or search_lower in key.lower()):
                            geom_count = len(self.controller.name_to_ids.get(key, []))
                            self.filtered_nodes.append((key, display, geom_count))
                
                # If list is empty initially (no search), show all
                if not self.search_text and not self.filtered_nodes:
                     if len(self.filtered_nodes) != len(self.controller.node_list):
                        self.filtered_nodes = [(k, self.controller.display_names.get(k, k), len(self.controller.name_to_ids.get(k,[]))) for k in self.controller.node_list]

                count = len(self.filtered_nodes)
                clipper = imgui.ListClipper()
                clipper.begin(count)
                while clipper.step():
                    for i in range(clipper.display_start, clipper.display_end):
                        key, display, geom_count = self.filtered_nodes[i]
                        label = f"{display} ({geom_count})"
                        is_selected = (key == self.selected_name)
                        if imgui.selectable(label, is_selected)[0]:
                            self.select_node(key)
                imgui.end_child()
                
            imgui.separator()
            if self.selected_name:
                imgui.text_colored((1, 0.8, 0.2, 1), f"Selected: {self.selected_name}")
                if imgui.button("Focus (F)"): self.focus_on_selected()
        imgui.end()

    # swap_transforms_task removed - motion blur disabled for performance

    def update_task(self, task):
        dt = globalClock.getDt()
        io = imgui.get_io()
            
        if io.want_capture_keyboard: return task.cont
        
        if self.selected_ids and self.controller:
            speed = 50 * dt
            acc = Vec3(0, 0, 0)
            if self.keys.get('arrow_up'): acc.z += speed
            if self.keys.get('arrow_down'): acc.z -= speed
            if self.keys.get('arrow_left'): acc.x -= speed
            if self.keys.get('arrow_right'): acc.x += speed
            if self.keys.get('z'): acc.y += speed
            if self.keys.get('x'): acc.y -= speed
            
            if acc.length_squared() > 0:
                for idx in self.selected_ids:
                    self.controller.move_object(idx, acc)
        
        return task.cont
