import OpenGL.GL as gl

class ShaderManager:
 def __init__(self):
  self.program = None
  self.bg_program = None
  self.pp_program = None

 def compile_shader(self, vertex_src: str, fragment_src: str) -> int:
  program = gl.glCreateProgram()
  
  def _compile(src, shader_type):
   shader = gl.glCreateShader(shader_type)
   gl.glShaderSource(shader, src)
   gl.glCompileShader(shader)
   
   if not gl.glGetShaderiv(shader, gl.GL_COMPILE_STATUS):
    info_log = gl.glGetShaderInfoLog(shader)
    if hasattr(info_log, "decode"):
     info_log = info_log.decode()
    raise RuntimeError(f"Shader compilation failed: {info_log}")
   return shader

  vert_shader = _compile(vertex_src, gl.GL_VERTEX_SHADER)
  frag_shader = _compile(fragment_src, gl.GL_FRAGMENT_SHADER)
  
  gl.glAttachShader(program, vert_shader)
  gl.glAttachShader(program, frag_shader)
  gl.glLinkProgram(program)
  
  if not gl.glGetProgramiv(program, gl.GL_LINK_STATUS):
   info_log = gl.glGetProgramInfoLog(program)
   if hasattr(info_log, "decode"):
    info_log = info_log.decode()
   raise RuntimeError(f"Program linking failed: {info_log}")
  
  gl.glDeleteShader(vert_shader)
  gl.glDeleteShader(frag_shader)
  
  self.program = program
  return program

 def update_uniforms(self, uniforms: dict):
  if self.program is None:
   return
  
  for name, value in uniforms.items():
   loc = gl.glGetUniformLocation(self.program, name)
   if loc == -1:
    continue
   
   if isinstance(value, float):
    gl.glUniform1f(loc, value)
   elif isinstance(value, int):
    gl.glUniform1i(loc, value)
   elif isinstance(value, (list, tuple)):
    if len(value) == 2:
     gl.glUniform2f(loc, value[0], value[1])
    elif len(value) == 3:
     gl.glUniform3f(loc, value[0], value[1], value[2])
    elif len(value) == 4:
     gl.glUniform4f(loc, value[0], value[1], value[2], value[3])

 def setup_background_shader(self):
  vertex_src = """
#version 330 core
const vec2 positions[4] = vec2[](
 vec2(-1.0, -1.0),
 vec2( 1.0, -1.0),
 vec2(-1.0,  1.0),
 vec2( 1.0,  1.0)
);
void main() {
 gl_Position = vec4(positions[gl_VertexID], 0.0, 1.0);
}
"""
  fragment_src = """
#version 330 core
uniform float u_time;
uniform vec2 u_resolution;
out vec4 FragColor;
void main() {
 vec2 uv = gl_FragCoord.xy / u_resolution.xy;
 vec3 col = 0.5 + 0.5 * cos(u_time + uv.xyx + vec3(0, 2, 4));
 FragColor = vec4(col, 1.0);
}
"""
  self.bg_program = self.compile_shader(vertex_src, fragment_src)

 def render_background(self, width, height, time):
  if not self.bg_program:
   return
  gl.glUseProgram(self.bg_program)
  u_time_loc = gl.glGetUniformLocation(self.bg_program, "u_time")
  if u_time_loc != -1:
   gl.glUniform1f(u_time_loc, float(time))
  u_res_loc = gl.glGetUniformLocation(self.bg_program, "u_resolution")
  if u_res_loc != -1:
   gl.glUniform2f(u_res_loc, float(width), float(height))
  gl.glDrawArrays(gl.GL_TRIANGLE_STRIP, 0, 4)
  gl.glUseProgram(0)

 def setup_post_process_shader(self):
  vertex_src = """
#version 330 core
const vec2 positions[4] = vec2[](
 vec2(-1.0, -1.0),
 vec2( 1.0, -1.0),
 vec2(-1.0,  1.0),
 vec2( 1.0,  1.0)
);
const vec2 uvs[4] = vec2[](
 vec2(0.0, 0.0),
 vec2(1.0, 0.0),
 vec2(0.0, 1.0),
 vec2(1.0, 1.0)
);
out vec2 v_uv;
void main() {
 gl_Position = vec4(positions[gl_VertexID], 0.0, 1.0);
 v_uv = uvs[gl_VertexID];
}
"""
  fragment_src = """
#version 330 core
in vec2 v_uv;
uniform sampler2D u_texture;
uniform float u_time;
out vec4 FragColor;
void main() {
 vec4 color = texture(u_texture, v_uv);
 float scanline = sin(v_uv.y * 800.0 + u_time * 2.0) * 0.04;
 color.rgb -= scanline;
 FragColor = color;
}
"""
  self.pp_program = self.compile_shader(vertex_src, fragment_src)

 def render_post_process(self, texture_id, width, height, time):
  if not self.pp_program:
   return
  gl.glUseProgram(self.pp_program)
  gl.glActiveTexture(gl.GL_TEXTURE0)
  gl.glBindTexture(gl.GL_TEXTURE_2D, texture_id)
  u_tex_loc = gl.glGetUniformLocation(self.pp_program, "u_texture")
  if u_tex_loc != -1:
   gl.glUniform1i(u_tex_loc, 0)
  u_time_loc = gl.glGetUniformLocation(self.pp_program, "u_time")
  if u_time_loc != -1:
   gl.glUniform1f(u_time_loc, float(time))
  gl.glDrawArrays(gl.GL_TRIANGLE_STRIP, 0, 4)
  gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
  gl.glUseProgram(0)