"""
volume_renderer.py — GPU Raymarching for Volumetric Density Grids.
Supports Isosurface, MIP, and Translucency rendering.
"""
from __future__ import annotations
from OpenGL.GL import *
import numpy as np
from typing import Optional, Tuple
from .shaders import VOLUME_VS, VOLUME_FS
from .utils import link_program

class VolumeRenderer:
    """
    Renders 3D density grids using single-pass raymarching.
    Handles grid normalization, 3D texture mapping, and isosurface extraction.
    """
    def __init__(self):
        self.prog = link_program(VOLUME_VS, VOLUME_FS)
        
        # Cube VAO (for raymarching entry/exit)
        self.vao = glGenVertexArrays(1)
        glBindVertexArray(self.vao)
        # 1x1x1 unit cube vertices (0 to 1 range)
        cube = np.array([
            0,0,0, 1,0,0, 1,1,0, 0,1,0,
            0,0,1, 1,0,1, 1,1,1, 0,1,1
        ], dtype=np.float32)
        indices = np.array([
            0,1,2, 0,2,3, 4,5,6, 4,6,7, # Front/Back
            0,4,7, 0,7,3, 1,5,6, 1,6,2, # Left/Right
            0,1,5, 0,5,4, 3,2,6, 3,6,7  # Bottom/Top
        ], dtype=np.uint32)
        
        vbo = glGenBuffers(1); glBindBuffer(GL_ARRAY_BUFFER, vbo)
        glBufferData(GL_ARRAY_BUFFER, cube.nbytes, cube, GL_STATIC_DRAW)
        glEnableVertexAttribArray(0)
        glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, None)
        
        ebo = glGenBuffers(1); glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ebo)
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, indices.nbytes, indices, GL_STATIC_DRAW)
        glBindVertexArray(0)
        
        self.tex = glGenTextures(1)
        self.has_data = False

    def upload_data(self, grid: np.ndarray, origin: np.ndarray, spacing: np.ndarray):
        """
        grid: (nx, ny, nz) float32
        """
        self.grid_shape = grid.shape
        self.origin = origin
        self.spacing = spacing
        self.size = np.array(grid.shape, dtype=np.float32) * spacing
        
        # Enforce [0..1] range for the shader if not already done
        grid_norm = grid.astype(np.float32)
        g_min, g_max = grid_norm.min(), grid_norm.max()
        if g_max > g_min:
            grid_norm = (grid_norm - g_min) / (g_max - g_min)
            
        glBindTexture(GL_TEXTURE_3D, self.tex)
        glTexImage3D(GL_TEXTURE_3D, 0, GL_R32F, grid.shape[2], grid.shape[1], grid.shape[0], 0, GL_RED, GL_FLOAT, grid_norm.flatten())
        glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
        glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
        glTexParameteri(GL_TEXTURE_3D, GL_TEXTURE_WRAP_R, GL_CLAMP_TO_EDGE)
        self.has_data = True

    def draw(self, view: np.ndarray, proj: np.ndarray, cam_pos: np.ndarray, 
             iso_level: float = 0.5, mode: int = 0, col: Tuple[float, float, float] = (0.4, 0.7, 1.0)):
        
        if not self.has_data: return

        glEnable(GL_BLEND)
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
        glDisable(GL_CULL_FACE)
        
        glUseProgram(self.prog)
        
        glUniformMatrix4fv(glGetUniformLocation(self.prog, "uView"), 1, GL_FALSE, view.T)
        glUniformMatrix4fv(glGetUniformLocation(self.prog, "uProj"), 1, GL_FALSE, proj.T)
        glUniform3f(glGetUniformLocation(self.prog, "uCamPos"), *cam_pos)
        glUniform3f(glGetUniformLocation(self.prog, "uOrigin"), *self.origin)
        glUniform3f(glGetUniformLocation(self.prog, "uSize"), *self.size)
        
        glUniform1f(glGetUniformLocation(self.prog, "uIsoLevel"), float(iso_level))
        glUniform1i(glGetUniformLocation(self.prog, "uMode"), int(mode))
        glUniform3f(glGetUniformLocation(self.prog, "uBaseColor"), *col)
        glUniform1f(glGetUniformLocation(self.prog, "uOpacity"), 0.5) # Default opacity
        glUniform1f(glGetUniformLocation(self.prog, "uStepSize"), 0.01) # Default step size
        
        glActiveTexture(GL_TEXTURE0)
        glBindTexture(GL_TEXTURE_3D, self.tex)
        glUniform1i(glGetUniformLocation(self.prog, "uVolume"), 0)

        glBindVertexArray(self.vao)
        glDrawElements(GL_TRIANGLES, 36, GL_UNSIGNED_INT, None)
        glBindVertexArray(0)
        
        glDisable(GL_BLEND)
        glEnable(GL_CULL_FACE)
