import numpy as np
from OpenGL.GL import *
from .utils import link_program
from .shaders import FULLSCREEN_VS, SSAO_FS, BLOOM_EXTRACT_FS, GAUSSIAN_BLUR_FS, GAUSSIAN_BILATERAL_FS, RESOLVE_FS

class PostProcessor:
    def __init__(self):
        self.width = 1
        self.height = 1
        
        # Shaders
        self.prog_ssao = link_program(FULLSCREEN_VS, SSAO_FS)
        self.prog_bloom_extract = link_program(FULLSCREEN_VS, BLOOM_EXTRACT_FS)
        self.prog_blur = link_program(FULLSCREEN_VS, GAUSSIAN_BLUR_FS)
        self.prog_bilateral_blur = link_program(FULLSCREEN_VS, GAUSSIAN_BILATERAL_FS)
        self.prog_resolve = link_program(FULLSCREEN_VS, RESOLVE_FS)
        
        # Screen Quad VAO
        self.vao = glGenVertexArrays(1)
        glBindVertexArray(self.vao)
        vbo = glGenBuffers(1)
        glBindBuffer(GL_ARRAY_BUFFER, vbo)
        data = np.array([-1,-1, 1,-1, -1,1, 1,1], dtype=np.float32)
        glBufferData(GL_ARRAY_BUFFER, data.nbytes, data, GL_STATIC_DRAW)
        glEnableVertexAttribArray(0)
        glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 0, None)
        glBindVertexArray(0)
        
        # Resources
        self.fbo_main = glGenFramebuffers(1)
        self.tex_main = glGenTextures(1)
        self.tex_depth = glGenTextures(1)
        self.tex_selection = glGenTextures(1)
        self.fbo_selection = glGenFramebuffers(1)
        
        self.fbo_ssao = glGenFramebuffers(2)
        self.tex_ssao = glGenTextures(2)
        
        self.fbo_bloom = glGenFramebuffers(2)
        self.tex_bloom = glGenTextures(2)

    def resize(self, w, h):
        if w == self.width and h == self.height: return
        self.width, self.height = max(1, w), max(1, h)
        
        # Main Scene Texture
        glBindTexture(GL_TEXTURE_2D, self.tex_main)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB16F, w, h, 0, GL_RGB, GL_FLOAT, None)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        
        # Depth Texture (for SSAO)
        glBindTexture(GL_TEXTURE_2D, self.tex_depth)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT32F, w, h, 0, GL_DEPTH_COMPONENT, GL_FLOAT, None)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        
        glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_main)
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, self.tex_main, 0)
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_TEXTURE_2D, self.tex_depth, 0)
        
        # Selection Mask Texture
        glBindTexture(GL_TEXTURE_2D, self.tex_selection)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_R8, w, h, 0, GL_RED, GL_UNSIGNED_BYTE, None)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        
        glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_selection)
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, self.tex_selection, 0)
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_TEXTURE_2D, self.tex_depth, 0) # Share depth!
        
        # SSAO Textures (Ping-pong for blur) - Use GL_R16F for precision
        for i in range(2):
            glBindTexture(GL_TEXTURE_2D, self.tex_ssao[i])
            glTexImage2D(GL_TEXTURE_2D, 0, GL_R16F, w, h, 0, GL_RED, GL_FLOAT, None)
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
            
            glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_ssao[i])
            glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, self.tex_ssao[i], 0)
        
        # Bloom Textures (Ping-pong)
        for i in range(2):
            glBindTexture(GL_TEXTURE_2D, self.tex_bloom[i])
            glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB16F, w//2, h//2, 0, GL_RGB, GL_FLOAT, None) # Half res
            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_EDGE)
            glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
            
            glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_bloom[i])
            glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, self.tex_bloom[i], 0)
            
        # Validation
        glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_main)
        status = glCheckFramebufferStatus(GL_FRAMEBUFFER)
        if status != GL_FRAMEBUFFER_COMPLETE:
             print(f"FBO MAIN: {hex(status)} (INCOMPLETE!)")
             raise RuntimeError(f"PostProcess: fbo_main incomplete!")
             
        for i in range(2):
            glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_ssao[i])
            if glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE:
                raise RuntimeError(f"PostProcess: fbo_ssao[{i}] incomplete!")
            
        for i in range(2):
            glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_bloom[i])
            if glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE:
                raise RuntimeError(f"PostProcess: fbo_bloom[{i}] incomplete!")

        glBindFramebuffer(GL_FRAMEBUFFER, 0)

    def begin(self):
        glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_main)
        glEnable(GL_DEPTH_TEST)
        glDepthMask(GL_TRUE)
        glClearColor(0.01, 0.01, 0.02, 1.0) # Match fog color
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

    def end(self, materials, show_ssao=True, show_bloom=True, show_fxaa=True, 
            near=0.1, far=100.0, viewport=None, is_ortho=False):
        """Final resolve pass with optional viewport for sub-region rendering."""
        prev_viewport = glGetIntegerv(GL_VIEWPORT)
        
        glBindFramebuffer(GL_FRAMEBUFFER, 0)
        glDisable(GL_DEPTH_TEST)
        glDisable(GL_BLEND)
        glDepthMask(GL_FALSE)
        
        if viewport:
            glViewport(*viewport)
        else:
            glViewport(0, 0, self.width, self.height)
        
        if show_ssao:
            if 'proj' not in materials:
                raise RuntimeError("PostProcess: SSAO requires 'proj' matrix in materials!")
                
            # Pass 1: AO Calculation
            glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_ssao[0])
            glViewport(0, 0, self.width, self.height)
            glUseProgram(self.prog_ssao)
            glActiveTexture(GL_TEXTURE0); glBindTexture(GL_TEXTURE_2D, self.tex_depth)
            glUniform1i(glGetUniformLocation(self.prog_ssao, "uDepthTex"), 0)
            
            glUniformMatrix4fv(glGetUniformLocation(self.prog_ssao, "uProj"), 1, GL_FALSE, materials['proj'].T)
            glUniformMatrix4fv(glGetUniformLocation(self.prog_ssao, "uProjInv"), 1, GL_FALSE, materials['projInv'].T)
            
            glUniform1f(glGetUniformLocation(self.prog_ssao, "uNear"), near)
            glUniform1f(glGetUniformLocation(self.prog_ssao, "uFar"), far)
            glUniform1f(glGetUniformLocation(self.prog_ssao, "uRadius"), materials.get('ssao_radius', 0.05))
            glUniform1f(glGetUniformLocation(self.prog_ssao, "uStrength"), materials.get('ssao_strength', 1.0))
            glUniform1f(glGetUniformLocation(self.prog_ssao, "uBias"), materials.get('ssao_bias', 0.01))
            glUniform1i(glGetUniformLocation(self.prog_ssao, "uIsOrtho"), 1 if is_ortho else 0)
            
            glBindVertexArray(self.vao)
            glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)
            
            # Pass 2: Bilateral Blur
            glUseProgram(self.prog_bilateral_blur)
            glUniform1f(glGetUniformLocation(self.prog_bilateral_blur, "uNear"), near)
            glUniform1f(glGetUniformLocation(self.prog_bilateral_blur, "uFar"), far)
            glActiveTexture(GL_TEXTURE1); glBindTexture(GL_TEXTURE_2D, self.tex_depth)
            glUniform1i(glGetUniformLocation(self.prog_bilateral_blur, "uDepthTex"), 1)
            
            # Horizontal pass (0 -> 1)
            glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_ssao[1])
            glUniform1i(glGetUniformLocation(self.prog_bilateral_blur, "uHorizontal"), 1)
            glActiveTexture(GL_TEXTURE0); glBindTexture(GL_TEXTURE_2D, self.tex_ssao[0])
            glUniform1i(glGetUniformLocation(self.prog_bilateral_blur, "uImage"), 0)
            glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)
            
            # Vertical pass (1 -> 0)
            glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_ssao[0])
            glUniform1i(glGetUniformLocation(self.prog_bilateral_blur, "uHorizontal"), 0)
            glActiveTexture(GL_TEXTURE0); glBindTexture(GL_TEXTURE_2D, self.tex_ssao[1])
            glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)
        
        # 2. Bloom Pass
        if show_bloom:
            # Extract
            glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_bloom[0])
            glViewport(0, 0, self.width//2, self.height//2)
            glUseProgram(self.prog_bloom_extract)
            glActiveTexture(GL_TEXTURE0); glBindTexture(GL_TEXTURE_2D, self.tex_main)
            glUniform1i(glGetUniformLocation(self.prog_bloom_extract, "uColorTex"), 0)
            glUniform1f(glGetUniformLocation(self.prog_bloom_extract, "uThreshold"), materials.get('bloom_threshold', 0.8))
            glBindVertexArray(self.vao)
            glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)
            
            # Blur (Ping-pong)
            glUseProgram(self.prog_blur)
            horizontal = True
            for _ in range(4): # 2 iterations (horizontal + vertical)
                glBindFramebuffer(GL_FRAMEBUFFER, self.fbo_bloom[1 if horizontal else 0])
                glUniform1i(glGetUniformLocation(self.prog_blur, "uHorizontal"), 1 if horizontal else 0)
                glActiveTexture(GL_TEXTURE0); glBindTexture(GL_TEXTURE_2D, self.tex_bloom[0 if horizontal else 1])
                glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)
                horizontal = not horizontal
            
        # 3. Final Resolve
        glBindFramebuffer(GL_FRAMEBUFFER, 0)
        if viewport:
            glViewport(*viewport)
        else:
            glViewport(0, 0, self.width, self.height)
            
        glUseProgram(self.prog_resolve)
        
        glActiveTexture(GL_TEXTURE0); glBindTexture(GL_TEXTURE_2D, self.tex_main)
        glActiveTexture(GL_TEXTURE1); glBindTexture(GL_TEXTURE_2D, self.tex_ssao[0])
        glActiveTexture(GL_TEXTURE2); glBindTexture(GL_TEXTURE_2D, self.tex_bloom[0])
        glActiveTexture(GL_TEXTURE3); glBindTexture(GL_TEXTURE_2D, self.tex_depth)
        glActiveTexture(GL_TEXTURE4); glBindTexture(GL_TEXTURE_2D, self.tex_selection)
        
        glUniform1i(glGetUniformLocation(self.prog_resolve, "uSceneTex"), 0)
        glUniform1i(glGetUniformLocation(self.prog_resolve, "uSSAO"), 1)
        glUniform1i(glGetUniformLocation(self.prog_resolve, "uBloom"), 2)
        glUniform1i(glGetUniformLocation(self.prog_resolve, "uDepthTex"), 3)
        glUniform1i(glGetUniformLocation(self.prog_resolve, "uSelectionMask"), 4)
        
        glUniform1i(glGetUniformLocation(self.prog_resolve, "uEnableSSAO"), 1 if show_ssao else 0)
        glUniform1i(glGetUniformLocation(self.prog_resolve, "uEnableBloom"), 1 if show_bloom else 0)
        glUniform1i(glGetUniformLocation(self.prog_resolve, "uEnableFXAA"), 1 if show_fxaa else 0)
        glUniform1i(glGetUniformLocation(self.prog_resolve, "uDebugMode"), materials.get('debug_mode', 0))
        glUniform1f(glGetUniformLocation(self.prog_resolve, "uBloomIntensity"), materials.get('bloom_intensity', 0.5))
        glUniform1f(glGetUniformLocation(self.prog_resolve, "uGamma"), materials.get('gamma', 1.0))
        glUniform1f(glGetUniformLocation(self.prog_resolve, "uNear"), near)
        glUniform1f(glGetUniformLocation(self.prog_resolve, "uFar"), far)
        
        glBindVertexArray(self.vao)
        glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)
        
        # Restore viewport and depth test for subsequent rendering (e.g. split view / next window)
        glBindVertexArray(0)
        glViewport(*prev_viewport)
        glEnable(GL_DEPTH_TEST)

