# ============================
# picking_renderer.py — Off-screen Atom Picking
# ============================
from __future__ import annotations
from OpenGL.GL import *
import numpy as np
from .shaders import PICK_VS, PICK_FS
from .utils import link_program

class PickingRenderer:
    """
    Renders unique color IDs to an off-screen buffer to detect which atom 
    is under the mouse cursor. Supports high-resolution displays.
    """
    def __init__(self):
        self.prog = link_program(PICK_VS, PICK_FS)
        
        # 1. Geometry Setup (Fullscreen Quad)
        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)

        # 2. Private FBO Setup
        self.fbo = glGenFramebuffers(1)
        self.color_tex = None
        self.depth_rb = None
        self._res = (0, 0)

    def _ensure_fbo_res(self, w: int, h: int):
        """Resizes the internal FBO if the window dimensions change."""
        if (w, h) == self._res: return
        self._res = (w, h)
        
        if self.color_tex is not None: glDeleteTextures([self.color_tex])
        if self.depth_rb is not None: glDeleteRenderbuffers(1, [self.depth_rb])
        
        glBindFramebuffer(GL_FRAMEBUFFER, self.fbo)
        
        # Color: 24-bit RGB (Matches glReadPixels request)
        self.color_tex = glGenTextures(1)
        glBindTexture(GL_TEXTURE_2D, self.color_tex)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, w, h, 0, GL_RGB, GL_UNSIGNED_BYTE, None)
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, self.color_tex, 0)
        
        # Depth: for standard occlusion
        self.depth_rb = glGenRenderbuffers(1)
        glBindRenderbuffer(GL_RENDERBUFFER, self.depth_rb)
        glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH_COMPONENT24, w, h)
        glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_RENDERBUFFER, self.depth_rb)
        
        status = glCheckFramebufferStatus(GL_FRAMEBUFFER)
        if status != GL_FRAMEBUFFER_COMPLETE:
            print(f"Picking FBO failure: {status}")
        glBindFramebuffer(GL_FRAMEBUFFER, 0)

    def pick(self, x: float, y: float, w: int, h: int, count: int, replica_count: int, 
             v_tbo: int, r_tbo: int, view: np.ndarray, proj: np.ndarray, is_ortho: bool):
        """Returns the index of the atom at screen coordinates (x, y)."""
        if count == 0: return None
        
        # 1. Isolated State Setup
        prev_fbo = glGetIntegerv(GL_FRAMEBUFFER_BINDING)
        self._ensure_fbo_res(w, h)
        glBindFramebuffer(GL_FRAMEBUFFER, self.fbo)
        
        glViewport(0, 0, w, h)
        glClearColor(0, 0, 0, 1)
        glEnable(GL_DEPTH_TEST)
        glDepthMask(GL_TRUE)
        glDepthFunc(GL_LESS)
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        
        # 2. Render IDs
        glUseProgram(self.prog)
        glUniformMatrix4fv(glGetUniformLocation(self.prog, "uView"), 1, GL_FALSE, view.T)
        glUniformMatrix4fv(glGetUniformLocation(self.prog, "uProj"), 1, GL_FALSE, proj.T)
        glUniform1i(glGetUniformLocation(self.prog, "uIsOrtho"), 1 if is_ortho else 0)
        glUniform1i(glGetUniformLocation(self.prog, "uAtomCount"), count)
        
        glActiveTexture(GL_TEXTURE0); glBindTexture(GL_TEXTURE_BUFFER, v_tbo)
        glActiveTexture(GL_TEXTURE1); glBindTexture(GL_TEXTURE_BUFFER, r_tbo)
        glUniform1i(glGetUniformLocation(self.prog, "uAtomPosRad"), 0)
        glUniform1i(glGetUniformLocation(self.prog, "uReplicaOffsets"), 1)
        
        glBindVertexArray(self.vao)
        glDrawArraysInstanced(GL_TRIANGLE_STRIP, 0, 4, count * replica_count)
        
        # 3. Read specific pixel from private FBO
        px = glReadPixels(int(x), int(h - y), 1, 1, GL_RGB, GL_UNSIGNED_BYTE)
        id_px = np.frombuffer(px, dtype=np.uint8)
        
        # 4. Cleanup and Restore State
        glBindFramebuffer(GL_FRAMEBUFFER, prev_fbo)
        
        if id_px.size == 3:
            # Reconstruct 24-bit ID from RGB
            idx = int(id_px[0]) + (int(id_px[1]) << 8) + (int(id_px[2]) << 16)
            return idx if idx < count else None
            
        return None
