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

#"pip install PyOpenGL"
import OpenGL.extensions
from OpenGL.GL import *
from OpenGL.constant import Constant as _C
import OpenGL.GL.ARB.framebuffer_object as fboarb
import OpenGL.GL.EXT.framebuffer_object as fboext
import OpenGL.raw.GL.EXT.framebuffer_blit as fboblitext

glGenRenderbuffers        = OpenGL.extensions.alternate( glGenRenderbuffers       , fboarb.glGenRenderbuffers   , fboext.glGenRenderbuffersEXT        )
glBindRenderbuffer        = OpenGL.extensions.alternate( glBindRenderbuffer       , fboarb.glBindRenderbuffer   , fboext.glBindRenderbufferEXT        )
glRenderbufferStorage     = OpenGL.extensions.alternate( glRenderbufferStorage    , fboarb.glRenderbufferStorage, fboext.glRenderbufferStorageEXT     )
glDeleteRenderbuffers     = OpenGL.extensions.alternate( glDeleteRenderbuffers    , fboarb.glDeleteRenderbuffers, fboext.glDeleteRenderbuffersEXT     )

assert fboarb.GL_RENDERBUFFER     == GL_RENDERBUFFER
assert fboext.GL_RENDERBUFFER_EXT == GL_RENDERBUFFER

glGenFramebuffers         = OpenGL.extensions.alternate( glGenFramebuffers        , fboarb.glGenFramebuffers        , fboext.glGenFramebuffersEXT         )
glDeleteFramebuffers      = OpenGL.extensions.alternate( glDeleteFramebuffers     , fboarb.glDeleteFramebuffers     , fboext.glDeleteFramebuffersEXT      )
glFramebufferTexture2D    = OpenGL.extensions.alternate( glFramebufferTexture2D   , fboarb.glFramebufferTexture2D   , fboext.glFramebufferTexture2DEXT    )
glFramebufferRenderbuffer = OpenGL.extensions.alternate( glFramebufferRenderbuffer, fboarb.glFramebufferRenderbuffer, fboext.glFramebufferRenderbufferEXT )
glBindFramebuffer         = OpenGL.extensions.alternate( glBindFramebuffer        , fboarb.glBindFramebuffer        , fboext.glBindFramebufferEXT         )
glBlitFramebuffer         = OpenGL.extensions.alternate( glBlitFramebuffer        , fboarb.glBlitFramebuffer        , fboblitext.glBlitFramebufferEXT     )
glCheckFramebufferStatus  = OpenGL.extensions.alternate( glCheckFramebufferStatus , fboarb.glCheckFramebufferStatus , fboext.glCheckFramebufferStatusEXT  )

assert fboarb.GL_FRAMEBUFFER     == GL_FRAMEBUFFER
assert fboext.GL_FRAMEBUFFER_EXT == GL_FRAMEBUFFER

assert fboarb    .GL_DRAW_FRAMEBUFFER     == GL_DRAW_FRAMEBUFFER
assert fboblitext.GL_DRAW_FRAMEBUFFER_EXT == GL_DRAW_FRAMEBUFFER
assert fboarb    .GL_READ_FRAMEBUFFER     == GL_READ_FRAMEBUFFER
assert fboblitext.GL_READ_FRAMEBUFFER_EXT == GL_READ_FRAMEBUFFER

