{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
module SJW.Module.Imports (
      Reference(..)
    , Tree(..)
    , parser
    , recurse
  ) where

import SJW.Source (Path(..))
import Control.Applicative ((<|>), many, optional)
import Data.Attoparsec.Text (
    Parser, char, count, digit, inClass, letter, sepBy, string, takeWhile
  )
import Data.Map (Map, foldlWithKey)
import qualified Data.Map as Map (empty, insert, lookup)
import qualified Data.Text as Text (pack)
import Prelude hiding (takeWhile)

data Reference =
    ModulePath {modulePath :: Path}
  | Object {modulePath :: Path, field :: String}
  deriving Show

data Tree = Tree {
      target :: Maybe Reference
    , children :: Map String Tree
  } deriving Show

data Mapping = Mapping {
      exposedName :: Path
    , reference :: Reference
  }

recurse :: (a -> [String] -> Reference -> a) -> a -> Tree -> a
recurse f initValue = recAux [] initValue
  where
    next _ value Nothing = value
    next stack value (Just ref) = f value (reverse stack) ref
    recAux stack value tree =
      let nextValue = next stack value (target tree) in
      foldlWithKey (\a k b -> recAux (k:stack) a b) nextValue (children tree)

space :: Parser ()
space = takeWhile (inClass " \t\r\n") *> pure ()

between :: Parser a -> (Parser b, Parser c) -> Parser a
between p (left, right) = left *> space *> p <* space <* right

keyword :: String -> Parser ()
keyword k = space <* string (Text.pack k) <* space

name :: Parser String
name = (:) <$> letter <*> many (letter <|> digit)

aliasedName :: Parser (Maybe String, String)
aliasedName =
    ((,) <$> (Just <$> name) <* keyword "as" <*> name)
  <|> ((\s -> (Just s, s)) <$> name)

buildMappings :: Maybe [(Maybe String, String)] -> Path -> [Mapping]
buildMappings Nothing modulePath =
  [Mapping modulePath (ModulePath modulePath)]
buildMappings (Just nameAssocs) modulePath = mappingOf <$> nameAssocs
  where
    mappingOf (Nothing, dest) = Mapping (Path [dest]) (ModulePath modulePath)
    mappingOf (Just source, dest) =
      Mapping (Path [dest]) (Object modulePath source)

mappingParser :: Parser [Mapping]
mappingParser =
  buildMappings <$> optional fromClause <*> (Path <$> name `sepBy` char '.')
  where
    fromClause =
      (count 1 (aliasedName <|> star) <|> namesBlock) <* keyword "from"
    namesBlock =
      (aliasedName `sepBy` (char ',' *> space)) `between` (char '{', char '}')
    star = (,) <$> (char '*' *> pure Nothing) <* keyword "as" <*> name

emptyTree :: Tree
emptyTree = Tree {
      target = Nothing
    , children = Map.empty
  }

insertMapping :: Tree -> Mapping -> Tree
insertMapping tmpTree (Mapping {exposedName, reference}) =
  insertAt components tmpTree
  where
    Path components = exposedName
    insertAt [] tree = tree {target = Just reference}
    insertAt (next:restOfPath) tree@(Tree {children}) =
      let subTree = maybe emptyTree id $ Map.lookup next children in tree {
        children = Map.insert next (insertAt restOfPath subTree) children
      }

parser :: Parser Tree
parser = foldl (foldl insertMapping) emptyTree <$> importParser `sepBy` blank
  where
    blank = takeWhile (inClass " \t\r\n")
    importParser = mappingParser `between` (string "import", char ';')