{-# LANGUAGE TemplateHaskell #-}
{-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}
module Data.Cursor.CLASE.Gen.Util where

import Control.Monad
import Data.Map (Map)
import Data.Maybe
import Language.Haskell.TH
import qualified Data.Map as Map

data DataType = DataType {
                  dtName :: Name,
                  dtCtrs :: [Constructor]
                }
              | WrapsList {
                  dtName :: Name,
                  dtConName :: Name,
                  dtChildName :: Name
                }

data Constructor = Ctr {
                     ctrName :: Name,
                     ctrKids :: [Child]
                   }

data Child 
  = NonNavigable  { childType :: Name }
  | Navigable     { childType :: Name }

isWrapsList :: DataType -> Bool
isWrapsList (WrapsList _ _ _) = True
isWrapsList _ = False

constructors :: DataType -> [Constructor]
constructors d@(DataType {}) = dtCtrs d
constructors (WrapsList _ conName kidName) = [Ctr conName [Navigable kidName]]


extractChildren :: DataType -> [Name]
extractChildren (DataType _ ctrs) = map ctrName ctrs
extractChildren (WrapsList _ cn _) = [cn]

isNavigable :: Child -> Bool
isNavigable (Navigable _) = True
isNavigable _ = False

buildMap :: [Name] -> Q (Map Name DataType)
buildMap okNames = do
  reifiedInfos <- mapM reify okNames
  let dts = map infoToDataType reifiedInfos 
  return . Map.fromList . zip okNames $ dts
  where
    infoToDataType :: Info -> DataType
    infoToDataType info
      | isSoloCtr &&
        onlyOneTypeInCtr &&
        (extractTypeNoList soloType) `elem` okNames &&
        isList soloType     = WrapsList name soloCtrName soloNonListType
      | otherwise = DataType name (map mkCtr ctrs)
      where
        ctrs = extractConstructors info
        name = getInfoName info

        mkCtr :: Con -> Constructor
        mkCtr con = Ctr conName (map mkChild children)
          where
            children = extractChildTypesRaw $ con
            mkChild :: Type -> Child
            mkChild n
              | (extractTypeNoList n) `elem` okNames = Navigable (extractTypeNoList n)
              | otherwise = NonNavigable (extractType n)

            conName = getConName con

        isSoloCtr = length ctrs == 1
        headCtrsTypes = extractChildTypesRaw . head $ ctrs
        soloCtrName = getConName . head $ ctrs
        onlyOneTypeInCtr = length headCtrsTypes == 1
        soloType = head headCtrsTypes
        soloNonListType = extractTypeNoList soloType

data ContextCtr 
  = ListCC {
    ctxCtrName :: String,
    ctxCtrTypeFrom :: Name,
    ctxCtrTypeTo :: Name,
    ctxCtrCtrTo :: Name
    }
  | NormalCC {
    ctxCtrName :: String,
    ctxCtrOffset_ :: Maybe Int,
    ctxCtrTypesBefore :: [Name],
    ctxCtrTypesAfter :: [Name],
    ctxCtrTypeFrom :: Name,
    ctxCtrTypeTo :: Name,
    ctxCtrCtrTo :: Name
    }
  deriving (Eq, Show)

numCCArgs :: ContextCtr -> Int
numCCArgs cc
  | ctxCtrIsList cc = 2 + val
  | otherwise = val
  where 
    val = length (ctxCtrArgsBefore cc) + length (ctxCtrArgsAfter cc)


downCtrName :: ContextCtr -> String
downCtrName cc = "M" ++ ctr ++ "To" ++ typ ++ moffset
  where
    ctr = nameBase . ctxCtrCtrTo $ cc
    typ = nameBase . ctxCtrTypeFrom $ cc
    moffset = maybe "" show . ctxCtrOffset $ cc

upCtrNameShown :: ContextCtr -> String
upCtrNameShown cc = mupctrshown
  where
    mupctrshown = "M" ++ typ ++ moffset ++ "To" ++ ctr

    ctr = nameBase . ctxCtrCtrTo $ cc
    typ = nameBase . ctxCtrTypeFrom $ cc
    moffset = maybe "" show . ctxCtrOffset $ cc

ctxCtrOffset :: ContextCtr -> Maybe Int
ctxCtrOffset (ListCC {}) = Nothing
ctxCtrOffset x = ctxCtrOffset_ x

ctxCtrIsList :: ContextCtr -> Bool
ctxCtrIsList (ListCC {}) = True
ctxCtrIsList _           = False

ctxCtrArgsBefore :: ContextCtr -> [String]
ctxCtrArgsBefore cc
  | ctxCtrIsList cc = []
  | otherwise = map nameBase . ctxCtrTypesBefore $ cc
  
ctxCtrArgsAfter :: ContextCtr -> [String]
ctxCtrArgsAfter cc
  | ctxCtrIsList cc = []
  | otherwise = map nameBase . ctxCtrTypesAfter $ cc

buildContextCtrs :: Map Name DataType -> [ContextCtr]
buildContextCtrs inMap = concatMap buildContextLines . Map.elems $ inMap
  where
    buildContextLines :: DataType -> [ContextCtr]
    buildContextLines (WrapsList dt cn ct) = [ListCC {
      ctxCtrName = (childName ++ "To" ++ conName),
      ctxCtrTypeFrom = ct,
      ctxCtrTypeTo   = dt,
      ctxCtrCtrTo    = cn
      }]
      where
        childName = nameBase ct
        conName   = nameBase cn
    buildContextLines (DataType dtName cons) = concatMap mkNormalCCs cons
      where
        mkNormalCCs :: Constructor -> [ContextCtr]
        mkNormalCCs (Ctr cn kids) = catMaybes $ mapWithContext kids maybeMkNormalCC
          where
            conName = nameBase cn

            maybeMkNormalCC :: [Child] -> Child -> [Child] -> Maybe ContextCtr
            maybeMkNormalCC _ (NonNavigable _) _ = Nothing
            maybeMkNormalCC before (Navigable n) after = Just $ NormalCC {
              ctxCtrName = (childName ++ "To" ++ conName ++ count),
              ctxCtrOffset_ = mcount,
              ctxCtrTypesBefore = map childType before,
              ctxCtrTypesAfter  = map childType after,
              ctxCtrTypeFrom    = n,
              ctxCtrTypeTo      = dtName,
              ctxCtrCtrTo       = cn
              }
              where
                childName = nameBase n
                count = maybe "" show mcount
                mcount
                  | null (sameSiblingsBefore ++ sameSiblingsAfter) = Nothing
                  | otherwise = Just . length $ sameSiblingsBefore
                  where
                    sameSiblingsBefore = filter ( (==n) . childType ) before
                    sameSiblingsAfter  = filter ( (==n) . childType ) after

mapWithContext :: [a] -> ([a] -> a -> [a] -> b) -> [b]
mapWithContext xs f = mapWithContext' xs f []
  where
    mapWithContext' [] _ _ = []
    mapWithContext' (x:xs) f r = (f r x xs):(mapWithContext' xs f (r ++ [x]))