assert fboarb.GL_FRAMEBUFFER_COMPLETE                          == GL_FRAMEBUFFER_COMPLETE
assert fboext.GL_FRAMEBUFFER_COMPLETE_EXT                      == GL_FRAMEBUFFER_COMPLETE
assert fboarb.GL_FRAMEBUFFER_INCOMPLETE_ATTACHMENT             == GL_FRAMEBUFFER_INCOMPLETE_ATTACHMENT
assert fboext.GL_FRAMEBUFFER_INCOMPLETE_ATTACHMENT_EXT         == GL_FRAMEBUFFER_INCOMPLETE_ATTACHMENT
assert fboarb.GL_FRAMEBUFFER_INCOMPLETE_MISSING_ATTACHMENT     == GL_FRAMEBUFFER_INCOMPLETE_MISSING_ATTACHMENT
assert fboext.GL_FRAMEBUFFER_INCOMPLETE_MISSING_ATTACHMENT_EXT == GL_FRAMEBUFFER_INCOMPLETE_MISSING_ATTACHMENT
assert fboarb.GL_FRAMEBUFFER_INCOMPLETE_DRAW_BUFFER            == GL_FRAMEBUFFER_INCOMPLETE_DRAW_BUFFER
assert fboext.GL_FRAMEBUFFER_INCOMPLETE_DRAW_BUFFER_EXT        == GL_FRAMEBUFFER_INCOMPLETE_DRAW_BUFFER
assert fboarb.GL_FRAMEBUFFER_INCOMPLETE_READ_BUFFER            == GL_FRAMEBUFFER_INCOMPLETE_READ_BUFFER
assert fboext.GL_FRAMEBUFFER_INCOMPLETE_READ_BUFFER_EXT        == GL_FRAMEBUFFER_INCOMPLETE_READ_BUFFER
assert fboarb.GL_FRAMEBUFFER_UNSUPPORTED                       == GL_FRAMEBUFFER_UNSUPPORTED
assert fboext.GL_FRAMEBUFFER_UNSUPPORTED_EXT                   == GL_FRAMEBUFFER_UNSUPPORTED
assert fboarb.GL_FRAMEBUFFER_INCOMPLETE_MULTISAMPLE            == GL_FRAMEBUFFER_INCOMPLETE_MULTISAMPLE
assert fboarb.GL_FRAMEBUFFER_UNDEFINED                         == GL_FRAMEBUFFER_UNDEFINED

#fboext.GL_FRAMEBUFFER_INCOMPLETE_DIMENSIONS_EXT
#fboext.GL_FRAMEBUFFER_INCOMPLETE_FORMATS_EXT

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



class AttachmentBase:
	__slots__ = ( "gl_attachpt", "gl_buffer", "res", "nspp" )

	def __init__( self:Self, res:Iterable[int],nspp:int, gl_attachpt:_C,gl_buffer:_C ):
		self.gl_attachpt = gl_attachpt
		self.gl_buffer   = gl_buffer

		self.res:tuple[int,int] = tuple(res)
		self.nspp = nspp



if _GLFBO_HAS_TEX:

	class AttachmentTextureBase(AttachmentBase):
		__slots__ = ( "texture", )

		def __init__( self:Self, tex:gl_texture.Texture2D, gl_attachpt:_C,gl_buffer:_C ):
			AttachmentBase.__init__( self, tex.res,tex.nspp, gl_attachpt,gl_buffer )
			self.texture = tex

	class AttachmentTextureColor(AttachmentTextureBase):
		__slots__ = ()

		def __init__( self:Self, tex:gl_texture.Texture2D, index:int ):
			arg = GL_COLOR_ATTACHMENT0 + index
			AttachmentTextureBase.__init__( self, tex, arg,arg )

	class AttachmentTextureDepthStencil(AttachmentTextureBase):
		__slots__ = ()

		def __init__( self:Self, tex:gl_texture.Texture2D ):
			AttachmentTextureBase.__init__( self, tex, GL_DEPTH_ATTACHMENT,GL_NONE )



class AttachmentRenderbufferBase(AttachmentBase):
	__slots__ = ( "_handle", )

	def __init__( self:Self, res:Iterable[int],nspp:int, gl_attachpt:_C,gl_buffer:_C, gl_fmt:_C ):
		AttachmentBase.__init__( self, res,nspp, gl_attachpt,gl_buffer )

		if nspp != 1 and not glRenderbufferStorageMultisample:
			raise RuntimeError("Multisample RBO not available!")

		self._handle = glGenRenderbuffers(1)
		glBindRenderbuffer( GL_RENDERBUFFER, self._handle )
		if nspp == 1:
			glRenderbufferStorage( GL_RENDERBUFFER, gl_fmt, res[0],res[1] )
		else:
			glRenderbufferStorageMultisample( GL_RENDERBUFFER, nspp, gl_fmt, res[0],res[1] )

	def __del__( self:Self ):
		glDeleteRenderbuffers( 1, [self._handle] )

class AttachmentRenderbufferColor(AttachmentRenderbufferBase):
	__slots__ = ()

	def __init__( self:Self, res:Iterable[int],nspp:int, gl_fmt:_C, index:int ):
		arg = GL_COLOR_ATTACHMENT0 + index
		AttachmentRenderbufferBase.__init__( self, res, nspp, arg,arg, gl_fmt )

