from typing import Any,Callable,ClassVar,Iterable,Self, overload
import weakref

#"pip install numpy"
try:
	import numpy as np
	_GLSHAD_HAS_NP = True
except ModuleNotFoundError:
	_GLSHAD_HAS_NP = False

#"pip install PyOpenGL"
import OpenGL.extensions
from OpenGL.GL import *
from OpenGL.constant import Constant as _C
from OpenGL.GL.ARB.shader_objects   import *
from OpenGL.GL.ARB.vertex_shader    import *
from OpenGL.GL.ARB.geometry_shader4 import *
from OpenGL.GL.ARB.fragment_shader  import *

glCreateShader       = OpenGL.extensions.alternate( glCreateShader      , glCreateShaderObjectARB   )
glCompileShader      = OpenGL.extensions.alternate( glCompileShader     , glCompileShaderARB        )
glGetShaderInfoLog   = OpenGL.extensions.alternate( glGetShaderInfoLog  , glGetInfoLogARB           )
glDeleteShader       = OpenGL.extensions.alternate( glDeleteShader      , glDeleteObjectARB         )
glShaderSource       = OpenGL.extensions.alternate( glShaderSource      , glShaderSourceARB         )

assert GL_VERTEX_SHADER_ARB == GL_VERTEX_SHADER
#PyOpenGL doesn't add `_ARB` to `GL_TESS_CONTROL_SHADER`?
#PyOpenGL doesn't add `_ARB` to `GL_TESS_EVALUATION_SHADER`?
assert GL_GEOMETRY_SHADER_ARB == GL_GEOMETRY_SHADER
assert GL_FRAGMENT_SHADER_ARB == GL_FRAGMENT_SHADER

glAttachShader       = OpenGL.extensions.alternate( glAttachShader      , glAttachObjectARB         )
glCreateProgram      = OpenGL.extensions.alternate( glCreateProgram     , glCreateProgramObjectARB  )
glDeleteProgram      = OpenGL.extensions.alternate( glDeleteProgram     , glDeleteObjectARB         )
glGetProgramiv       = OpenGL.extensions.alternate( glGetProgramiv      , glGetObjectParameterivARB )
glGetProgramInfoLog  = OpenGL.extensions.alternate( glGetProgramInfoLog , glGetInfoLogARB           )
glLinkProgram        = OpenGL.extensions.alternate( glLinkProgram       , glLinkProgramARB          )
glValidateProgram    = OpenGL.extensions.alternate( glValidateProgram   , glValidateProgramARB      )
glUseProgram         = OpenGL.extensions.alternate( glUseProgram        , glUseProgramObjectARB     )

glGetUniformLocation = OpenGL.extensions.alternate( glGetUniformLocation, glGetUniformLocationARB   )
glUniform1f          = OpenGL.extensions.alternate( glUniform1f         , glUniform1fARB            )
glUniform2f          = OpenGL.extensions.alternate( glUniform2f         , glUniform2fARB            )
glUniform3f          = OpenGL.extensions.alternate( glUniform3f         , glUniform3fARB            )
glUniform4f          = OpenGL.extensions.alternate( glUniform4f         , glUniform4fARB            )
glUniform1i          = OpenGL.extensions.alternate( glUniform1i         , glUniform1iARB            )
glUniform2i          = OpenGL.extensions.alternate( glUniform2i         , glUniform2iARB            )
glUniform3i          = OpenGL.extensions.alternate( glUniform3i         , glUniform3iARB            )
glUniform4i          = OpenGL.extensions.alternate( glUniform4i         , glUniform4iARB            )
glUniformMatrix2fv   = OpenGL.extensions.alternate( glUniformMatrix2fv  , glUniformMatrix2fvARB     )
#glUniformMatrix3x2fv (no alternates)
#glUniformMatrix4x2fv
#glUniformMatrix2x3fv
glUniformMatrix3fv   = OpenGL.extensions.alternate( glUniformMatrix3fv  , glUniformMatrix3fvARB     )
#glUniformMatrix4x3fv (no alternates)
#glUniformMatrix2x4fv
#glUniformMatrix3x4fv
glUniformMatrix4fv   = OpenGL.extensions.alternate( glUniformMatrix4fv  , glUniformMatrix4fvARB     )

assert GL_OBJECT_LINK_STATUS_ARB     == GL_LINK_STATUS
assert GL_OBJECT_VALIDATE_STATUS_ARB == GL_VALIDATE_STATUS

try:
	#(another basecode file, optional)
	from . import gl_buffer
	_GLSHAD_HAS_BUF = True
except ModuleNotFoundError:
	_GLSHAD_HAS_BUF = False