extractConstructors :: Info -> [Con]
extractConstructors (TyConI dec) = getConstructors dec
extractConstructors _ = []

getConstructors :: Dec -> [Con]
getConstructors (DataD _ _ _ cons _)    = cons
getConstructors (NewtypeD _ _ _ con _)  = [con]
getConstructors _                       = []

getInfoName :: Info -> Name
getInfoName (TyConI (DataD _ n _ _ _))    = n
getInfoName (TyConI (NewtypeD _ n _ _ _)) = n
getInfoName x                    = error $ "Don't know how to get name for: " ++ show x

getConName :: Con -> Name
getConName (NormalC n _) = n
getConName (RecC n _)    = n
getConName x = error $ "GetConName: " ++ show x

extractChildTypesRaw :: Con -> [Type]
extractChildTypesRaw con = case con of
  (RecC _ types)    -> map thd $ types
  (NormalC _ types) -> map snd $ types
  _ -> []

extractTypeNoList :: Type -> Name
extractTypeNoList (ConT name) = name
extractTypeNoList (AppT (ConT lst) rhs)
  | lst == ''[] = (extractTypeNoList rhs)
extractTypeNoList x = error $ show x

isList :: Type -> Bool
isList (AppT (ConT lst) _) = lst == ''[]
isList _ = False

thd :: (a,b,c) -> c
thd (_,_,c) = c

extractType :: Type -> Name
extractType (ConT name) = name
extractType (AppT (ConT lst) rhs)
  | lst == ''[] = mkName $ "[" ++ nameBase (extractType rhs) ++ "]"
extractType x = error $ "Extract type; " ++ show x

extractChildTypes :: Con -> [Name]
extractChildTypes = catMaybes . map ect . extractChildTypesRaw 
  where
    ect :: Type -> Maybe Name
    ect (ConT name) = Just name
    ect (AppT (ConT lst) rhs)
      | lst == ''[] = ect rhs
    ect _ = Nothing