module Graphics.LambdaCube.Loader.SkeletonXML where

import Graphics.LambdaCube.Skeleton
import Graphics.LambdaCube.Types

import Control.Applicative
import Data.List as L
import Data.Map as Map
import Data.Maybe
import Data.Tree
import Data.Vector as V
import Prelude as P
import Text.XML.Light

readFloat :: String -> String -> Element -> Float
readFloat n v e = read $ fromMaybe v $ findAttr (unqual n) e

readStr :: String -> String -> Element -> String
readStr n v e = fromMaybe v $ findAttr (unqual n) e

readBone :: Element -> (String, Bone)
readBone x = (name,Bone (p pos) (rotU (p axis) (readFloat "angle" "angle" rot)))
  where
    name = readStr "name" "boneName" x
    Just pos = findElement (unqual "position") x
    Just rot = findElement (unqual "rotation") x
    Just axis = findElement (unqual "axis") rot
    p a = Vec3 (readFloat "x" "x" a) (readFloat "y" "y" a) (readFloat "z" "z" a)

readKeyFrame :: Element -> KeyFrame
readKeyFrame x = KeyFrame (p pos) (rotU (p axis) (readFloat "angle" "angle" rot)) (readFloat "time" "time" x)
  where
    Just pos = findElement (unqual "translate") x
    Just rot = findElement (unqual "rotate") x
    Just axis = findElement (unqual "axis") rot
    p a = Vec3 (readFloat "x" "x" a) (readFloat "y" "y" a) (readFloat "z" "z" a)

readTrack :: Element -> (String, [KeyFrame])
readTrack x = (bone,keyframes)
  where
    bone = readStr "bone" "trackBone" x
    keyframes = P.map readKeyFrame $ findElements (unqual "keyframe") keyframesn
    Just keyframesn = findElement (unqual "keyframes") x

readAnimation :: [String] -> Element -> (String, Animation)
readAnimation nameList x = (name,Animation tracks len name)
  where
    name = readStr "name" "animName" x
    len = readFloat "length" "animLength" x
    tracks = V.fromList $ P.map snd $ sortBy (\(a,_) (b,_) -> a `compare` b) tl
    tl = [(fromJust $ L.elemIndex n nameList, t) | (n,t) <- P.map readTrack $ findElements (unqual "track") x]
    --Just tracksl = findElement (unqual "tracks") x

readSkeleton :: Element -> Skeleton
readSkeleton x = Skeleton (V.fromList boneList) boneTree animations
  where
    Just bones = findElement (unqual "bones") x
    (nameList,boneList) = L.unzip $ P.map readBone $ findElements (unqual "bone") bones
    --Just hierarchy = findElement (unqual "bonehierarchy") x
    [boneNTree] = L.foldl' add [Node n [] | n <- nameList] $ findElements (unqual "boneparent") x
    boneTree = fmap (\n -> fromJust $ L.elemIndex n nameList) boneNTree
    bon n = readStr "bone" "bone" n
    par n = readStr "parent" "parent" n
    add [e] _ = [e]
    add l n = (ins p):cs
      where
        ([p],ps) = L.partition (\tn -> L.elem (par n) $ flatten tn) l
        ([c],cs) = L.partition (\(Node a _) -> a == bon n) ps
        ins (Node nn nt) | nn == par n = Node nn (c:nt)
                         | otherwise = Node nn $ L.map ins nt
    Just anims = findElement (unqual "animations") x
    animations = Map.fromList $ P.map (readAnimation nameList) $ findElements (unqual "animation") anims

parseSkeleton :: String -> Skeleton
parseSkeleton doc = readSkeleton $ fromMaybe (error "fromJust 6") $ findElement (unqual "skeleton") x
  where
    Just x = parseXMLDoc doc

loadSkeleton :: FilePath -> IO Skeleton
loadSkeleton fileName = parseSkeleton <$> readFile fileName