try:
	#(another basecode file, optional)
	from . import gl_texture
	_GLSHAD_HAS_TEX = True
except ModuleNotFoundError:
	_GLSHAD_HAS_TEX = False



def _get_processed_log( log:str|bytes ) -> str:
	if type(log) == bytes:
		return _get_processed_log( log.decode() )
	assert type(log) == str

	return "\n".join([ line for line in log.split("\n") if line ])



class ShaderBase:
	__slots__ = ( "type", "source", "_handle" )

	def __init__( self:Self, type:_C, src:str ):
		self.type = type

		self.source = src

		if not glCreateShader:
			raise RuntimeError( "Shader object not available!" )
		self._handle = glCreateShader(type)

		glShaderSource( self._handle, [src] ) # Note must be array of strings
		glCompileShader(self._handle)
		if glGetShaderiv:
			status = glGetShaderiv( self._handle, GL_COMPILE_STATUS )
		else:
			status = glGetObjectParameterivARB( self._handle, GL_OBJECT_COMPILE_STATUS_ARB )

		if status != GL_TRUE:
			raise RuntimeError( "Shader compile failed:\n" + self.get_log_str() )

	def __del__( self:Self ):
		glDeleteShader(self._handle)

	def get_log_str( self:Self ):
		return _get_processed_log( glGetShaderInfoLog(self._handle) )

class ShaderVert(ShaderBase):
	__slots__ = ()

	def __init__( self:Self, src:str ):
		ShaderBase.__init__( self, GL_VERTEX_SHADER, src )

class ShaderTessCtrl(ShaderBase):
	__slots__ = ()

	def __init__( self:Self, src:str ):
		ShaderBase.__init__( self, GL_TESS_CONTROL_SHADER, src )

class ShaderTessEval(ShaderBase):
	__slots__ = ()

	def __init__( self:Self, src:str ):
		ShaderBase.__init__( self, GL_TESS_EVALUATION_SHADER, src )

class ShaderGeom(ShaderBase):
	__slots__ = ()

	def __init__( self:Self, src:str ):
		ShaderBase.__init__( self, GL_GEOMETRY_SHADER, src )

class ShaderFrag(ShaderBase):
	__slots__ = ()

	def __init__( self:Self, src:str ):
		ShaderBase.__init__( self, GL_FRAGMENT_SHADER, src )



#We need to ensure that `Program.current` doesn't hold a reference to a program, as this can cause
#	horrendous bugs (e.g. trying to `glDeleteProgram(⋯)` on Python shutdown.  Thus, we store a weak
#	reference `Program._current`.  This metaclass makes it appear as `Program.current` and avoiding
#	the parens to access the value.
class _ProgramMeta(type):
	__slots__ = ()

	def __getattr__( cls:Self, item:str ):
		if   item == "current":
			if cls._current == None: return None
			else                   : return cls._current()
		else:
			super().__getattr__(item)
	def __setattr__( cls:Self, item:str, value ):
		if   item == "current":
			if value == None: cls._current = None
			else            : cls._current = weakref.ref(value)
		else:
			super().__setattr__( item, value )

