import os
import sys
from typing import Callable, ClassVar, Iterable, Self
import weakref

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

#"pip install PyOpenGL"
import OpenGL.extensions
from OpenGL.GL  import *
from OpenGL.GLU import *
from OpenGL.constant import Constant as _C
from OpenGL.GL.EXT.texture_filter_anisotropic import *
#Could have sworn there was `glGenerateMipmapARB(⋯)` somewhere too . . .

from OpenGL.GL.ARB.framebuffer_object import glGenerateMipmap as glGenerateMipmap_ARB
from OpenGL.GL.EXT.framebuffer_object import glGenerateMipmapEXT
glGenerateMipmap = OpenGL.extensions.alternate( glGenerateMipmap, glGenerateMipmap_ARB, glGenerateMipmapEXT )
_GTTEX_INITED:bool = False
_GLTEX_HAS_ANIS:bool|None = None

#"pip install pygame-ce"
with open( os.devnull, "w" ) as f:
	oldstdout=sys.stdout; sys.stdout=f
	try:
		import pygame
		_GLTEX_HAS_PYGAME = True
	except ModuleNotFoundError:
		_GLTEX_HAS_PYGAME = False
	sys.stdout = oldstdout



class TextureFormatInfo:
	__slots__ = ( "fmt_base", "type_transfer", "num_chans", "pixel_sz" )

	def __init__( self:Self, fmt_base:_C, type_transfer:_C, num_chans:int, pixel_sz:int ):
		self.fmt_base = fmt_base
		self.type_transfer = type_transfer
		self.num_chans = num_chans
		self.pixel_sz = pixel_sz



