# ============================
# sphere_renderer.py — GPU Impostor Spheres
# ============================
from __future__ import annotations
from OpenGL.GL import *
import numpy as np
from typing import Optional, Tuple, List, Dict
from .shaders import SPHERE_VS, SPHERE_FS
from .utils import link_program

class SphereRenderer:
    """
    High-performance GPU renderer for atomic spheres using
    impostor billboards and raycasting. Supports per-atom and per-element
    material properties through TBOs.
    """
    def __init__(self):
        self.prog = link_program(SPHERE_VS, SPHERE_FS)
        
        # Attribute Buffers (Screen Quad)
        self.quad_vao = glGenVertexArrays(1)
        glBindVertexArray(self.quad_vao)
        quad = np.array([[-1,-1],[1,-1],[-1,1],[1,1]], dtype=np.float32)
        vbo = glGenBuffers(1)
        glBindBuffer(GL_ARRAY_BUFFER, vbo)
        glBufferData(GL_ARRAY_BUFFER, quad.nbytes, quad, GL_STATIC_DRAW)
        glEnableVertexAttribArray(0)
        glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 0, None)
        glBindVertexArray(0)
        
        # Texture Buffer Objects (TBOs)
        self.pos_bo, self.pos_tbo = glGenBuffers(1), glGenTextures(1)
        self.col_bo, self.col_tbo = glGenBuffers(1), glGenTextures(1)
        self.rep_bo, self.rep_tbo = glGenBuffers(1), glGenTextures(1)
        self.mat_bo, self.mat_tbo = glGenBuffers(1), glGenTextures(1)
        self.sel_bo, self.sel_tbo = glGenBuffers(1), glGenTextures(1)

        # Shadow Mapping Pass
        self.shadow_size = 4096
        self.shadow_fbo = glGenFramebuffers(1)
        self.shadow_tex = glGenTextures(1)
        self.shadow_compare = True # Hardware comparison by default
        glBindTexture(GL_TEXTURE_2D, self.shadow_tex)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT, self.shadow_size, self.shadow_size, 0, GL_DEPTH_COMPONENT, GL_FLOAT, None)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_BORDER)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_BORDER)
        glTexParameterfv(GL_TEXTURE_2D, GL_TEXTURE_BORDER_COLOR, [1.0, 1.0, 1.0, 1.0])
        
        # Add shadow comparison mode for HW-accelerated percentage-closer filtering (PCF)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_MODE, GL_COMPARE_REF_TO_TEXTURE)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_FUNC, GL_LEQUAL)

        glBindFramebuffer(GL_FRAMEBUFFER, self.shadow_fbo)
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_TEXTURE_2D, self.shadow_tex, 0)
        glDrawBuffer(GL_NONE)
        glReadBuffer(GL_NONE)
        
        if glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE:
            print("Warning: Shadow FBO incomplete!")
        glBindFramebuffer(GL_FRAMEBUFFER, 0)
        
        # Shadow Shader (Specific for depth pass)
        from .shaders import SHADOW_VS, SHADOW_FS
        self.shadow_prog = link_program(SHADOW_VS, SHADOW_FS)
        
        # Material Data TBOs (M1: Amb/Spec/Shin/Rim, M2: Sheen/CC/Env/Gloss)
        self.mat_bo, self.mat_tbo = glGenBuffers(1), glGenTextures(1)
        self.mat_bo2, self.mat_tbo2 = glGenBuffers(1), glGenTextures(1)
        
        # Selection Pass Data
        self.sel_bo, self.sel_tbo = glGenBuffers(1), glGenTextures(1)
        self.sel_full_bo, self.sel_full_tbo = glGenBuffers(1), glGenTextures(1)
        
        # Visibility Pass Data
        self.vis_bo, self.vis_tbo = glGenBuffers(1), glGenTextures(1)

    def upload_visible_pairs(self, pairs: np.ndarray):
        """Uploads list of (atomIdx, replicaIdx) currently visible."""
        if pairs.size == 0:
            p_data = np.zeros((1, 2), dtype=np.uint32)
            self.visible_count = 0
        else:
            p_data = pairs.astype(np.uint32)
            self.visible_count = len(pairs)
        glBindBuffer(GL_TEXTURE_BUFFER, self.vis_bo)
        glBufferData(GL_TEXTURE_BUFFER, p_data.nbytes, p_data, GL_STREAM_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.vis_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RG32UI, self.vis_bo)

    def draw_shadows(self, view: np.ndarray, proj: np.ndarray, 
                     visible_pairs: np.ndarray, atom_scale: float):
        """Pass 1: Render depth to shadow map."""
        count = visible_pairs.shape[0]
        if count == 0: return
        self.upload_visible_pairs(visible_pairs)

        # Save current viewport
        prev_viewport = glGetIntegerv(GL_VIEWPORT)

        glBindFramebuffer(GL_FRAMEBUFFER, self.shadow_fbo)
        glViewport(0, 0, self.shadow_size, self.shadow_size)
        glClear(GL_DEPTH_BUFFER_BIT)
        glEnable(GL_DEPTH_TEST)
        glDepthFunc(GL_LESS) 
        glUseProgram(self.shadow_prog)
        glEnable(GL_CULL_FACE); glCullFace(GL_FRONT)
        
        glUniformMatrix4fv(glGetUniformLocation(self.shadow_prog, "uView"), 1, GL_FALSE, view.T)
        glUniformMatrix4fv(glGetUniformLocation(self.shadow_prog, "uProj"), 1, GL_FALSE, proj.T)
        glUniform1i(glGetUniformLocation(self.shadow_prog, "uVisibleCount"), count)
        glUniform1f(glGetUniformLocation(self.shadow_prog, "uAtomScale"), atom_scale)
        
        glActiveTexture(GL_TEXTURE0); glBindTexture(GL_TEXTURE_BUFFER, self.pos_tbo)
        glActiveTexture(GL_TEXTURE2); glBindTexture(GL_TEXTURE_BUFFER, self.rep_tbo)
        glActiveTexture(GL_TEXTURE5); glBindTexture(GL_TEXTURE_BUFFER, self.vis_tbo)
        
        glUniform1i(glGetUniformLocation(self.shadow_prog, "uAtomPosRad"), 0)
        glUniform1i(glGetUniformLocation(self.shadow_prog, "uReplicaOffsets"), 2)
        glUniform1i(glGetUniformLocation(self.shadow_prog, "uVisiblePairs"), 5)

        glBindVertexArray(self.quad_vao)
        glDrawArraysInstanced(GL_TRIANGLE_STRIP, 0, 4, count)
        glBindVertexArray(0)
        
        glDisable(GL_CULL_FACE)
        # Restore viewport
        glViewport(prev_viewport[0], prev_viewport[1], prev_viewport[2], prev_viewport[3])

    def upload_materials(self, mat_data1: np.ndarray, mat_data2: np.ndarray):
        """Uploads per-atom material constants across two TBOs."""
        # Buffer 1: [Ambient, Specular, Shininess, Rim]
        glBindBuffer(GL_TEXTURE_BUFFER, self.mat_bo)
        glBufferData(GL_TEXTURE_BUFFER, mat_data1.nbytes, mat_data1, GL_DYNAMIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.mat_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RGBA32F, self.mat_bo)
        
        # Buffer 2: [Sheen, Clearcoat, EnvIntensity, Glossiness]
        glBindBuffer(GL_TEXTURE_BUFFER, self.mat_bo2)
        glBufferData(GL_TEXTURE_BUFFER, mat_data2.nbytes, mat_data2, GL_DYNAMIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.mat_tbo2)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RGBA32F, self.mat_bo2)

    def upload_selection(self, selected_ids: list, total_count: int):
        """Uploads list of atom indices currently selected for glowing."""
        if not selected_ids:
            data = np.array([0], dtype=np.uint32)
            full_data = np.zeros(total_count, dtype=np.uint8)
            self.selected_count = 0
        else:
            data = np.array(selected_ids, dtype=np.uint32)
            full_data = np.zeros(total_count, dtype=np.uint8)
            full_data[selected_ids] = 1
            self.selected_count = len(selected_ids)
            
        glBindBuffer(GL_TEXTURE_BUFFER, self.sel_bo)
        glBufferData(GL_TEXTURE_BUFFER, data.nbytes, data, GL_DYNAMIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.sel_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_R32UI, self.sel_bo)

        glBindBuffer(GL_TEXTURE_BUFFER, self.sel_full_bo)
        glBufferData(GL_TEXTURE_BUFFER, full_data.nbytes, full_data, GL_DYNAMIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.sel_full_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_R8UI, self.sel_full_bo)

    def upload_data(self, positions: np.ndarray, radii: np.ndarray, colors: np.ndarray):
        """Uploads positions/radii and colors to GPU buffers."""
        N = positions.shape[0]
        if N == 0:
            pos_rad = np.zeros((1, 4), dtype=np.float32)
            rgba = np.zeros((1, 4), dtype=np.float32)
            self.atom_count = 0
        else:
            pos_rad = np.hstack([positions, radii[:, None]]).astype(np.float32)
            if colors.shape[1] == 4:
                rgba = colors.astype(np.float32)
            else:
                rgba = np.ones((N, 4), dtype=np.float32)
                rgba[:, :3] = colors
            self.atom_count = N
        glBindBuffer(GL_TEXTURE_BUFFER, self.pos_bo)
        glBufferData(GL_TEXTURE_BUFFER, pos_rad.nbytes, pos_rad, GL_DYNAMIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.pos_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RGBA32F, self.pos_bo)
        
        glBindBuffer(GL_TEXTURE_BUFFER, self.col_bo)
        glBufferData(GL_TEXTURE_BUFFER, rgba.nbytes, rgba, GL_DYNAMIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.col_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RGBA32F, self.col_bo)

    def upload_replicas(self, offsets: np.ndarray):
        offsets_padded = np.hstack([offsets, np.zeros((offsets.shape[0], 1))]).astype(np.float32)
        glBindBuffer(GL_TEXTURE_BUFFER, self.rep_bo)
        glBufferData(GL_TEXTURE_BUFFER, offsets_padded.nbytes, offsets_padded, GL_STATIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.rep_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RGBA32F, self.rep_bo)

    def change_shadow_res(self, res: int):
        """Re-initializes shadow texture with new resolution."""
        if res == self.shadow_size: return
        self.shadow_size = res
        glBindTexture(GL_TEXTURE_2D, self.shadow_tex)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT, res, res, 0, GL_DEPTH_COMPONENT, GL_FLOAT, None)
        glBindTexture(GL_TEXTURE_2D, 0)

    def set_shadow_debug(self, enabled: bool):
        """Toggles hardware comparison. True = Normal shadows, False = Raw Depth Debug."""
        if self.shadow_compare == (not enabled): return
        self.shadow_compare = not enabled
        glBindTexture(GL_TEXTURE_2D, self.shadow_tex)
        if enabled:
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_MODE, GL_NONE)
        else:
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_MODE, GL_COMPARE_REF_TO_TEXTURE)
        glBindTexture(GL_TEXTURE_2D, 0)

    def get_visible_offsets(self, view: np.ndarray, proj: np.ndarray, 
                            all_replica_offsets: np.ndarray, 
                            scene_bounds: Optional[Tuple[np.ndarray, np.ndarray]] = None) -> List[np.ndarray]:
        """Returns list of offsets that are likely visible within the frustum."""
        vp = proj @ view
        visible_offsets = []
        mn, mx = scene_bounds if scene_bounds else (np.zeros(3), np.ones(3)*10)
        cell_size = mx - mn
        center_cell = mn + cell_size * 0.5
        radius_cell = np.linalg.norm(cell_size) * 0.6
        
        for off in all_replica_offsets:
            c = off + center_cell
            cp = vp @ np.append(c, 1.0)
            if cp[3] <= 0:
                if np.linalg.norm(c - view[3,:3]) > radius_cell * 2: continue
            ndc = cp[:3]/cp[3]
            if np.all(np.abs(ndc) < 1.0 + (radius_cell/max(1.0, cp[3]))): 
                visible_offsets.append(off)
        return visible_offsets

    def draw(self, view, proj, visible_pairs, 
             ortho_scale, viewport_height,
             lights=[], materials={},
             atom_scale=1.0, fov=45.0, is_ortho=False,
             skip_px=0.5, point_px=4.0, pretty_render=False,
             shadow_matrix=None, shadow_map=0, shadow_size=512,
             selection_only=False,
             reflection_intensity=0.0, reflection_gloss=0.5,
             fog_density=0.0, time=0.0):
        """Renders spheres with multi-light support and per-atom materials."""
        count = visible_pairs.shape[0]
        if count == 0: return
        self.upload_visible_pairs(visible_pairs)
        
        glUseProgram(self.prog)
        glEnable(GL_DEPTH_TEST)
        glDepthFunc(GL_LEQUAL)
        glEnable(GL_BLEND)
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
        
        # 1. Lights
        glUniform1i(glGetUniformLocation(self.prog, "uNumLights"), len(lights))
        vr = view[:3, :3]
        for i, l in enumerate(lights[:4]):
             lv = vr @ np.array(l['pos'], dtype=np.float32)
             glUniform3f(glGetUniformLocation(self.prog, f"uLightDir[{i}]"), *lv)
             glUniform3f(glGetUniformLocation(self.prog, f"uLightCol[{i}]"), *l['col'])
            
        # 2. Materials & VFX
        glUniform1f(glGetUniformLocation(self.prog, "uExposure"), materials.get('exposure', 1.5))
        glUniform1f(glGetUniformLocation(self.prog, "uHeadlight"), materials.get('headlight', 0.8))
        glUniform1i(glGetUniformLocation(self.prog, "uShowShadows"), 1 if materials.get('show_shadows', True) else 0)
        glUniform1f(glGetUniformLocation(self.prog, "uShadowIntensity"), materials.get('shadow_intensity', 0.6))
        glUniform1f(glGetUniformLocation(self.prog, "uShadowBias"), materials.get('shadow_bias', 0.002))
        
        softness = materials.get('shadow_softness', 1.5) / float(self.shadow_size)
        glUniform1f(glGetUniformLocation(self.prog, "uShadowSoftness"), softness)

        glUniform1f(glGetUniformLocation(self.prog, "uReflectionIntensity"), reflection_intensity)
        glUniform1f(glGetUniformLocation(self.prog, "uReflectionGloss"), reflection_gloss)
        glUniform1f(glGetUniformLocation(self.prog, "uSelectionGlow"), materials.get('selection_glow', 1.0))
        glUniform1f(glGetUniformLocation(self.prog, "uSheen"), materials.get('sheen', 0.0))
        glUniform1f(glGetUniformLocation(self.prog, "uClearcoat"), materials.get('clearcoat', 0.0))
        glUniform1f(glGetUniformLocation(self.prog, "uFogDensity"), fog_density)
        
        # Atom Motion
        glUniform1f(glGetUniformLocation(self.prog, "uPulseSpeed"), materials.get('pulse_speed', 4.0))
        glUniform1f(glGetUniformLocation(self.prog, "uPulseAmp"), materials.get('pulse_amp', 0.0))
        glUniform1f(glGetUniformLocation(self.prog, "uTime"), time)

        # 3. Matrices
        glUniformMatrix4fv(glGetUniformLocation(self.prog, "uView"), 1, GL_FALSE, view.T)
        glUniformMatrix4fv(glGetUniformLocation(self.prog, "uProj"), 1, GL_FALSE, proj.T)
        if shadow_matrix is not None:
             glUniformMatrix4fv(glGetUniformLocation(self.prog, "uShadowMatrix"), 1, GL_FALSE, shadow_matrix.T)
        view_inv = np.linalg.inv(view)
        glUniformMatrix4fv(glGetUniformLocation(self.prog, "uViewInv"), 1, GL_FALSE, view_inv.T)
        
        import math
        ps = viewport_height / (ortho_scale if is_ortho else (2.0 * math.tan(math.radians(fov)/2.0)))
        glUniform1f(glGetUniformLocation(self.prog, "uProjScale"), ps)

        # 4. State
        glUniform1i(glGetUniformLocation(self.prog, "uVisibleCount"), count)
        glUniform1i(glGetUniformLocation(self.prog, "uIsOrtho"), 1 if is_ortho else 0)
        glUniform1f(glGetUniformLocation(self.prog, "uAtomScale"), atom_scale)
        glUniform1f(glGetUniformLocation(self.prog, "uViewportHeight"), float(viewport_height))
        glUniform1f(glGetUniformLocation(self.prog, "uSkipPx"), skip_px)
        glUniform1f(glGetUniformLocation(self.prog, "uPointPx"), point_px)
        glUniform1i(glGetUniformLocation(self.prog, "uPrettyRender"), 1 if pretty_render else 0)
        glUniform1i(glGetUniformLocation(self.prog, "uSelectedCount"), getattr(self, 'selected_count', 0))

        # 5. Textures (Units 0-7)
        glActiveTexture(GL_TEXTURE0); glBindTexture(GL_TEXTURE_BUFFER, self.pos_tbo)
        glActiveTexture(GL_TEXTURE1); glBindTexture(GL_TEXTURE_BUFFER, self.col_tbo)
        glActiveTexture(GL_TEXTURE2); glBindTexture(GL_TEXTURE_BUFFER, self.rep_tbo)
        glActiveTexture(GL_TEXTURE3); glBindTexture(GL_TEXTURE_BUFFER, self.mat_tbo)
        glActiveTexture(GL_TEXTURE4); glBindTexture(GL_TEXTURE_2D, shadow_map if shadow_map != 0 else self.shadow_tex)
        glActiveTexture(GL_TEXTURE5); glBindTexture(GL_TEXTURE_BUFFER, self.vis_tbo)
        glActiveTexture(GL_TEXTURE6); glBindTexture(GL_TEXTURE_BUFFER, self.sel_tbo)
        glActiveTexture(GL_TEXTURE7); glBindTexture(GL_TEXTURE_BUFFER, self.mat_tbo2) # New material TBO
        glActiveTexture(GL_TEXTURE8); glBindTexture(GL_TEXTURE_BUFFER, self.sel_full_tbo)

        glUniform1i(glGetUniformLocation(self.prog, "uAtomPosRad"), 0)
        glUniform1i(glGetUniformLocation(self.prog, "uAtomColor"), 1)
        glUniform1i(glGetUniformLocation(self.prog, "uReplicaOffsets"), 2)
        glUniform1i(glGetUniformLocation(self.prog, "uAtomMaterial"), 3)
        glUniform1i(glGetUniformLocation(self.prog, "uShadowMap"), 4)
        glUniform1i(glGetUniformLocation(self.prog, "uVisiblePairs"), 5)
        glUniform1i(glGetUniformLocation(self.prog, "uSelectedAtoms"), 6)
        glUniform1i(glGetUniformLocation(self.prog, "uAtomMaterial2"), 7) # New uniform for second material TBO
        glUniform1i(glGetUniformLocation(self.prog, "uSelectionMaskFull"), 8)
        
        glUniform1i(glGetUniformLocation(self.prog, "uSelectionOnly"), 1 if selection_only else 0)

        # 6. Draw
        glBindVertexArray(self.quad_vao)
        glDrawArraysInstanced(GL_TRIANGLE_STRIP, 0, 4, count)
        glBindVertexArray(0)
