from OpenGL.GL import *
import array

class VertexBuffer:
  def __init__(self, componentsPerVertex=3):
    self.componentsPerVertex = componentsPerVertex
    self.vertices = []
    self.colors = []
    self.uvs = []
    self.data = None
    self.colorData = None
    self.uvData = None

  def buildData(self):
    hasColors = len(self.colors) > 0
    hasUvs = len(self.uvs) > 0

    vertexCount = len(self.vertices) // self.componentsPerVertex

    dataList = []
    colorList = []
    uvList = []

    for i in range(vertexCount):
      vStart = i * self.componentsPerVertex
      dataList.extend(self.vertices[vStart:vStart+self.componentsPerVertex])

      if hasColors:
        cStart = i * 4  # Assuming RGBA
        colorList.extend(self.colors[cStart:cStart+4])

      if hasUvs:
        uStart = i * 2  # Assuming UV
        uvList.extend(self.uvs[uStart:uStart+2])

    self.data = array.array('f', dataList)

    if hasColors:
      self.colorData = array.array('f', colorList)
    else:
      self.colorData = None

    if hasUvs:
      self.uvData = array.array('f', uvList)
    else:
      self.uvData = None

  def draw(self, mode=GL_TRIANGLES, count=-1):
    if count == -1:
      count = len(self.data) // self.componentsPerVertex

    if count == 0:
      return

    glEnableClientState(GL_VERTEX_ARRAY)
    glVertexPointer(self.componentsPerVertex, GL_FLOAT, 0, self.data.tobytes())

    if self.colorData:
      glEnableClientState(GL_COLOR_ARRAY)
      glColorPointer(4, GL_FLOAT, 0, self.colorData.tobytes())

    if self.uvData:
      glEnableClientState(GL_TEXTURE_COORD_ARRAY)
      glTexCoordPointer(2, GL_FLOAT, 0, self.uvData.tobytes())

    glDrawArrays(mode, 0, count)

    glDisableClientState(GL_VERTEX_ARRAY)
    if self.colorData:
      glDisableClientState(GL_COLOR_ARRAY)
    if self.uvData:
      glDisableClientState(GL_TEXTURE_COORD_ARRAY)

  def clear(self):
    self.vertices = []
    self.colors = []
    self.uvs = []
    self.data = array.array('f')
    self.colorData = None
    self.uvData = None