"""
bond_renderer.py — GPU Impostor Cylinders for Molecular Bonds.
Uses instanced raycasting for high-performance cylinder rendering.
"""
from __future__ import annotations
from OpenGL.GL import *
import numpy as np
from typing import Optional, Tuple, List, Dict
from .shaders import CYLINDER_VS, CYLINDER_FS, SHADOW_CYLINDER_VS, SHADOW_CYLINDER_FS
from .utils import link_program

class BondRenderer:
    """
    Renders bonds as cylinders using GPU raycasting into impostor billboards.
    Supports periodic boundaries and instanced rendering.
    """
    def __init__(self):
        self.prog = link_program(CYLINDER_VS, CYLINDER_FS)
        self.prog_shadow = link_program(SHADOW_CYLINDER_VS, SHADOW_CYLINDER_FS)
        
        # Screen Quad VAO
        self.vao = glGenVertexArrays(1)
        glBindVertexArray(self.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.bond_bo, self.bond_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)

    def upload_data(self, pos: np.ndarray, bonds: np.ndarray, 
                    colors: np.ndarray, lattice: Optional[np.ndarray] = None):
        """
        pos: (N, 3) Å
        bonds: (Nb, 5) indices (i, j) + fractional shift (dx, dy, dz)
        colors: (N, 4) half-bond colors
        lattice: (3, 3) Å for shift conversion
        """
        self.num_bonds = len(bonds)
        if self.num_bonds == 0: return

        # Atoms positions
        glBindBuffer(GL_TEXTURE_BUFFER, self.pos_bo)
        glBufferData(GL_TEXTURE_BUFFER, pos.nbytes, pos, GL_DYNAMIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.pos_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RGB32F, self.pos_bo)

        # Bond definitions (i, j, shift_x, shift_y, shift_z)
        # We store them as float32 for the TBO access even if indices are ints
        bond_data = bonds.astype(np.float32)
        glBindBuffer(GL_TEXTURE_BUFFER, self.bond_bo)
        glBufferData(GL_TEXTURE_BUFFER, bond_data.nbytes, bond_data, GL_DYNAMIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.bond_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RGB32F, self.bond_bo) # Use RGB32F for (i,j,dx)? No, maybe 1x5?

        # Colors
        glBindBuffer(GL_TEXTURE_BUFFER, self.col_bo)
        glBufferData(GL_TEXTURE_BUFFER, colors.nbytes, colors, GL_DYNAMIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.col_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RGBA32F, self.col_bo)

    def upload_replicas(self, offsets: List[np.ndarray]):
        offsets_arr = np.array(offsets, dtype=np.float32)
        self.replica_count = len(offsets)
        glBindBuffer(GL_TEXTURE_BUFFER, self.rep_bo)
        glBufferData(GL_TEXTURE_BUFFER, offsets_arr.nbytes, offsets_arr, GL_STATIC_DRAW)
        glBindTexture(GL_TEXTURE_BUFFER, self.rep_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RGB32F, self.rep_bo)

    def draw(self, view: np.ndarray, proj: np.ndarray, vis_pairs: np.ndarray, 
             lights: list, materials: dict, bond_scale: float = 0.2,
             pretty_render: bool = False, shadow_matrix: Optional[np.ndarray] = None,
             shadow_map: Optional[int] = None, shadow_size: int = 4096,
             reflection_intensity: float = 1.0, reflection_gloss: float = 1.0,
             fog_density: float = 0.002, time: float = 0.0,
             mat_tbo: Optional[int] = None, mat_tbo2: Optional[int] = None):
        
        if self.num_bonds == 0 or vis_pairs.shape[0] == 0: return

        glUseProgram(self.prog)
        
        # Matrix Uniforms
        glUniformMatrix4fv(glGetUniformLocation(self.prog, "uView"), 1, GL_FALSE, view.T)
        glUniformMatrix4fv(glGetUniformLocation(self.prog, "uProj"), 1, GL_FALSE, proj.T)
        glUniformMatrix4fv(glGetUniformLocation(self.prog, "uViewInv"), 1, GL_FALSE, np.linalg.inv(view).T)
        
        glUniform1f(glGetUniformLocation(self.prog, "uBondScale"), bond_scale)
        glUniform1f(glGetUniformLocation(self.prog, "uTime"), time)
        glUniform1i(glGetUniformLocation(self.prog, "uPrettyRender"), int(pretty_render))
        glUniform1f(glGetUniformLocation(self.prog, "uFogDensity"), fog_density)
        glUniform1f(glGetUniformLocation(self.prog, "uReflectionIntensity"), reflection_intensity)
        
        if shadow_matrix is not None:
            glUniformMatrix4fv(glGetUniformLocation(self.prog, "uShadowMatrix"), 1, GL_FALSE, shadow_matrix.T)
            glUniform1i(glGetUniformLocation(self.prog, "uShowShadows"), 1)
            glUniform1f(glGetUniformLocation(self.prog, "uShadowIntensity"), materials.get('shadow_intensity', 0.6))
            glUniform1f(glGetUniformLocation(self.prog, "uShadowSoftness"), materials.get('shadow_softness', 2.0))
            glUniform1f(glGetUniformLocation(self.prog, "uShadowBias"), materials.get('shadow_bias', 0.001))
            
            glActiveTexture(GL_TEXTURE6) # Use a higher slot for shadow map
            glBindTexture(GL_TEXTURE_2D, shadow_map)
            glUniform1i(glGetUniformLocation(self.prog, "uShadowMap"), 6)
        else:
            glUniform1i(glGetUniformLocation(self.prog, "uShowShadows"), 0)

        # Lights (Directional from light pos)
        if lights:
            l_dir = np.array(lights[0]['pos'], dtype=np.float32)
            l_dir_vs = (view[:3, :3] @ l_dir).tolist()
            glUniform3f(glGetUniformLocation(self.prog, "uLightDirVS"), *l_dir_vs)
        
        # TBOs
        for i, (loc, tbo) in enumerate([
            ("uAtomPos", self.pos_tbo),
            ("uBondData", self.bond_tbo),
            ("uAtomCol", self.col_tbo),
            ("uReplicaOffsets", self.rep_tbo),
            ("uAtomMaterial", mat_tbo),
            ("uAtomMaterial2", mat_tbo2)
        ]):
            if tbo is not None:
                glActiveTexture(GL_TEXTURE0 + i)
                glBindTexture(GL_TEXTURE_BUFFER, tbo)
                glUniform1i(glGetUniformLocation(self.prog, loc), i)

        # Bond visibility pairs (instanced input)
        vis_bo = glGenBuffers(1)
        glBindBuffer(GL_TEXTURE_BUFFER, vis_bo)
        glBufferData(GL_TEXTURE_BUFFER, vis_pairs.nbytes, vis_pairs, GL_STREAM_DRAW)
        vis_tbo = glGenTextures(1)
        glBindTexture(GL_TEXTURE_BUFFER, vis_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RG32UI, vis_bo) 
        
        glActiveTexture(GL_TEXTURE7)
        glBindTexture(GL_TEXTURE_BUFFER, vis_tbo)
        glUniform1i(glGetUniformLocation(self.prog, "uVisiblePairs"), 7)

        glBindVertexArray(self.vao)
        glDrawArraysInstanced(GL_TRIANGLE_STRIP, 0, 4, vis_pairs.shape[0])
        
        glDeleteTextures(1, [vis_tbo]); glDeleteBuffers(1, [vis_bo])

    def draw_shadows(self, view: np.ndarray, proj: np.ndarray, vis_pairs: np.ndarray, 
                     scale: float, shadow_fbo: int, shadow_size: int):
        if self.num_bonds == 0 or vis_pairs.shape[0] == 0: return

        glBindFramebuffer(GL_FRAMEBUFFER, shadow_fbo)
        glViewport(0, 0, shadow_size, shadow_size)
        glUseProgram(self.prog_shadow)
        
        glUniformMatrix4fv(glGetUniformLocation(self.prog_shadow, "uView"), 1, GL_FALSE, view.T)
        glUniformMatrix4fv(glGetUniformLocation(self.prog_shadow, "uProj"), 1, GL_FALSE, proj.T)
        glUniform1f(glGetUniformLocation(self.prog_shadow, "uBondScale"), scale)
        
        # TBOs
        glActiveTexture(GL_TEXTURE0); glBindTexture(GL_TEXTURE_BUFFER, self.pos_tbo)
        glUniform1i(glGetUniformLocation(self.prog_shadow, "uAtomPos"), 0)
        glActiveTexture(GL_TEXTURE1); glBindTexture(GL_TEXTURE_BUFFER, self.bond_tbo)
        glUniform1i(glGetUniformLocation(self.prog_shadow, "uBondData"), 1)
        glActiveTexture(GL_TEXTURE2); glBindTexture(GL_TEXTURE_BUFFER, self.rep_tbo)
        glUniform1i(glGetUniformLocation(self.prog_shadow, "uReplicaOffsets"), 2)

        vis_bo = glGenBuffers(1); glBindBuffer(GL_TEXTURE_BUFFER, vis_bo)
        glBufferData(GL_TEXTURE_BUFFER, vis_pairs.nbytes, vis_pairs, GL_STREAM_DRAW)
        vis_tbo = glGenTextures(1); glBindTexture(GL_TEXTURE_BUFFER, vis_tbo)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RG32UI, vis_bo)
        glActiveTexture(GL_TEXTURE3); glBindTexture(GL_TEXTURE_BUFFER, vis_tbo)
        glUniform1i(glGetUniformLocation(self.prog_shadow, "uVisiblePairs"), 3)

        glBindVertexArray(self.vao)
        glDrawArraysInstanced(GL_TRIANGLE_STRIP, 0, 4, vis_pairs.shape[0])
        
        glDeleteTextures(1, [vis_tbo]); glDeleteBuffers(1, [vis_bo])
