import array
from enum import Enum
from typing import Any,Iterable,Self, overload

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

#"pip install PyOpenGL"
from OpenGL.GL  import *
from OpenGL.GLU import *



#Note we use `GL_PIXEL_UNPACK_BUFFER` as an intermediary for copying and mapping.  Don't need
#	support for `GL_COPY_WRITE_BUFFER`, but maybe a little confusing and worse performance.  Note we
#	can't be messy with leaving buffers bound there; screws up textures.

class BUFFER_HINT(Enum):
	StreamToGPU = GL_STREAM_DRAW,
	StreamToCPU = GL_STREAM_READ,
	StreamOnGPU = GL_STREAM_COPY,

	DynamicToGPU = GL_DYNAMIC_DRAW,
	DynamicToCPU = GL_DYNAMIC_READ,
	DynamicOnGPU = GL_DYNAMIC_COPY,

	SetOnceUseRepeated = GL_STATIC_DRAW,
	#`GL_STATIC_READ`, `GL_STATIC_COPY` don't really make sense, so aren't exposed

#See "Array Handling" here: https://mcfletch.github.io/pyopengl/documentation/development.html

class BufferGPU:
	__slots__ = ( "size_bytes", "usage_hint", "_handle" )

	size_bytes:int
	usage_hint:BUFFER_HINT
	_handle:int

	def _create_handle( self:Self ):
		self._handle = glGenBuffers(1)
		assert self._handle!=0, "Invalid handle; probably forgot to bind a context!"
	def _remake_store( self:Self ):
		glBindBuffer( GL_PIXEL_UNPACK_BUFFER, self._handle )
		glBufferData( GL_PIXEL_UNPACK_BUFFER, self.size_bytes,None, self.usage_hint.value[0] )
		glBindBuffer( GL_PIXEL_UNPACK_BUFFER, 0 )

	@staticmethod
	def _arr_size( arr:Any ): #Iterable[float]|array.array|np.ndarray
		if type(arr) == array.array:
			return arr.buffer_info()[1] * arr.itemsize
		elif _GLMISC_HAS_NP and type(arr)==np.ndarray:
			return arr.nbytes
		else:
			return sizeof(GLfloat) * len(arr)

	@overload
	def __init__( self:Self, other:Self ): ...
	@overload
	def __init__( self:Self, size_bytes:int, usage_hint:BUFFER_HINT=BUFFER_HINT.SetOnceUseRepeated ): ...
	@overload
	def __init__( self:Self, initial_data:Iterable[float], usage_hint:BUFFER_HINT=BUFFER_HINT.SetOnceUseRepeated ): ...
	@overload
	def __init__( self:Self, initial_data:array.array, usage_hint:BUFFER_HINT=BUFFER_HINT.SetOnceUseRepeated ): ...
	if _GLMISC_HAS_NP:
		@overload
		def __init__( self:Self, initial_data:np.ndarray, usage_hint:BUFFER_HINT=BUFFER_HINT.SetOnceUseRepeated ): ...
	def __init__( self:Self, *args ):
		def _setup( size_bytes:int, usage_hint:BUFFER_HINT ):
			self.size_bytes = size_bytes
			self.usage_hint = usage_hint
			self._create_handle()
			self._remake_store()
		if   len(args) == 2:
			match args:
				case [ int() as size_bytes, BUFFER_HINT() as usage_hint ]:
					_setup( size_bytes, usage_hint )
					return
				case [ initial_data, BUFFER_HINT() as usage_hint ]:
					_setup( BufferGPU._arr_size(initial_data), usage_hint )
					self.data_write(initial_data)
					return
		elif len(args) == 1:
			match args:
				case [ BufferGPU() as other ]:
					_setup( other.size_bytes, other.usage_hint )
					BufferGPU.data_copy( other, self )
					return
				case [ int() as size_bytes ]:
					_setup( size_bytes, BUFFER_HINT.SetOnceUseRepeated )
					return
				case [ initial_data ]:
					_setup( BufferGPU._arr_size(initial_data), BUFFER_HINT.SetOnceUseRepeated )
					self.data_write(initial_data)
					return
		raise TypeError("Invalid argument type(s)!")
	def __del__( self:Self ):
		glDeleteBuffers( 1, [self._handle] )

	def resize_nocopy_invalidate( self:Self,
		new_size_bytes:int, new_usage_hint:BUFFER_HINT=BUFFER_HINT.SetOnceUseRepeated
	):
		if self.size_bytes==new_size_bytes and self.usage_hint==new_usage_hint: return

		#Note: safe to delete 0
		glDeleteBuffers( 1, [self._handle] )
		self._create_handle()

		self.size_bytes = new_size_bytes
		self.usage_hint = new_usage_hint
		self._remake_store()

	@overload
	@staticmethod
	def data_copy( src:Self, dst:Self ): ...
	@overload
	@staticmethod
	def data_copy( src:Self, dst:Self, offset_src:int,offset_dst:int,size_bytes:int=-1 ): ...
	@staticmethod
	def data_copy( src:Self, dst:Self, *args:int ):
		if len(args) == 0:
			assert src.size_bytes==dst.size_bytes, "Size mismatch!"
			offset_src = offset_dst = 0
			size_bytes = src.size_bytes
		else:
			offset_src, offset_dst, size_bytes = args
			if size_bytes == -1:
				size_bytes = min( src.size_bytes-offset_src, dst.size_bytes-offset_dst )
			assert size_bytes>=0, "Access out of bounds!"

		tmp = src.data_read( offset_src,size_bytes )
		dst.draw_write( tmp, offset_dst,size_bytes )

	def data_read( self:Self, offset:int=0,size_bytes:int=-1 ) -> arrays.ArrayDatatype:
		if size_bytes == -1:
			size_bytes = self.size_bytes - offset
			assert size_bytes>=0, "Access out of bounds!"
		else:
			assert offset+size_bytes<=self.size_bytes, "Access out of bounds!"

		glBindBuffer( GL_PIXEL_UNPACK_BUFFER, self._handle )
		data = glGetBufferSubData( GL_PIXEL_UNPACK_BUFFER, offset, size_bytes )
		glBindBuffer( GL_PIXEL_UNPACK_BUFFER, 0 )
		return data
	@overload
	def data_write( self:Self, data:Iterable[float], offset:int=0,size_bytes:int=-1 ): ...
	@overload
	def data_write( self:Self, data:array.array, offset:int=0,size_bytes:int=-1 ): ...
	if _GLMISC_HAS_NP:
		@overload
		def data_write( self:Self, data:np.ndarray, offset:int=0,size_bytes:int=-1 ): ...
	def data_write( self:Self, data:Any, offset:int=0,size_bytes:int=-1 ):
		if type(data) == array.array:
			data_sz = data.buffer_info()[1] * data.itemsize
			ptr     = data.tobytes()
		elif _GLMISC_HAS_NP and type(data)==np.ndarray:
			data_sz = data.nbytes
			ptr     = data
		else:
			return self.data_write( array.array("f",data), offset,size_bytes )

		if size_bytes == -1:
			size_bytes = data_sz - offset
			assert size_bytes>=0, "Access out of bounds!"
		else:
			assert offset+size_bytes<=self.size_bytes and offset+size_bytes<=data_sz, "Access out of bounds!"

		glBindBuffer( GL_PIXEL_PACK_BUFFER, self._handle )
		glBufferSubData( GL_PIXEL_PACK_BUFFER, offset,size_bytes, ptr )
		glBindBuffer( GL_PIXEL_PACK_BUFFER, 0 )

	def map( self:Self, readable:bool,writable:bool, offset:int=0,index_granularity:int=4 ):
		if   readable: access = GL_READ_WRITE if writable else GL_READ_ONLY
		elif writable: access = GL_WRITE_ONLY
		else: raise ValueError("Must be readable, writable, or both!")
		return BufferMapping( self, access, offset,index_granularity )