class Program( metaclass=_ProgramMeta ):
	_current:ClassVar[weakref.ref[Self]|None] = None
	current:ClassVar[Self|None]

	#__slots__ = ( ... ) #Need weak-refs

	def __init__( self:Self,
		shaders:list[ShaderBase],
		locs_attribs:dict[str,int] = {},
		locs_outputs:dict[str,int] = {}
	):
		""" Compile a shader program from a list of individual `shaders`.

		Hardcoding the locations of the vertex attribute(s) and fragment output(s) in the shaders
		(e.g. `layout( location=0 )`, GL 3.3+) is best, but you can pass them here.  For example:
			```py
			prog = Program( shaders, {"verts":0,"norms":1}, {"color":0} )
			```
		It is important that the locations be specified somehow, because what happens if you don't
		is undefined.  E.g. NVIDIA seems to number attributes and outputs sequentially, while Intel
		discrete puts them in no discernible order.
		"""

		assert len(shaders) != 0
		for shader in shaders:
			assert isinstance( shader, ShaderBase )

		self.shaders = shaders

		self._handle = glCreateProgram()

		for shader in shaders:
			glAttachShader( self._handle, shader._handle )

		for name, loc in locs_attribs.items():
			glBindAttribLocation( self._handle, loc, name ) #GL 2.0
		for name, loc in locs_outputs.items():
			glBindFragDataLocation( self._handle, loc, name ) #GL 3.0

		self.relink()

	def __del__( self:Self ):
		#Note: if you get weird errors here at program exit, it's because a program reference still
		#	existed at the end of the program.  You must ensure all programs are deleted.
		if Program.current == self:
			Program.use(None)
		glDeleteProgram(self._handle)

	def relink( self:Self ):
		glLinkProgram(self._handle)
		status = glGetProgramiv( self._handle, GL_LINK_STATUS )
		if status != GL_TRUE:
			raise RuntimeError( "Program link failed:\n" + self.get_log_str() )

		#Don't `glValidateProgram(⋯)` here; OpenGL state might not yet be set up to execute this
		#	shader (e.g. texture bindings might imply bogus multisampling use).  See
		#	`.validate_could_execute_now()`.

		self._locs_attribs   :dict[str,int] = {}
		self._locs_fragdatas :dict[str,int] = {}
		self._locs_unifs     :dict[str,int] = {}
		self._blockinds_ssbos:dict[str,int] = {}
		self._blockinds_ubos :dict[str,int] = {}

	def get_log_str( self:Self ):
		return _get_processed_log( glGetProgramInfoLog(self._handle) )

	def save_sources( self:Self, path:str ):
		with open( path, "w" ) as file:
			for shader in self.shaders:
				file.write(shader.source)
				file.write("\n")

	def validate_could_execute_now( self:Self ):
		glValidateProgram(self._handle)
		status = glGetProgramiv( self._handle, GL_VALIDATE_STATUS )
		if status != GL_TRUE:
			raise RuntimeError( "Program validation failed:\n" + self.get_log_str() )

	def get_loc_attrib( self:Self, symbol:str ):
		if not symbol in self._locs_attribs.keys():
			loc = glGetAttribLocation( self._handle, symbol )
			# Note result will be -1 if not found / optimized out
			self._locs_attribs[symbol] = int(loc)
		return self._locs_attribs[symbol]
	def get_loc_fragdata( self:Self, symbol:str ):
		if not symbol in self._locs_fragdatas.keys():
			loc = glGetFragDataLocation( self._handle, symbol )
			# Note result will be -1 if not found / optimized out
			self._locs_fragdatas[symbol] = int(loc)
		return self._locs_fragdatas[symbol]
	def get_loc_unif  ( self:Self, symbol:str ):
		if not symbol in self._locs_unifs.keys():
			loc = glGetUniformLocation( self._handle, symbol )
			# Note result will be -1 if not found / optimized out
			self._locs_unifs[symbol] = int(loc)
		return self._locs_unifs[symbol]
	def get_blockind_ssbo( self:Self, block_symbol:str ):
		if not block_symbol in self._blockinds_ssbos.keys():
			block_ind = glGetProgramResourceIndex( # GL 4.3+
				self._handle, GL_SHADER_STORAGE_BLOCK, block_symbol
			)
			self._blockinds_ssbos[block_symbol] = block_ind
		return self._blockinds_ssbos[block_symbol]
	def get_blockind_ubo ( self:Self, block_symbol:str ):
		if not block_symbol in self._blockinds_ubos.keys():
			#block_ind = glGetProgramResourceIndex(
			#	self._handle, GL_UNIFORM_BLOCK, block_symbol
			#)
			block_ind = glGetUniformBlockIndex( self._handle, block_symbol ) # GL 3.1
			self._blockinds_ubos[block_symbol] = block_ind
		return self._blockinds_ubos[block_symbol]



	def _pass_unif( self:Self, glfn:Callable, symbol:str, *args ):
		assert Program.current == self
		loc = self.get_loc_unif(symbol)
		glfn( loc, *args )
		return loc



	# Pass Scalars

	def pass_1f( self:Self, symbol:str, val:float ):
		assert type(val) in (float,int)
		return self._pass_unif( glUniform1f, symbol, val )

	def pass_1b( self:Self, symbol:str, val:bool ):
		assert type(val) == bool
		return self.pass_1i( symbol, int(val) )

	def pass_1i( self:Self, symbol:str, val:int ):
		assert type(val) == int
		return self._pass_unif( glUniform1i, symbol, val )

	def pass_1u( self:Self, symbol:str, val:int ):
		assert type(val)==int and val>=0
		return self.pass_1i( symbol, val )



	# Pass Vectors

	@overload
	def pass_2f( self:Self, symbol:str, vec:Iterable[float] ) -> int: ...
	@overload
	def pass_2f( self:Self, symbol:str, x:float,y:float ) -> int: ...
	def pass_2f( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self._pass_unif( glUniform2f, symbol, args[0],args[1] )
	@overload
	def pass_3f( self:Self, symbol:str, vec:Iterable[float] ) -> int: ...
	@overload
	def pass_3f( self:Self, symbol:str, x:float,y:float,z:float ) -> int: ...
	def pass_3f( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self._pass_unif( glUniform3f, symbol, args[0],args[1],args[2] )
	@overload
	def pass_4f( self:Self, symbol:str, vec:Iterable[float] ) -> int: ...
	@overload
	def pass_4f( self:Self, symbol:str, x:float,y:float,z:float,w:float ) -> int: ...
	def pass_4f( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self._pass_unif( glUniform4f, symbol, args[0],args[1],args[2],args[3] )

	@overload
	def pass_2b( self:Self, symbol:str, vec:Iterable[bool] ) -> int: ...
	@overload
	def pass_2b( self:Self, symbol:str, x:bool,y:bool ) -> int: ...
	def pass_2b( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self.pass_2i( symbol, int(args[0]),int(args[1]) )
	@overload
	def pass_3b( self:Self, symbol:str, vec:Iterable[bool] ) -> int: ...
	@overload
	def pass_3b( self:Self, symbol:str, x:bool,y:bool,z:bool ) -> int: ...
	def pass_3b( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self.pass_3i( symbol, int(args[0]),int(args[1]),int(args[2]) )
	@overload
	def pass_4b( self:Self, symbol:str, vec:Iterable[bool] ) -> int: ...
	@overload
	def pass_4b( self:Self, symbol:str, x:bool,y:bool,z:bool,w:bool ) -> int: ...
	def pass_4b( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self.pass_4i( symbol, int(args[0]),int(args[1]),int(args[2]),int(args[3]) )

	@overload
	def pass_2i( self:Self, symbol:str, vec:Iterable[int] ) -> int: ...
	@overload
	def pass_2i( self:Self, symbol:str, x:int,y:int ) -> int: ...
	def pass_2i( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self._pass_unif( glUniform2i, symbol, args[0],args[1] )
	@overload
	def pass_3i( self:Self, symbol:str, vec:Iterable[int] ) -> int: ...
	@overload
	def pass_3i( self:Self, symbol:str, x:int,y:int,z:int ) -> int: ...
	def pass_3i( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self._pass_unif( glUniform3i, symbol, args[0],args[1],args[2] )
	@overload
	def pass_4i( self:Self, symbol:str, vec:Iterable[int] ) -> int: ...
	@overload
	def pass_4i( self:Self, symbol:str, x:int,y:int,z:int,w:int ) -> int: ...
	def pass_4i( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self._pass_unif( glUniform4i, symbol, args[0],args[1],args[2],args[3] )

	@overload
	def pass_2u( self:Self, symbol:str, vec:Iterable[int] ) -> int: ...
	@overload
	def pass_2u( self:Self, symbol:str, x:int,y:int ) -> int: ...
	def pass_2u( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self.pass_2i( symbol, int(args[0]),int(args[1]) )
	@overload
	def pass_3u( self:Self, symbol:str, vec:Iterable[int] ) -> int: ...
	@overload
	def pass_3u( self:Self, symbol:str, x:int,y:int,z:int ) -> int: ...
	def pass_3u( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self.pass_3i( symbol, int(args[0]),int(args[1]),int(args[2]) )
	@overload
	def pass_4u( self:Self, symbol:str, vec:Iterable[int] ) -> int: ...
	@overload
	def pass_4u( self:Self, symbol:str, x:int,y:int,z:int,w:int ) -> int: ...
	def pass_4u( self:Self, symbol:str, *args ):
		if len(args) == 1: args = args[0]
		return self.pass_4i( symbol, int(args[0]),int(args[1]),int(args[2]),int(args[3]) )



	# Pass Matrices
	# Note e.g. "3x4" means 3 rows, 4 columns, as is standard in mathematics.

	def _pass_unif_matr(
		self:Self, symbol:str, rowmaj_matr:Any,shape:tuple[int,int], transpose:bool, glfn:Callable
	):
		#TODO: check shape
		if _GLSHAD_HAS_NP and type(rowmaj_matr)==np.ndarray:
			rowmaj_data = list( rowmaj_matr.reshape(-1) ) #TODO: no cast to list
		else:
			rowmaj_data = list(rowmaj_matr)
		transpose_arg = not transpose #OpenGL wants it in evil column-major format
		return self._pass_unif( glfn, symbol, 1, transpose_arg, rowmaj_data )

	def pass_2x2f( self:Self, symbol:str, rowmaj_matr:Any, transpose:bool=False ):
		self._pass_unif_matr( symbol, rowmaj_matr,(2,2), transpose, glUniformMatrix2fv )
	def pass_2x3f( self:Self, symbol:str, rowmaj_matr:Any, transpose:bool=False ):
		self._pass_unif_matr( symbol, rowmaj_matr,(2,3), transpose, glUniformMatrix3x2fv )
	def pass_2x4f( self:Self, symbol:str, rowmaj_matr:Any, transpose:bool=False ):
		self._pass_unif_matr( symbol, rowmaj_matr,(2,4), transpose, glUniformMatrix4x2fv )
	def pass_3x2f( self:Self, symbol:str, rowmaj_matr:Any, transpose:bool=False ):
		self._pass_unif_matr( symbol, rowmaj_matr,(3,2), transpose, glUniformMatrix2x3fv )
	def pass_3x3f( self:Self, symbol:str, rowmaj_matr:Any, transpose:bool=False ):
		self._pass_unif_matr( symbol, rowmaj_matr,(3,3), transpose, glUniformMatrix3fv )
	def pass_3x4f( self:Self, symbol:str, rowmaj_matr:Any, transpose:bool=False ):
		self._pass_unif_matr( symbol, rowmaj_matr,(3,4), transpose, glUniformMatrix4x3fv )
	def pass_4x2f( self:Self, symbol:str, rowmaj_matr:Any, transpose:bool=False ):
		self._pass_unif_matr( symbol, rowmaj_matr,(4,2), transpose, glUniformMatrix2x4fv )
	def pass_4x3f( self:Self, symbol:str, rowmaj_matr:Any, transpose:bool=False ):
		self._pass_unif_matr( symbol, rowmaj_matr,(4,3), transpose, glUniformMatrix3x4fv )
	def pass_4x4f( self:Self, symbol:str, rowmaj_matr:Any, transpose:bool=False ):
		self._pass_unif_matr( symbol, rowmaj_matr,(4,4), transpose, glUniformMatrix4fv )
	def pass_4x4f_gl2modelview( self:Self, symbol:str ):
		colmaj_matr = glGetFloatv(GL_MODELVIEW_MATRIX)
		return self.pass_4x4f( symbol, colmaj_matr, True ) #Note transpose
	def pass_4x4f_gl2proj     ( self:Self, symbol:str ):
		colmaj_matr = glGetFloatv(GL_PROJECTION_MATRIX)
		return self.pass_4x4f( symbol, colmaj_matr, True ) #Note transpose



	# Pass other kinds of data

	if _GLSHAD_HAS_TEX:
		def pass_bind_tex( self:Self, symbol:str, tex:gl_texture.TextureBase, tex_unit_ind:int=0 ):
			gl_texture.TextureBase.bind( tex, tex_unit_ind )
			self.pass_1i( symbol, tex_unit_ind )

	if _GLSHAD_HAS_BUF:
		def pass_bind_ssbo( self:Self, buffer:gl_buffer.BufferGPU, bind_loc:int, block_symbol:str|None ):
			""" Pass SSBO to shader (GL 4.3+)

			Binds the SSBO to binding point `bind_loc`.  With GL 4.2+, you can specify the same bind
			location in the shader, for example:
				```glsl
				layout( binding = 0 ) buffer Foo { /*...*/ };
				```
			in which case we're done and `block_symbol` should be `None`.  Otherwise, `block_symbol`
			should give the name of the SSBO block (`"Foo"`, for the above example).  From this, the
			shader will be queried to find the SSBO's block index to connect to the binding point.
			"""
			if block_symbol != None:
				glShaderStorageBlockBinding(
					self._handle, self.get_blockind_ssbo(block_symbol), bind_loc
				)

			glBindBufferBase( GL_SHADER_STORAGE_BUFFER, bind_loc, buffer._handle )

		def pass_bind_ubo ( self:Self, buffer:gl_buffer.BufferGPU, bind_loc:int, block_symbol:str|None ):
			""" Pass UBO to shader (GL 3.1+)

			Binds the UBO to binding point `bind_loc`.  With GL 4.2+, you can specify the same bind
			location in the shader, for example:
				```glsl
				layout( binding = 0 ) uniform Foo { /*...*/ };
				```
			in which case we're done and `block_symbol` should be `None`.  Otherwise, `block_symbol`
			should give the name of the UBO block (`"Foo"`, for the above example).  From this, the
			shader will be queried to find the UBO's block index to connect to the binding point.
			"""
			if block_symbol != None:
				glUniformBlockBinding( self._handle, self.get_blockind_ubo(block_symbol), bind_loc )

			glBindBufferBase( GL_UNIFORM_BUFFER, bind_loc, buffer._handle )



	@staticmethod
	def use( program:Self|None=None ):
		if program == Program.current: return

		handle = 0 if program==None else program._handle
		glUseProgram(handle)
		Program.current = program