class TextureBase:
	#__slots__ = ( "_handle", "gl_type","gl_fmt_real" ) #Need weakrefs

	_active_tex_unit_ind:ClassVar[int] = 0

	num_units:ClassVar[int]
	#We need to ensure that `TextureBase.current` doesn't hold a reference to a texture, as this can
	#	cause horrendous bugs (e.g. trying to `glDeleteTextures(⋯)` on Python shutdown, or
	#	`TextureBase.bind(⋯)` with a non-None texture immediately causing `None` to be bound right
	#	after.
	current:ClassVar[weakref.WeakValueDictionary[ int, Self ]] = weakref.WeakValueDictionary()

	base_fmt:ClassVar[dict[ _C, TextureFormatInfo ]] = {
		fmt_real : TextureFormatInfo( fmt_base, type_transfer, num_chans, pixel_sz )

		for fmt_real, fmt_base, type_transfer, num_chans, pixel_sz in [
			( GL_DEPTH_COMPONENT16 , GL_DEPTH_COMPONENT , GL_UNSIGNED_SHORT, 1,  2 ),
			( GL_DEPTH_COMPONENT24 , GL_DEPTH_COMPONENT , GL_UNSIGNED_INT  , 1,  4 ),
			( GL_DEPTH_COMPONENT32 , GL_DEPTH_COMPONENT , GL_UNSIGNED_INT  , 1,  4 ),
			( GL_DEPTH_COMPONENT32F, GL_DEPTH_COMPONENT , GL_FLOAT         , 1,  4 ),
			( GL_DEPTH24_STENCIL8  , GL_DEPTH_STENCIL   , GL_UNSIGNED_INT  , 2,  4 ),
			( GL_DEPTH32F_STENCIL8 , GL_DEPTH_STENCIL   , GL_FLOAT         , 2,  5 ),
			( GL_STENCIL_INDEX8    , GL_STENCIL_INDEX   , GL_UNSIGNED_BYTE , 1,  1 ),

			( GL_R8                , GL_RED             , GL_UNSIGNED_BYTE , 1,  1 ),
			( GL_R8_SNORM          , GL_RED             , GL_BYTE          , 1,  1 ),
			( GL_R8UI              , GL_RED_INTEGER     , GL_UNSIGNED_BYTE , 1,  1 ),
			( GL_R8I               , GL_RED_INTEGER     , GL_BYTE          , 1,  1 ),
			( GL_R16               , GL_RED             , GL_UNSIGNED_SHORT, 1,  2 ),
			( GL_R16_SNORM         , GL_RED             , GL_SHORT         , 1,  2 ),
			( GL_R16UI             , GL_RED_INTEGER     , GL_UNSIGNED_SHORT, 1,  2 ),
			( GL_R16I              , GL_RED_INTEGER     , GL_SHORT         , 1,  2 ),
			( GL_R16F              , GL_RED             , GL_HALF_FLOAT    , 1,  2 ),
			( GL_R32UI             , GL_RED_INTEGER     , GL_UNSIGNED_INT  , 1,  4 ),
			( GL_R32I              , GL_RED_INTEGER     , GL_INT           , 1,  4 ),
			( GL_R32F              , GL_RED             , GL_FLOAT         , 1,  4 ),

			( GL_RG8               , GL_RG              , GL_UNSIGNED_BYTE , 2,  2 ),
			( GL_RG8_SNORM         , GL_RG              , GL_BYTE          , 2,  2 ),
			( GL_RG8UI             , GL_RG_INTEGER      , GL_UNSIGNED_BYTE , 2,  2 ),
			( GL_RG8I              , GL_RG_INTEGER      , GL_BYTE          , 2,  2 ),
			( GL_RG16              , GL_RG              , GL_UNSIGNED_SHORT, 2,  4 ),
			( GL_RG16_SNORM        , GL_RG              , GL_SHORT         , 2,  4 ),
			( GL_RG16UI            , GL_RG_INTEGER      , GL_UNSIGNED_SHORT, 2,  4 ),
			( GL_RG16I             , GL_RG_INTEGER      , GL_SHORT         , 2,  4 ),
			( GL_RG16F             , GL_RG              , GL_HALF_FLOAT    , 2,  4 ),
			( GL_RG32UI            , GL_RG_INTEGER      , GL_UNSIGNED_INT  , 2,  8 ),
			( GL_RG32I             , GL_RG_INTEGER      , GL_INT           , 2,  8 ),
			( GL_RG32F             , GL_RG              , GL_FLOAT         , 2,  8 ),

			( GL_R3_G3_B2          , GL_RGB             , GL_UNSIGNED_BYTE , 3,  1 ),
			( GL_RGB4              , GL_RGB             , GL_UNSIGNED_BYTE , 3,  2 ),
			( GL_RGB5              , GL_RGB             , GL_UNSIGNED_BYTE , 3,  2 ),
			( GL_RGB8              , GL_RGB             , GL_UNSIGNED_BYTE , 3,  3 ),
			( GL_RGB8_SNORM        , GL_RGB             , GL_BYTE          , 3,  3 ),
			( GL_RGB8UI            , GL_RGB_INTEGER     , GL_UNSIGNED_BYTE , 3,  3 ),
			( GL_RGB8I             , GL_RGB_INTEGER     , GL_BYTE          , 3,  3 ),
			( GL_RGB9_E5           , GL_RGB             , GL_FLOAT         , 3,  4 ),
			( GL_RGB10             , GL_RGB             , GL_UNSIGNED_SHORT, 3,  4 ),
			( GL_R11F_G11F_B10F    , GL_RGB             , GL_FLOAT         , 3,  4 ),
			( GL_RGB12             , GL_RGB             , GL_UNSIGNED_SHORT, 3,  6 ),
			( GL_RGB16             , GL_RGB             , GL_UNSIGNED_SHORT, 3,  6 ),
			( GL_RGB16_SNORM       , GL_RGB             , GL_SHORT         , 3,  6 ),
			( GL_RGB16UI           , GL_RGB_INTEGER     , GL_UNSIGNED_SHORT, 3,  6 ),
			( GL_RGB16I            , GL_RGB_INTEGER     , GL_SHORT         , 3,  6 ),
			( GL_RGB16F            , GL_RGB             , GL_HALF_FLOAT    , 3,  6 ),
			( GL_RGB32UI           , GL_RGB_INTEGER     , GL_UNSIGNED_INT  , 3, 12 ),
			( GL_RGB32I            , GL_RGB_INTEGER     , GL_INT           , 3, 12 ),
			( GL_RGB32F            , GL_RGB             , GL_FLOAT         , 3, 12 ),

			( GL_RGBA2             , GL_RGBA            , GL_UNSIGNED_BYTE , 4,  1 ), # Assuming docs meant RGBA
			( GL_RGBA4             , GL_RGBA            , GL_UNSIGNED_BYTE , 4,  2 ), # Assuming docs meant RGBA
			( GL_RGB5_A1           , GL_RGBA            , GL_UNSIGNED_BYTE , 4,  2 ),
			( GL_RGBA8             , GL_RGBA            , GL_UNSIGNED_BYTE , 4,  4 ),
			( GL_RGBA8_SNORM       , GL_RGBA            , GL_BYTE          , 4,  4 ),
			( GL_RGBA8UI           , GL_RGBA_INTEGER    , GL_UNSIGNED_BYTE , 4,  4 ),
			( GL_RGBA8I            , GL_RGBA_INTEGER    , GL_BYTE          , 4,  4 ),
			( GL_RGB10_A2          , GL_RGBA            , GL_UNSIGNED_BYTE , 4,  4 ),
			( GL_RGB10_A2UI        , GL_RGBA            , GL_UNSIGNED_BYTE , 4,  4 ),
			( GL_RGBA12            , GL_RGBA            , GL_UNSIGNED_BYTE , 4,  8 ),
			( GL_RGBA16            , GL_RGBA            , GL_UNSIGNED_SHORT, 4,  8 ),
			( GL_RGBA16_SNORM      , GL_RGBA            , GL_SHORT         , 4,  8 ),
			( GL_RGBA16UI          , GL_RGBA_INTEGER    , GL_UNSIGNED_SHORT, 4,  8 ),
			( GL_RGBA16I           , GL_RGBA_INTEGER    , GL_SHORT         , 4,  8 ),
			( GL_RGBA16F           , GL_RGBA            , GL_HALF_FLOAT    , 4,  8 ),
			( GL_RGBA32UI          , GL_RGBA_INTEGER    , GL_UNSIGNED_INT  , 4, 16 ),
			( GL_RGBA32I           , GL_RGBA_INTEGER    , GL_INT           , 4, 16 ),
			( GL_RGBA32F           , GL_RGBA            , GL_FLOAT         , 4, 16 ),

			( GL_SRGB8             , GL_RGB             , GL_UNSIGNED_BYTE , 3,  3 ),
			( GL_SRGB8_ALPHA8      , GL_RGBA            , GL_UNSIGNED_BYTE , 4,  4 )
		]
	}

	def __init__( self:Self, gl_type:_C, gl_fmt:_C ):
		global _GTTEX_INITED, _GLTEX_HAS_ANIS
		if not _GTTEX_INITED:
			TextureBase.num_units = glGetIntegerv(GL_MAX_COMBINED_TEXTURE_IMAGE_UNITS)

			_GLTEX_HAS_ANIS = OpenGL.extensions.hasExtension("GL_EXT_texture_filter_anisotropic")
			if _GLTEX_HAS_ANIS:
				Texture2D.max_anis = glGetFloatv(GL_MAX_TEXTURE_MAX_ANISOTROPY_EXT)

			_GTTEX_INITED = True

		self._handle = glGenTextures(1)

		self.gl_type = gl_type #E.g. `GL_TEXTURE_2D`, `GL_TEXTURE_2D_MULTISAMPLE`

		if   gl_fmt == GL_DEPTH_COMPONENT: self.gl_fmt_real=GL_DEPTH_COMPONENT32F
		elif gl_fmt == GL_DEPTH_STENCIL  : self.gl_fmt_real=GL_DEPTH24_STENCIL8
		elif gl_fmt == GL_RED            : self.gl_fmt_real=GL_R8
		elif gl_fmt == GL_RG             : self.gl_fmt_real=GL_RG8
		elif gl_fmt == GL_RGB            : self.gl_fmt_real=GL_SRGB8
		elif gl_fmt == GL_RGBA           : self.gl_fmt_real=GL_SRGB8_ALPHA8
		else                             : self.gl_fmt_real=gl_fmt

	def __del__( self:Self ):
		#Note: if you get weird errors here at program exit, it's because a texture reference still
		#	existed at the end of the program.  You must ensure all textures are deleted.
		for tex_unit_ind in self.get_bound_tex_units():
			TextureBase.bind( None, tex_unit_ind )
		glDeleteTextures( 1, [self._handle] )

	@staticmethod
	def bind( tex:Self|None, tex_unit_ind:int=0 ):
		if TextureBase.current.get( tex_unit_ind, None ) == tex:
			return

		if tex_unit_ind != TextureBase._active_tex_unit_ind:
			glActiveTexture( GL_TEXTURE0 + tex_unit_ind )
			TextureBase._active_tex_unit_ind = tex_unit_ind

		if tex != None:
			glBindTexture( tex.gl_type, tex._handle )
			TextureBase.current[tex_unit_ind] = tex

		else:
			glBindTexture( GL_TEXTURE_2D, 0 )
			TextureBase.current.pop( tex_unit_ind, None )

	def get_bound_tex_units( self:Self ):
		""" Gets the set of texture units on which this texture is bound. """
		s:set[int] = set()
		for tex_unit_ind, tex in TextureBase.current.items():
			if tex == self:
				s.add(tex_unit_ind)
		return s

	def regen_mips( self:Self ):
		TextureBase.bind(self)

		if not glGenerateMipmap:
			raise RuntimeError( "Mipmapping not available!" )

		glGenerateMipmap(self.gl_type)