class AttachmentRenderbufferDepthStencil(AttachmentRenderbufferBase):
	__slots__ = ()

	def __init__( self:Self, res:Iterable[int],nspp:int, gl_fmt:_C ):
		AttachmentRenderbufferBase.__init__( self, res, nspp, GL_DEPTH_ATTACHMENT,GL_NONE, gl_fmt )



class FBO2D:
	max_attachments:ClassVar[int] = -1

	current_read:Self|None = None
	current_draw:Self|None = None

	__slots__ = ( "res", "nspp", "attachments_color","attachment_depthstencil", "_handle" )

	def __init__( self:Self, res:Iterable[int], nspp:int=1 ):
		if FBO2D.max_attachments == -1:
			FBO2D.max_attachments = glGetIntegerv(GL_MAX_COLOR_ATTACHMENTS)

		assert (type(res)==tuple or type(res)==list) and len(res)==2 and res[0]>0 and res[1]>0
		self.res:tuple[int,int] = tuple(res)

		assert type(nspp)==int and nspp>=1
		self.nspp = nspp

		self.attachments_color      :list[AttachmentBase|None] = [ None ]*8
		self.attachment_depthstencil:AttachmentBase|None       = None

		self._handle = glGenFramebuffers(1)

	def __del__( self:Self ):
		for attachment in self.attachments_color + [ self.attachment_depthstencil ]:
			if attachment != None:
				del attachment

		glDeleteFramebuffers( 1, [self._handle] )

	@overload
	def set_attachment_color( self:Self, renbuf_glfmt:_C, index:int ) -> None: ...
	if _GLFBO_HAS_TEX:
		@overload
		def set_attachment_color( self:Self, tex_attach:gl_texture.Texture2D, index:int ) -> None: ...
	def set_attachment_color( self:Self, arg:Any, index:int ):
		if _GLFBO_HAS_TEX and type(arg)==gl_texture.Texture2D:
			tex = arg
			self.attachments_color[index] = AttachmentTextureColor( tex, index )
			self._set_attachment_texture( self.attachments_color[index] )
		else:
			gl_fmt = arg
			self.attachments_color[index] = AttachmentRenderbufferColor(
				self.res,self.nspp, gl_fmt, index
			)
			self._set_attachment_renderbuffer( self.attachments_color[index] )

	@overload
	def set_attachment_depthstencil( self:Self, renbuf_glfmt:_C ) -> None: ...
	if _GLFBO_HAS_TEX:
		@overload
		def set_attachment_depthstencil( self:Self, tex_attach:gl_texture.Texture2D ) -> None: ...
	def set_attachment_depthstencil( self:Self, arg:Any ):
		if _GLFBO_HAS_TEX and type(arg)==gl_texture.Texture2D:
			tex = arg
			self.attachment_depthstencil = AttachmentTextureDepthStencil(tex)
			self._set_attachment_texture(self.attachment_depthstencil)
		else:
			gl_fmt = arg
			self.attachment_depthstencil = AttachmentRenderbufferDepthStencil(
				self.res,self.nspp, gl_fmt
			)
			self._set_attachment_renderbuffer(self.attachment_depthstencil)

	if _GLFBO_HAS_TEX:
		def _set_attachment_texture     ( self:Self, attachment:AttachmentTextureBase      ):
			assert attachment.texture.nspp==self.nspp,\
				"Expected texture samples per texel ({attachment.texture.nspp}) to match FBO's ({self.nspp})!"
			FBO2D.set_readdraw(self)
			glFramebufferTexture2D( GL_FRAMEBUFFER,
				attachment.gl_attachpt, attachment.texture.gl_type, attachment.texture._handle, 0
			)

	def _set_attachment_renderbuffer( self:Self, attachment:AttachmentRenderbufferBase ):
		FBO2D.set_readdraw(self)
		glFramebufferRenderbuffer( GL_FRAMEBUFFER,
			attachment.gl_attachpt, GL_RENDERBUFFER, attachment._handle
		)

	@staticmethod
	def set_readdraw( fbo:Self|None=None ):
		FBO2D.current_read = fbo
		FBO2D.current_draw = fbo
		handle = 0 if fbo==None else fbo._handle
		glBindFramebuffer( GL_FRAMEBUFFER, handle )

	@staticmethod
	def set_read( fbo_read:Self|None=None, color_index:int|str|None="first" ):
		assert type(fbo_read   ) in (FBO2D,  ) or fbo_read   ==None
		assert type(color_index) in (int,str,) or color_index==None

		FBO2D.current_read = fbo_read
		handle = 0 if fbo_read==None else fbo_read._handle

		glBindFramebuffer( GL_READ_FRAMEBUFFER, handle )

		if fbo_read == None:
			glReadBuffer(GL_BACK)
		else:
			if type(color_index) == int:
				glReadBuffer( fbo_read.attachments_color[color_index].gl_buffer )
			elif color_index == None:
				glReadBuffer(GL_NONE)
			else:
				assert len(fbo_read.attachments_color)>0, "FBO has no color attachments!"
				if color_index == "first":
					color_attachment = fbo_read.attachments_color[ 0].gl_buffer
				elif color_index == "last":
					color_attachment = fbo_read.attachments_color[-1].gl_buffer
				else:
					raise ValueError(f"Unknown color index \"{color_index}\"")
				#Note OpenGL does not support multiple read buffers
				glReadBuffer(color_attachment)

	@staticmethod
	def set_draw( fbo_draw:Self|None=None, color_indices:list[int]|str="all" ):
		assert type(fbo_draw     ) in (FBO2D,       ) or fbo_draw==None
		assert type(color_indices) in (list[int],str)

		FBO2D.current_draw = fbo_draw
		handle = 0 if fbo_draw==None else fbo_draw._handle

		glBindFramebuffer( GL_DRAW_FRAMEBUFFER, handle )

		if fbo_draw == None:
			glDrawBuffer(GL_BACK)
		else:
			if color_indices == []:
				glDrawBuffer(GL_NONE)
			elif type(color_indices) == list:
				color_attachments = [ GL_COLOR_ATTACHMENT0+ind for ind in color_indices ]
				glDrawBuffers(color_attachments)
			else:
				assert color_indices == "all"
				color_attachments = []
				for attachment_color in fbo_draw.attachments_color:
					if attachment_color == None: continue
					color_attachments.append(attachment_color.gl_buffer)
				glDrawBuffers(color_attachments)

	@staticmethod
	def copy(
		src_rect:tuple[ int,int, int,int ],
		dst_rect:tuple[ int,int, int,int ],
		kind:_C = GL_COLOR_BUFFER_BIT,
		filtering:_C = GL_NEAREST
	):
		glBlitFramebuffer(
			src_rect[0], src_rect[1], src_rect[0]+src_rect[2], src_rect[1]+src_rect[3],
			dst_rect[0], dst_rect[1], dst_rect[0]+dst_rect[2], dst_rect[1]+dst_rect[3],
			kind,
			filtering
		)

	def get_status_str( self:Self ):
		FBO2D.set_readdraw(self)

		status = glCheckFramebufferStatus(GL_FRAMEBUFFER)
		if   status == GL_FRAMEBUFFER_COMPLETE                     :
			return    "GL_FRAMEBUFFER_COMPLETE"
		elif status == GL_FRAMEBUFFER_UNDEFINED                    :
			return    "GL_FRAMEBUFFER_UNDEFINED"
		elif status == GL_FRAMEBUFFER_INCOMPLETE_ATTACHMENT        :
			return    "GL_FRAMEBUFFER_INCOMPLETE_ATTACHMENT"
		elif status == GL_FRAMEBUFFER_INCOMPLETE_MISSING_ATTACHMENT:
			return    "GL_FRAMEBUFFER_INCOMPLETE_MISSING_ATTACHMENT"
		elif status == GL_FRAMEBUFFER_INCOMPLETE_DRAW_BUFFER       :
			return    "GL_FRAMEBUFFER_INCOMPLETE_DRAW_BUFFER"
		elif status == GL_FRAMEBUFFER_INCOMPLETE_READ_BUFFER       :
			return    "GL_FRAMEBUFFER_INCOMPLETE_READ_BUFFER"
		elif status == GL_FRAMEBUFFER_UNSUPPORTED                  :
			return    "GL_FRAMEBUFFER_UNSUPPORTED"
		elif status == GL_FRAMEBUFFER_INCOMPLETE_MULTISAMPLE       :
			return    "GL_FRAMEBUFFER_INCOMPLETE_MULTISAMPLE"
		elif status == GL_FRAMEBUFFER_INCOMPLETE_LAYER_TARGETS     :
			return    "GL_FRAMEBUFFER_INCOMPLETE_LAYER_TARGETS"
		else                                                       :
			return f"(Unknown FBO error {status})"
