{-# LANGUAGE TypeSynonymInstances #-}
module Main (main) where

import Control.Monad (when)
import Foreign (Ptr, nullPtr, castPtr, plusPtr, advancePtr, allocaBytes, with, withArray, peek, poke, pokeByteOff)
import Foreign.C (CUChar)
import Foreign.C.String (withCString, peekCString)
import Foreign.ForeignPtr (withForeignPtr)
import System.Exit (exitSuccess)

import Graphics.GL
import Graphics.Rendering.OpenGL.Capture (capturePPM)
import Graphics.UI.GLUT (getArgsAndInitialize, createWindow, displayCallback, ($=), swapBuffers, reportErrors, mainLoop)
import qualified Data.ByteString as BS
import Data.Array.Repa.IO.DevIL (runIL, readImage, Image(RGBA))
import Data.Array.Repa.Repr.ForeignPtr (toForeignPtr)

import Paths_butterflies (getDataFileName)

import Geometry.Flat.TwoD.Space
import Geometry.Flat.TwoD.Tessellation.Triangular

swarm' p q = tessellate p q 3 (Point (-12) (-12)) (Point 12 12)

class GLPoke t where glPoke :: t -> Ptr t -> IO (Ptr t)
instance GLPoke t => GLPoke [t] where
  glPoke [] p = return p
  glPoke (x:xs) p = glPoke x (castPtr p) >>= glPoke xs . castPtr

instance GLPoke GLfloat where glPoke x p = poke p x >> return (advancePtr p 1)
instance GLPoke Double where
  glPoke x p = do
    p <- glPoke (realToFrac x :: GLfloat) (castPtr p)
    return (castPtr p)
instance GLPoke Point where
  glPoke (Point x y) p = do
    p <- glPoke x (castPtr p)
    p <- glPoke y p
    return (castPtr p)
instance GLPoke Vertex where
  glPoke (Vertex vp vt h0 h1 h2) p = do
    p <- glPoke [vp, vt] (castPtr p)
    p <- glPoke [h0, h1, h2] (castPtr p)
    return (castPtr p)
instance GLPoke a => GLPoke (Triangle a) where
  glPoke (Triangle a b c) p = do
    p <- glPoke [a,b,c] (castPtr p)
    return (castPtr p)

vert = unlines
  [ "#version 120"
  , "uniform mat4 mvp;"
  , "attribute vec2 p0;"
  , "attribute vec2 t0;"
  , "attribute vec3 c0;"
  , "varying vec2 t;"
  , "varying vec3 c;"
  , "void main() {"
  , "  gl_Position = vec4(vec2(vec4(p0, 0.0, 1.0) * mvp), 0.0, 1.0);"
  , "  t = t0;"
  , "  c = c0;"
  , "}"
  ]

frag = unlines
  [ "#version 120"
  , "uniform sampler2D tex;"
  , "uniform sampler1D pal;"
  , "varying vec2 t;"
  , "varying vec3 c;"
  , "const float phi1 = (sqrt(5.0) - 1.0) / 2.0;"
  , "const float phi2 = (sqrt(5.0) - 2.0) / 2.0;"
  , "vec4 colour(float i) {"
  , "  float j = 1.0/6.0 - phi1 * i;"
  , "  float k = phi2 * i;"
  , "  j -= floor(j);"
  , "  k -= floor(k);"
  , "  k *= -0.5;"
  , "  k +=  1.0;"
  , "  return vec4(texture1D(pal, j).rgb * k, 1.0);"
  , "}"
  , "void main() {"
  , "  vec4 w = texture2D(tex, t);"
  , "  vec4 f = vec4(0.0);"
  , "  if (w.a <= 0.5) {"
  , "    f = vec4(0.5, 0.5, 0.5, 1.0);"
  , "  } else if (w.r >= 0.5 && w.g >= 0.5 && w.b >= 0.5) {"
  , "    f = vec4(1.0, 1.0, 1.0, 1.0);"
  , "  } else if (w.r <= 0.5 && w.g <= 0.5 && w.b <= 0.5) {"
  , "    f = vec4(0.0, 0.0, 0.0, 1.0);"
  , "  } else {"
  , "    float k = 0.0;"
  , "    if (w.r >= 0.5) {"
  , "      k = c.r;"
  , "    } else if (w.g >= 0.5) {"
  , "      k = c.g;"
  , "    } else if (w.b >= 0.5) {"
  , "      k = c.b;"
  , "    }"
  , "    f = colour(k);"
  , "  }"
  , "  gl_FragColor = f;"
  , "}"
  ]

main = do
  (_, [sp, sq]) <- getArgsAndInitialize
  _ <- createWindow "butterflies"
  program <- compileProgram
  glUseProgram program
  loadTextures
  let ip = read sp
      iq = read sq
      swarm = swarm' ip iq
      count = length swarm * 3
      stride = 4 * (2 + 2 + 3)
      bytes = count * fromIntegral stride
  allocaBytes bytes $ \p -> do
    glPoke swarm p
    vbo <- with 0 $ \q -> glGenBuffers 1 q >> peek q
    glBindBuffer GL_ARRAY_BUFFER vbo
    glBufferData GL_ARRAY_BUFFER (fromIntegral bytes) p GL_STATIC_DRAW
    att <- withCString "p0" $ glGetAttribLocation program
    glVertexAttribPointer (fromIntegral att) 2 GL_FLOAT (fromIntegral GL_FALSE) stride (plusPtr nullPtr 0)
    glEnableVertexAttribArray (fromIntegral att)
    att <- withCString "t0" $ glGetAttribLocation program
    glVertexAttribPointer (fromIntegral att) 2 GL_FLOAT (fromIntegral GL_FALSE) stride (plusPtr nullPtr (2 * 4))
    glEnableVertexAttribArray (fromIntegral att)
    att <- withCString "c0" $ glGetAttribLocation program
    glVertexAttribPointer (fromIntegral att) 3 GL_FLOAT (fromIntegral GL_FALSE) stride (plusPtr nullPtr (4 * 4))
    glEnableVertexAttribArray (fromIntegral att)
  let s = 10
      l = -s
      r = s
      t = -s
      b = s
      n = -1
      f = 1
      ortho = [ 2 / (r - l), 0, 0, -(r + l) / (r - l)
              , 0, 2 / (t - b), 0, -(t + b) / (t - b)
              , 0, 0, 2 / (f - n), -(f + n) / (f - n)
              , 0, 0, 0, 1 ]
  withArray ortho $ \p -> do
    loc <- withCString "mvp" $ glGetUniformLocation program
    glUniformMatrix4fv loc 1 (fromIntegral GL_FALSE) p
  loc <- withCString "tex" $ glGetUniformLocation program
  glUniform1i loc 0
  loc <- withCString "pal" $ glGetUniformLocation program
  glUniform1i loc 1
  glClearColor 0.5 0.5 0.5 1
  displayCallback $= do
    glClear GL_COLOR_BUFFER_BIT
    glDrawArrays GL_TRIANGLES 0 (fromIntegral count)
    swapBuffers
    let ipq = ip * ip + iq * iq + ip * iq
    capturePPM >>= BS.writeFile ("tessellation-" ++ (if ipq < 10 then "0" else "") ++ show ipq ++ "-" ++ show ip ++ "-" ++ show iq ++ ".ppm")
    reportErrors
    exitSuccess
  mainLoop

loadTextures = do
  RGBA img <- getDataFileName "butterfly.png" >>= runIL . readImage
  withForeignPtr (toForeignPtr img) $ \p -> do
    tex <- with 0 $ \q -> glGenTextures 1 q >> peek q
    glBindTexture GL_TEXTURE_2D tex
    glTexImage2D GL_TEXTURE_2D 0 (fromIntegral GL_RGBA) 1024 1024 0 GL_RGBA GL_UNSIGNED_BYTE p
    glTexParameteri GL_TEXTURE_2D GL_TEXTURE_MIN_FILTER (fromIntegral GL_LINEAR_MIPMAP_LINEAR)
    glTexParameteri GL_TEXTURE_2D GL_TEXTURE_MAG_FILTER (fromIntegral GL_LINEAR)
    glGenerateMipmap GL_TEXTURE_2D
  glActiveTexture GL_TEXTURE1
  RGBA img <- getDataFileName "palette.png" >>= runIL . readImage
  withForeignPtr (toForeignPtr img) $ \p -> do
    tex <- with 0 $ \q -> glGenTextures 1 q >> peek q
    glBindTexture GL_TEXTURE_1D tex
    glTexImage1D GL_TEXTURE_1D 0 (fromIntegral GL_RGBA) 256 0 GL_RGBA GL_UNSIGNED_BYTE p
    glTexParameteri GL_TEXTURE_1D GL_TEXTURE_MIN_FILTER (fromIntegral GL_NEAREST)
    glTexParameteri GL_TEXTURE_1D GL_TEXTURE_MAG_FILTER (fromIntegral GL_NEAREST)

compileProgram = do
  program <- glCreateProgram
  compileShader program GL_VERTEX_SHADER vert
  compileShader program GL_FRAGMENT_SHADER frag
  glLinkProgram program
  debugProgram program
  return program

compileShader program t src = do
  shader <- glCreateShader t
  withCString src $ \srcp -> with srcp $ \srcpp -> glShaderSource shader 1 srcpp nullPtr
  glCompileShader shader
  glAttachShader program shader
  debugShader shader
  glDeleteShader shader

debugProgram program = do
  if program /= 0
    then do
      linked <- with 0 $ \p -> glGetProgramiv program GL_LINK_STATUS p >> peek p
      when (linked /= fromIntegral GL_TRUE) $ putStrLn "link failed"
      len <- with 0 $ \p -> glGetProgramiv program GL_INFO_LOG_LENGTH p >> peek p
      when (len > 1) $ do
        allocaBytes (fromIntegral len + 1) $ \p -> glGetProgramInfoLog program len nullPtr p >> pokeByteOff p (fromIntegral len) (0 :: CUChar) >> peekCString p >>= putStrLn
    else putStrLn "no program"

debugShader shader = do
  if shader /= 0
    then do
      compiled <- with 0 $ \p -> glGetShaderiv shader GL_COMPILE_STATUS p >> peek p
      when (compiled /= fromIntegral GL_TRUE) $ putStrLn "compile failed"
      len <- with 0 $ \p -> glGetShaderiv shader GL_INFO_LOG_LENGTH p >> peek p
      when (len > 1) $ do
        allocaBytes (fromIntegral len + 1) $ \p -> glGetShaderInfoLog shader len nullPtr p >> pokeByteOff p (fromIntegral len) (0 :: CUChar) >> peekCString p >>= putStrLn
    else putStrLn "no shader"