class BufferMapping:
	""" Encapsulates a buffer mapping of a `BufferGPU`.

	This is difficult, since Python is not typesafe and also does not support assignment
	overloading.  However, there's still a good amount of syntactic sugar.  To use, take your
	buffer, make a mapping (with an index granularity, in bytes, for convenience), offset into it,
	and then read / write the field for the type you want.

	For example, the following writes a single `float` to the 5th float in the buffer (offset 20):
		```py
		mapping = buffer.map( False,True, 0,4 )
		mapping[5].f = 1.0
		del mapping
		```
	You should write the biggest possible blocks at once for performance.  E.g. to write sixteen
	`float`s at once:
		```py
		mapping[5].f_arr = array.array( "f", [ 1,2,3,4, 5,6,7,8, 9,10,11,12, 13,14,15,16 ] )
		```
	"""

	__slots__ = ( "base_addr", "index_granularity" )

	def __init__( self:Self, buffer:BufferGPU, access:GLenum, offset:int=0,index_granularity:int=4 ):
		glBindBuffer( GL_PIXEL_PACK_BUFFER, buffer._handle )
		self.base_addr = glMapBuffer( GL_PIXEL_PACK_BUFFER, access ) + offset
		#self.arr = ( GLfloat * count ).from_address(self.base_addr)
		self.index_granularity = index_granularity
	def __del__( self:Self ):
		glUnmapBuffer(GL_PIXEL_PACK_BUFFER)
		glBindBuffer( GL_PIXEL_PACK_BUFFER, 0 )

	class Ptr:
		__slots__ = ( "offset_addr", )

		def __init__( self:Self, offset_addr:int ):
			self.offset_addr = offset_addr

		def _f_get( self:Self ):
			return ( GLfloat * 1 ).from_address(self.offset_addr)[0]
		def _f_set( self:Self, val:float ):
			( GLfloat * 1 ).from_address(self.offset_addr)[0] = val

		def _vec2f_get( self:Self ):
			bufdat = ( GLfloat * 2 ).from_address(self.offset_addr)
			return ( bufdat[0], bufdat[1] )
		def _vec2f_set( self:Self, vec:Iterable[float] ):
			bufdat = ( GLfloat * 2 ).from_address(self.offset_addr)
			bufdat[0]=vec[0]; bufdat[1]=vec[1]

		def _vec3f_get( self:Self ):
			bufdat = ( GLfloat * 3 ).from_address(self.offset_addr)
			return ( bufdat[0], bufdat[1], bufdat[2] )
		def _vec3f_set( self:Self, vec:Iterable[float] ):
			bufdat = ( GLfloat * 3 ).from_address(self.offset_addr)
			bufdat[0]=vec[0]; bufdat[1]=vec[1]; bufdat[2]=vec[2]

		def _vec4f_get( self:Self ):
			bufdat = ( GLfloat * 4 ).from_address(self.offset_addr)
			return ( bufdat[0], bufdat[1], bufdat[2], bufdat[3] )
		def _vec4f_set( self:Self, vec:Iterable[float] ):
			bufdat = ( GLfloat * 4 ).from_address(self.offset_addr)
			bufdat[0]=vec[0]; bufdat[1]=vec[1]; bufdat[2]=vec[2]; bufdat[3]=vec[3]

		def _matrNxMf_get( self:Self,                       N:int,M:int, Nsrc:int ):
			bufdat = ( GLfloat * (Nsrc*M) ).from_address(self.offset_addr)
			matr = []
			for     j in range(N):
				for i in range(M):
					matr.append( bufdat[ i*Nsrc + j ] )
			return matr
		def _matrNxMf_set( self:Self, matr:Iterable[float], N:int,M:int, Ndst:int ):
			bufdat = ( GLfloat * (Ndst*M) ).from_address(self.offset_addr)
			it = iter(matr) #profiled faster than indexing
			for     j in range(N):
				for i in range(M):
					#bufdat[ i*Ndst + j ] = matr[ j*M + i ]
					bufdat[ i*Ndst + j ] = next(it)
		def _matr2x2f_get( self:Self                       ): return self._matrNxMf_get(       2,2, 2 )
		def _matr2x3f_get( self:Self                       ): return self._matrNxMf_get(       2,3, 2 )
		def _matr2x4f_get( self:Self                       ): return self._matrNxMf_get(       2,4, 2 )
		def _matr3x2f_get( self:Self                       ): return self._matrNxMf_get(       3,2, 3+1 )
		def _matr3x3f_get( self:Self                       ): return self._matrNxMf_get(       3,3, 3+1 )
		def _matr3x4f_get( self:Self                       ): return self._matrNxMf_get(       3,4, 3+1 )
		def _matr4x2f_get( self:Self                       ): return self._matrNxMf_get(       4,2, 4 )
		def _matr4x3f_get( self:Self                       ): return self._matrNxMf_get(       4,3, 4 )
		def _matr4x4f_get( self:Self                       ): return self._matrNxMf_get(       4,4, 4 )
		def _matr2x2f_set( self:Self, matr:Iterable[float] ): return self._matrNxMf_set( matr, 2,2, 2 )
		def _matr2x3f_set( self:Self, matr:Iterable[float] ): return self._matrNxMf_set( matr, 2,3, 2 )
		def _matr2x4f_set( self:Self, matr:Iterable[float] ): return self._matrNxMf_set( matr, 2,4, 2 )
		def _matr3x2f_set( self:Self, matr:Iterable[float] ): return self._matrNxMf_set( matr, 3,2, 3+1 )
		def _matr3x3f_set( self:Self, matr:Iterable[float] ): return self._matrNxMf_set( matr, 3,3, 3+1 )
		def _matr3x4f_set( self:Self, matr:Iterable[float] ): return self._matrNxMf_set( matr, 3,4, 3+1 )
		def _matr4x2f_set( self:Self, matr:Iterable[float] ): return self._matrNxMf_set( matr, 4,2, 4 )
		def _matr4x3f_set( self:Self, matr:Iterable[float] ): return self._matrNxMf_set( matr, 4,3, 4 )
		def _matr4x4f_set( self:Self, matr:Iterable[float] ): return self._matrNxMf_set( matr, 4,4, 4 )

		def _f_arr_set( self:Self, vals:Iterable[float] ):
			bufdat = ( GLfloat * len(n) ).from_address(self.offset_addr)
			for k in range( len(n) ): bufdat[k] = vals[k]

		def _i_get( self:Self ):
			return ( GLfloat * 1 ).from_address(self.offset_addr)[0]
		def _i_set( self:Self, val:float ):
			( GLint * 1 ).from_address(self.offset_addr)[0] = val

		f        = property( fget=_f_get       , fset=_f_set        )
		vec2f    = property( fget=_vec2f_get   , fset=_vec2f_set    )
		vec3f    = property( fget=_vec3f_get   , fset=_vec3f_set    )
		vec4f    = property( fget=_vec4f_get   , fset=_vec4f_set    )
		matr2x2f = property( fget=_matr2x2f_get, fset=_matr2x2f_set )
		matr2x3f = property( fget=_matr2x3f_get, fset=_matr2x3f_set )
		matr2x4f = property( fget=_matr2x4f_get, fset=_matr2x4f_set )
		matr3x2f = property( fget=_matr3x2f_get, fset=_matr3x2f_set )
		matr3x3f = property( fget=_matr3x3f_get, fset=_matr3x3f_set )
		matr3x4f = property( fget=_matr3x4f_get, fset=_matr3x4f_set )
		matr4x2f = property( fget=_matr4x2f_get, fset=_matr4x2f_set )
		matr4x3f = property( fget=_matr4x3f_get, fset=_matr4x3f_set )
		matr4x4f = property( fget=_matr4x4f_get, fset=_matr4x4f_set )
		f_arr    = property(                     fset=_f_arr_set    )
		i        = property( fget=_i_get       , fset=_i_set        )

	def __getitem__( self:Self, offset:int ):
		return BufferMapping.Ptr( self.base_addr + self.index_granularity*offset )
	def __setitem__( self:Self, offset:int, val:Any ):
		raise RuntimeError(
			"Python doesn't support assignment; assign the type's property in the returned pointer"+
			"instead!"
		)