class Texture2D(TextureBase):
	max_anis:ClassVar[float] = 1.0

	__slots__ = ( "nspp", "res" )

	@property
	def w( self:Self ): return self.res[0]
	@property
	def h( self:Self ): return self.res[1]

	#TODO: use overload to refactor construction

	def __init__( self:Self, res:Iterable[int], gl_fmt:_C, nspp:int=1 ):
		assert len(res)==2 and res[0]>0 and res[1]>0
		assert nspp >= 1

		if nspp == 1: gl_type=GL_TEXTURE_2D
		else        : gl_type=GL_TEXTURE_2D_MULTISAMPLE
		TextureBase.__init__( self, gl_type, gl_fmt )

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

		TextureBase.bind(self)
		if nspp == 1:
			fmt_info = TextureBase.base_fmt[gl_fmt]
			glTexImage2D( self.gl_type,
				0,
				self.gl_fmt_real,
				res[0], res[1],
				0,
				fmt_info.fmt_base, fmt_info.type_transfer, None
			)
			glTexParameteri( self.gl_type, GL_TEXTURE_MIN_FILTER, GL_NEAREST )
			glTexParameteri( self.gl_type, GL_TEXTURE_MAG_FILTER, GL_NEAREST )
		else:
			#OpenGL 3.2+
			glTexImage2DMultisample( self.gl_type,
				nspp,
				self.gl_fmt_real,
				res[0], res[1],
				False
			)

	@staticmethod
	def from_empty( res:Iterable[int], gl_fmt:_C=GL_SRGB8_ALPHA8, nspp:int=1 ):
		return Texture2D( res, gl_fmt, nspp )

	if _GLTEX_HAS_NP:
		@staticmethod
		def from_arr( arr:np.ndarray, gl_fmt:_C|None=None ):
			assert type(arr) == np.ndarray

			dim = len(arr.shape)
			if   dim == 1:
				return Texture2D.from_arr( arr.view().reshape(1,arr.shape[0],1), gl_fmt )
			elif dim == 2:
				return Texture2D.from_arr( arr.view().reshape(arr.shape[0],arr.shape[1],1), gl_fmt )
			elif dim == 3:
				assert arr.shape[2] in (1,2,3,4) # Unknown number of channels
			else:
				assert False # Unknown dimensionality

			assert arr.dtype in ( "float","float32","float64", "uint8" )
			if gl_fmt == None:
				if   arr.shape[2] == 1:
					if arr.dtype == "uint8": gl_fmt=GL_R8
					else                   : gl_fmt=GL_R32F
				elif arr.shape[2] == 2:
					if arr.dtype == "uint8": gl_fmt=GL_RG8
					else                   : gl_fmt=GL_RG32F
				elif arr.shape[2] == 3:
					if arr.dtype == "uint8": gl_fmt=GL_SRGB8_ALPHA8
					else                   : gl_fmt=GL_RGB32F
				else:
					if arr.dtype == "uint8": gl_fmt=GL_SRGB8_ALPHA8
					else                   : gl_fmt=GL_RGBA32F

			tex = Texture2D.from_empty( arr.shape[:2], gl_fmt )
			tex.set_from_arr(arr)
			return tex

		@staticmethod
		def to_arr( tex:Self, miplvl:int=0 ):
			TextureBase.bind(tex)

			fmt_info = TextureBase.base_fmt[tex.gl_fmt_real]
			glPixelStorei( GL_PACK_ALIGNMENT, 1 )
			data:np.ndarray = glGetTexImage( tex.gl_type,
				miplvl, fmt_info.fmt_base,fmt_info.type_transfer, None, -1
			)
			assert type(data) == np.ndarray

			dim = len(data.shape)
			if   dim == 1: data.shape=( 1, data.shape[0], 1 )
			if   dim == 2: data.shape=( data.shape[0], data.shape[1], 1 )
			elif dim == 3: pass
			else         : assert False # Unknown dimensionality

			assert data.shape[0]==tex.res[0] and data.shape[1]==tex.res[1] and data.shape[2]==fmt_info.num_chans

			return data

		@staticmethod
		def to_bytes( tex:Self, miplvl:int=0 ):
			arr = Texture2D.to_arr( tex, miplvl )
			return arr.tobytes()

		def set_from_arr( self:Self, arr:np.ndarray ):
			assert type(arr) == np.ndarray
			assert arr.size == self.res[1]*self.res[0]

			dim = len(arr.shape)
			if   dim == 1:
				return self.set_from_arr( arr.view().reshape(1,arr.shape[0],1) )
			elif dim == 2:
				return self.set_from_arr( arr.view().reshape(arr.shape[0],arr.shape[1],1) )
			elif dim == 3:
				assert arr.shape[2] in (1,2,3,4) # Unknown number of channels
			else:
				assert False # Unknown dimensionality
			assert arr.dtype in ( "float","float32","float64", "uint8" )

			fmt_info = TextureBase.base_fmt[self.gl_fmt_real]

			TextureBase.bind(self)
			glPixelStorei( GL_UNPACK_ALIGNMENT, 1 )
			glTexSubImage2D( self.gl_type,
				0,
				0,0, self.res[0],self.res[1],
				fmt_info.fmt_base, fmt_info.type_transfer, arr
			)

	if _GLTEX_HAS_PYGAME:
		@staticmethod
		def from_path( path:str, gl_fmt:_C|None=None ):
			surf = pygame.image.load(path)
			return Texture2D.from_surf( surf, gl_fmt )

		@staticmethod
		def from_surf( surf:pygame.Surface, gl_fmt:_C|None=None ):
			if gl_fmt == None:
				num_chans = surf.get_bytesize()
				if num_chans == 4: gl_fmt = GL_SRGB8_ALPHA8
				else:
					assert num_chans in ( 1, 2, 3 ) # Unknown size
					gl_fmt = GL_SRGB8
			else:
				num_chans = TextureBase.base_fmt[gl_fmt].num_chans
				assert num_chans in ( 3, 4 ) # Unknown size for surface

			if num_chans == 4: fmt_str = "RGBA"
			else             : fmt_str = "RGB"
			data = pygame.image.tobytes( surf, fmt_str, True )

			tex = Texture2D.from_empty( surf.get_size(), gl_fmt )
			tex.set_from_bytes(data)
			return tex

		@staticmethod
		def to_surf( tex:Self, miplvl:int=0 ):
			data = Texture2D.to_bytes( tex, miplvl )

			num_chans = TextureBase.base_fmt[tex.gl_fmt_real].num_chans
			if num_chans == 3: fmt_str = "RGB"
			else             : fmt_str = "RGBA"
			surf = pygame.image.frombytes( data, tex.res, fmt_str, True )

			return surf

		def set_from_surf( self:Self, surf:pygame.Surface ):
			assert surf.get_size() == self.res

			num_chans = TextureBase.base_fmt[self.gl_fmt_real].num_chans
			if num_chans == 3: fmt_str = "RGB"
			else             : fmt_str = "RGBA"
			data = pygame.image.tobytes( surf, fmt_str, True )

			self.set_from_bytes(data)

	def set_from_bytes( self:Self, data:bytes ): #TODO miplvl:int=0
		assert type(data) == bytes

		fmt_info = TextureBase.base_fmt[self.gl_fmt_real]
		assert len(data) == self.res[1]*self.res[0]*fmt_info.pixel_sz

		TextureBase.bind(self)
		glPixelStorei( GL_UNPACK_ALIGNMENT, 1 )
		glTexSubImage2D( self.gl_type,
			0,
			0,0, self.res[0],self.res[1],
			fmt_info.fmt_base, fmt_info.type_transfer, data
		)

	def set_edge( self:Self, repeat_s:bool,repeat_t:bool ):
		if repeat_s: wrap_s = GL_REPEAT
		else       : wrap_s = GL_CLAMP
		if repeat_t: wrap_t = GL_REPEAT
		else       : wrap_t = GL_CLAMP
		TextureBase.bind(self)
		glTexParameteri( self.gl_type, GL_TEXTURE_WRAP_S,wrap_s )
		glTexParameteri( self.gl_type, GL_TEXTURE_WRAP_T,wrap_t )

	def set_filters( self:Self, filt_min_0near_1lin_2mip:int,filt_max_0near_1lin:int ):
		assert filt_min_0near_1lin_2mip in ( 0, 1, 2 )
		assert filt_max_0near_1lin      in ( 0, 1    )
		enums = [ GL_NEAREST, GL_LINEAR, GL_LINEAR_MIPMAP_LINEAR ]
		TextureBase.bind(self)
		glTexParameteri( self.gl_type, GL_TEXTURE_MIN_FILTER, enums[filt_min_0near_1lin_2mip] )
		glTexParameteri( self.gl_type, GL_TEXTURE_MAG_FILTER, enums[filt_max_0near_1lin     ] )

	def set_anis( self:Self, level:float|bool=True ):
		if _GLTEX_HAS_ANIS:
			if type(level) == bool:
				level = Texture2D.max_anis if level else 0.0
			else:
				assert type(level)==float and level>=0.0
				level = min( level, Texture2D.max_anis )
			glTexParameterf( self.gl_type, GL_TEXTURE_MAX_ANISOTROPY_EXT, level )
			return level
		else:
			raise RuntimeError("Anisotropic filtering not supported!")

	def set_depthmap_compare( self:Self, cmp:GLint=GL_LESS ):
		TextureBase.bind(self)
		glTexParameteri( self.gl_type, GL_TEXTURE_COMPARE_MODE,GL_COMPARE_REF_TO_TEXTURE )
		glTexParameteri( self.gl_type, GL_TEXTURE_COMPARE_FUNC,cmp                       )

	def set_nicest( self:Self ):
		self.set_filters( 2, 1 )

		if _GLTEX_HAS_ANIS:
			self.set_anis()

		self.regen_mips()



