import Data.Generics.Uniplate.Data  
import Data.Generics.Biplate hiding (transformBi, childrenBi)
import Language.Haskell.Exts.Annotated 
import Language.Haskell.Exts.SrcLoc
import System.Environment


main :: IO ()
main = do
     [fileName, inputPath, outputPath] <- getArgs
     res <- parseFileWithComments (defaultParseMode { parseFilename = fileName }) inputPath
     case res of 
          ParseFailed _ _ -> readFile inputPath >>= return . (header ++) >>= writeFile  outputPath
            where
              header = "{-# LINE 1 \"" ++ fileName ++ "\" #-}\n"
          ParseOk (mod, cs) -> do
                                   let (mod', cs') = case trans fileName mod of
                                                      (Nothing, _) -> (mod, cs) 
                                                      (Just l, m)  -> (m, linePragma fileName (l+1) l : move l cs)
                                   writeFile outputPath $ exactPrint mod' cs'
  


-- move :: Int -> Module SrcSpanInfo -> Module SrcSpanInfo
move boundary = transformBi move'
  where
    move' :: SrcSpan -> SrcSpan
    -- move' = id
    move' s@(SrcSpan n l c l' c')
        | boundary <= l           = SrcSpan n (l+2) c (l'+2) c' 
        | otherwise               = s

trans fileName m@(Module l h ps is ds)
    | is == [] && ds == []    = (Nothing, m)
    | otherwise               = (Just insertLine, Module l h ps is' (map (move insertLine) ds))
  where
    insertLine = if length is > 0 then lineOf (head is)
                                  else lineOf (head ds)

    is' = new : map (move insertLine) is 
    new = ImportDecl imp (ModuleName mod "Prelude") False False Nothing Nothing Nothing

    imp = infoSpan start [start]
    start = mkSrcSpan (SrcLoc fileName insertLine 1) (SrcLoc fileName insertLine 7)

    mod = infoSpan mod' []
    mod' = mkSrcSpan (SrcLoc fileName insertLine 8) (SrcLoc fileName insertLine 15)

lineOf :: Annotated ast => ast SrcSpanInfo -> Int
lineOf i = l
  where
    (l, _) = srcSpanStart (srcInfoSpan (ann i))

linePragma :: String -> Int -> Int -> Comment
linePragma fileName realLinePos virtualLinePos = Comment True pos commentString
  where
    commentString = "# LINE " ++ show virtualLinePos ++ " \"" ++ fileName ++ "\""++  " #"
    pos = mkSrcSpan (SrcLoc fileName realLinePos 1) (SrcLoc fileName realLinePos (length commentString + 4))