{-# LANGUAGE FlexibleContexts #-}
module Wavefront where

import Wavefront.Types
import Wavefront.Lex

import Control.Arrow
import Control.Monad.State
import qualified Data.ByteString.Lazy.Char8 as L
import qualified Data.DList                 as DList
         ;import Data.DList                 (DList)
import Data.Functor.Identity
import qualified Data.IntMap                as IntMap
import Data.Text.Encoding                   (decodeUtf8)
import qualified Data.Vector                as Vector
         ;import Data.Vector                (Vector)
import qualified Rank2

type WavefrontOBJ = OBJ Vector

newtype Count x = Count Int


incrementCount :: MonadState s m => (s -> Count x) -> (Count x -> s -> s) -> m ()
incrementCount field setField = do
    Count c0 <- gets field
    let c = succ c0
    c `seq` modify (setField $ Count c)

fixupRef :: Count x -> Int -> Int
fixupRef (Count n) x | x >  0    = x - 1 -- Renumber from 0.
                     | otherwise = n + x -- Negative values are relative.

fixupTriple :: OBJ Count -> RefTriple -> RefTriple
fixupTriple o (RefTriple v t n) =
    RefTriple (fixupRef (objLocations o) v)
              (fixupRef (objTexCoords o) <$> t)
              (fixupRef (objNormals   o) <$> n)

objBookKeeping :: Monad m => ObjBuilder m -> ObjBuilder (StateT (OBJ Count) m)
objBookKeeping builder = (lift Rank2.<$> builder)
    { vertex = \xs -> do lift $ vertex builder xs
                         incrementCount objLocations $ \x o -> o { objLocations = x }
    , vertexT = \xs -> do lift $ vertexT builder xs
                          incrementCount objTexCoords $ \x o -> o { objTexCoords = x }
    , vertexN = \xs -> do lift $ vertexN builder xs
                          incrementCount objNormals $ \x o -> o { objNormals = x }
    , points = \xs -> do
        n <- gets objLocations
        lift $ points builder $ fixupRef n <$> xs
    , line = \ts -> do
        o <- get
        lift $ line builder $ fixupTriple o <$> ts
    , face = \ts -> do
        o <- get
        lift $ face builder $ fixupTriple o <$> ts
    }

mkv :: [Double] -> Location
mkv cs = Location x y z w where (x:y:z:w:_) = map realToFrac cs

mkt :: [Double] -> TexCoord
mkt cs = TexCoord x y z where (x:y:z:_) = map realToFrac cs

mkn :: [Double] -> Normal
mkn cs = Normal x y z where (x:y:z:_) = map realToFrac cs

mkl :: RefTriple -> RefTriple -> Line
mkl (RefTriple a at _) (RefTriple b bt _) = Line (LineIndex a at) (LineIndex b bt)

-- I'd have thought these would be Coercible, but I guess not.
mkF :: RefTriple -> FaceIndex
mkF (RefTriple a at an) = FaceIndex a at an

elemental :: Element () -> x -> Element x
elemental element x = fmap (const x) element

modifyFirst :: MonadState (c, d) m => (c -> c) -> m ()
modifyFirst = modify' . first

buildOBJ :: ObjBuilder (State (OBJ DList,Element ()))
buildOBJ = nullBuilder
    { vertex = \xs -> modifyFirst $ \o -> o { objLocations = objLocations o `DList.snoc` mkv xs }
    , vertexT = \xs -> modifyFirst $ \o -> o { objTexCoords = objTexCoords o `DList.snoc` mkt xs }
    , vertexN = \xs -> modifyFirst $ \o -> o { objNormals = objNormals o `DList.snoc` mkn xs }
    , points = \xs -> do
        let p = map Point xs :: [Point]
        (pts,element) <- gets (objPoints *** elemental)
        modifyFirst $ \o -> o { objPoints = pts `DList.append` fmap element (DList.fromList p) }
    , line = \xs -> do
        (lns,element) <- gets (objLines *** elemental)
        let l = zipWith mkl xs (tail xs)
        -- Line requires at least two points.  We'll ignore it otherwise.
        when (not $ null l) $
            modifyFirst $ \o -> o { objLines = lns `DList.append` fmap element (DList.fromList l) }
    , face = \xs -> do
        (fcs,element) <- gets (objFaces *** elemental)
        case map mkF xs of
          a:b:c:ds -> modifyFirst $ \o -> o { objFaces = fcs `DList.snoc` element (Face a b c ds) }
          _        -> return () -- Ignore faces with less than 3 indices.
    , mtllib = \xs -> do
        let l = map decodeUtf8 xs
        libs <- gets (objMtlLibs . fst)
        modifyFirst $ \o -> o { objMtlLibs = libs `DList.append` DList.fromList l }
    , groups = \xs -> do
        let g = map decodeUtf8 xs
        modify' $ second $ \e -> e { elGroups = g }
    , objectName = \x -> do
        let o = decodeUtf8 x
        modify' $ second $ \e -> e { elObject = Just o }
    , usemtl = \x -> do
        let mtl = decodeUtf8 x
        modify' $ second $ \e -> e { elMtl = Just mtl }
    , smoothingGroup = \x -> when (x > 0) $ do
        modify $ second $ \e -> e { elSmoothingGroup = fromIntegral x }
    }

parse :: L.ByteString -> OBJ Vector
parse bs = Rank2.fmap (Vector.fromList . DList.toList) obj
 where
    b = objBookKeeping buildOBJ
    ls = OBJ
            { objLocations = DList.empty
            , objTexCoords = DList.empty
            , objNormals   = DList.empty
            , objPoints    = DList.empty
            , objLines     = DList.empty
            , objFaces     = DList.empty
            , objMtlLibs   = DList.empty
            }
    c = Rank2.fmap (const $ Count 0) ls :: OBJ Count
    el = Element
            { elObject         = Nothing
            , elGroups         = []
            , elMtl            = Nothing
            , elSmoothingGroup = 0
            , elValue          = ()
            }
    substvars = ObjConfig IntMap.empty
    (obj,_) = execState (runStateT (parseOBJ b substvars bs) c) (ls,el)