if _GLTEX_HAS_PYGAME:
	def surface_grayscale_convert8( grayscale_surf:pygame.Surface ):
		"""Convert 24-bit surface with grayscale data to 8-bit palettal image."""
		#This is currently much harder than it should be.
		#	See https://github.com/pygame-community/pygame-ce/issues/3029
		ret = pygame.Surface( grayscale_surf.get_size(), depth=8 )
		ret.set_palette([ (gr,gr,gr) for gr in range(256) ])
		pygame.surfarray.pixels2d(ret)[:,:] = pygame.surfarray.pixels_red(grayscale_surf)[:,:]
		return ret

	def surface_calc_data_rect( surf:pygame.Surface, fn_data_at:Callable[[pygame.Color],bool] ):
		"""Calculates the rect within the given surface that contains data.

		Args:
			surf:       The surface to inspect.
			fn_data_at: Predicate that decides whether a particular pixel value constitutes "data".

		Returns:
			The rectangle where data is contained.
		"""

		x=0; y=0
		w, h = surf.get_size()

		def _at( i:int, j:int ):
			return fn_data_at( surf.get_at(( i, j )) )
		def _col( i:int ):
			nonlocal y, h
			for j in range( y, y+h, 1 ):
				if _at( i, j ): return True
			return False
		def _row( j:int ):
			nonlocal x, w
			for i in range( x, x+w, 1 ):
				if _at( i, j ): return True
			return False

		while w > 0:
			if _col( w - 1 ): break
			w -= 1
		while w > 0:
			if _col(x): break
			x += 1
			w -= 1

		while h > 0:
			if _row( h - 1 ): break
			h -= 1
		while h > 0:
			if _row(y): break
			y += 1
			h -= 1

		return pygame.Rect( x,y, w,h )
