{-# LANGUAGE FlexibleInstances
           , FlexibleContexts
           , UndecidableInstances
           , NamedFieldPuns 
           #-}

-- | Pretty printing the AST
module Language.Haskell.Tools.PrettyPrint (prettyPrint) where

import FastString (fsLit)
import SrcLoc

import Language.Haskell.Tools.AST (SrcTemplateStage, SourceInfoTraversal(..))
import Language.Haskell.Tools.PrettyPrint.RoseTree (RoseSourceInfo(..), RoseTree(..), toRoseTree)
import Language.Haskell.Tools.Transform.SourceTemplate

import Control.Monad.State
import Data.Foldable (Foldable(..), concat)
import Data.List as List
import Data.List.Split (splitOn)
import Data.Sequence hiding (null, replicate)

-- | Pretty prints an AST by using source templates stored as node info
prettyPrint :: (SourceInfoTraversal node) => node dom SrcTemplateStage -> String
prettyPrint = toList . printRose . toRoseTree

printRose :: RoseTree SrcTemplateStage -> Seq Char      
printRose rt = evalState (printRose' startLoc rt) startLoc
  where startLoc = mkRealSrcLoc (fsLit "") 1 1

type PPState = State RealSrcLoc

-- | Pretty prints a rose tree according to the source templates remaining from the original AST
printRose' :: RealSrcLoc -> RoseTree SrcTemplateStage -> PPState (Seq Char)
-- simple implementation could be optimized a bit
-- warning: the length of the file should not exceed maxbound::Int
printRose' parent (RoseTree (RoseSpan (SourceTemplateNode rng elems minInd relInd)) children) 
  = do slide <- calculateSlide rng
       let printTemplateElems :: [SourceTemplateElem] -> [RoseTree SrcTemplateStage] -> PPState (Seq Char)
           printTemplateElems (TextElem txt : rest) children = putString slide min txt >+< printTemplateElems rest children
           printTemplateElems (ChildElem : rest) (child : children) = printRose' parent child >+< printTemplateElems rest children
           printTemplateElems [] [] = return empty
           printTemplateElems _ [] = error $ "More child elem in template than actual children (elems: " ++ show elems ++ ", children: " ++ show children ++ ")"
           printTemplateElems [] _ = error $ "Not all children are used to pretty printing. (elems: " ++ show elems ++ ", children: " ++ show children ++ ")"

           min = minInd `max` getPosByRelative parent relInd

       printTemplateElems elems children
  
printRose' _ (RoseTree (RoseList (SourceTemplateList {})) []) = return empty
printRose' parent (RoseTree (RoseList (SourceTemplateList rng bef aft defSep indented seps minInd relInd)) children) 
    = do slide <- calculateSlide rng
         actRng <- get
         let min = minInd `max` getPosByRelative parent relInd
         putString slide min bef 
           >+< (if indented then printListWithSepsIndented else printListWithSeps) actRng slide min actualSeps children 
           >+< putString slide min aft
  where actualSeps = case seps of [] -> repeat defSep
                                  _  -> seps ++ repeat (last seps)

printRose' _ (RoseTree (RoseOptional (SourceTemplateOpt {})) []) = return empty
printRose' parent (RoseTree (RoseOptional (SourceTemplateOpt rng bef aft minInd relInd)) [child]) 
  = do slide <- calculateSlide rng
       actRng <- get
       let min = minInd `max` getPosByRelative parent relInd
       putString slide min bef >+< printRose' actRng child >+< putString slide min aft
printRose' _ (RoseTree (RoseOptional _) _) = error "More than one child element in an optional node."
    
getPosByRelative :: RealSrcLoc -> Maybe Int -> Int
getPosByRelative sp (Just i) = srcLocCol sp + i - 1
getPosByRelative _ _ = 0

calculateSlide :: SrcSpan -> PPState Int
calculateSlide (RealSrcSpan originalSpan) = do 
  actualSpan <- get
  return $ srcLocCol actualSpan - srcLocCol (realSrcSpanStart originalSpan)
calculateSlide _ = return 0

putString :: Int -> Int -> String -> PPState (Seq Char)
putString slide minInd s 
  = do modify $ advanceStr newStr
       return (fromList newStr)
  where start:rest = splitOn "\n" s
        newStr = concat $ intersperse ("\n" ++ replicate slide ' ') (start : map (extendToNSpaces minInd) rest)
        extendToNSpaces n str = replicate n ' ' ++ (List.dropWhile (== ' ') $ List.take n str) ++ List.drop n str
                  
advanceStr :: String -> RealSrcLoc -> RealSrcLoc
advanceStr s loc = foldl advanceSrcLoc loc s

untilReaches :: String -> RealSrcLoc -> RealSrcLoc -> (String, Int)
untilReaches s start end 
  = let ls = splitOn "\n" s 
     in case ls of _:_:_ -> (unlines (init ls) ++) 
                              `mapFst` untilReaches' (last ls) (advanceSrcLoc start '\n') end 
                   _ -> (s, srcLocCol start)
  where
    untilReaches' [] curr _ = ([], srcLocCol curr)
    untilReaches' (c:rest) curr until | srcLocCol advancedLoc <= srcLocCol until
      = (c:) `mapFst` untilReaches' rest advancedLoc until
      where advancedLoc = advanceSrcLoc curr c
    untilReaches' _ curr _ = ([], srcLocCol curr)
    
mapFst :: (a -> b) -> (a, x) -> (b, x)
mapFst f (a, x) = (f a, x)

(>+<) :: PPState (Seq Char) -> PPState (Seq Char) -> PPState (Seq Char)
(>+<) = liftM2 (><)
    
printListWithSeps :: RealSrcLoc -> Int -> Int -> [String] -> [RoseTree SrcTemplateStage] -> PPState (Seq Char)
printListWithSeps = printListWithSeps' putString

-- | Prints the elements of a list where the elements must be printed in the same line (do stmts, case alts, let binds, ...)
printListWithSepsIndented :: RealSrcLoc -> Int -> Int -> [String] -> [RoseTree SrcTemplateStage] -> PPState (Seq Char)
printListWithSepsIndented parent slide minInd seps children
  = do base <- get
       let putCorrectSep _ min s = do curr <- get
                                      let (shortened, currCol) = untilReaches s curr base
                                      putString 0 min $ shortened ++ replicate (srcLocCol base - currCol) ' '
       printListWithSeps' putCorrectSep parent slide minInd seps children
    
printListWithSeps' :: (Int -> Int -> String -> PPState (Seq Char)) -> RealSrcLoc -> Int -> Int -> [String] -> [RoseTree SrcTemplateStage] -> PPState (Seq Char)
printListWithSeps' _ _ _ _ _ [] = return empty
printListWithSeps' _ parent _ _ _ [child] = printRose' parent child
printListWithSeps' putCorrectSep parent slide minInd (sep:seps) (child:children) 
  = printRose' parent child >+< putCorrectSep slide minInd sep >+< printListWithSeps' putCorrectSep parent slide minInd seps children
printListWithSeps' _ _ _ _ [] _ = error "printListWithSeps': the number of elements and separators does not match"