module Test.DocTest.Parse (
  DocTest(..),
  Expression,
  Interaction,
  parseComment,
  ) where

import Test.DocTest.Location (Located(Located), unLoc)
import Test.DocTest.Base

import Data.List (stripPrefix, isPrefixOf)
import Data.Maybe (fromMaybe)
import Data.Char (isSpace)



data DocTest = Example Expression ExpectedResult | Property Expression
   deriving (Eq, Show)

type Expression = String

type Interaction = (Expression, ExpectedResult)

parseComment :: [Located pos String] -> [Located pos DocTest]
parseComment nLines = properties ++ examples
  where
    examples   = map (fmap $ uncurry Example) (parseInteractions nLines)
    properties = map (fmap          Property) (parseProperties   nLines)

-- | Extract all properties from given Haddock comment.
parseProperties :: [Located pos String] -> [Located pos Expression]
parseProperties = go
  where
    isPrompt :: Located pos String -> Bool
    isPrompt = isPrefixOf "prop>" . dropWhile isSpace . unLoc

    go xs = case dropWhile (not . isPrompt) xs of
      prop:rest -> stripPrompt `fmap` prop : go rest
      [] -> []

    stripPrompt = strip . drop 5 . dropWhile isSpace

-- | Extract all interactions from given Haddock comment.
parseInteractions :: [Located pos String] -> [Located pos Interaction]
parseInteractions = go
  where
    isPrompt :: Located pos String -> Bool
    isPrompt = isPrefixOf ">>>" . dropWhile isSpace . unLoc

    isBlankLine :: Located pos String -> Bool
    isBlankLine  = null . dropWhile isSpace . unLoc

    isEndOfInteraction :: Located pos String -> Bool
    isEndOfInteraction x = isPrompt x || isBlankLine x


    go :: [Located pos String] -> [Located pos Interaction]
    go xs = case dropWhile (not . isPrompt) xs of
      prompt:rest ->
        case (words (drop 3 (dropWhile isSpace (unLoc prompt))),
              break isBlankLine rest) of
          (":{" : _, (ys,zs)) -> toInteraction prompt ys : go zs
          _ ->
            let (ys,zs) = break isEndOfInteraction rest
            in toInteraction prompt ys : go zs
      [] -> []

-- | Create an `Interaction`, strip superfluous whitespace as appropriate.
--
-- also merge lines between :{ and :}, preserving whitespace inside
-- the block (since this is useful for avoiding {;}).
toInteraction :: Located pos String -> [Located pos String] -> Located pos Interaction
toInteraction (Located loc x) xs = Located loc $
  (
    (strip   cleanedE)  -- we do not care about leading and trailing
                        -- whitespace in expressions, so drop them
  , map mkExpectedLine result_
  )
  where
    -- 1. drop trailing whitespace from the prompt, remember the prefix
    (prefix, e) = span isSpace x
    (ePrompt, eRest) = splitAt 3 e

    -- 2. drop, if possible, the exact same sequence of whitespace
    -- characters from each result line
    unindent pre = map (tryStripPrefix pre . unLoc)

    cleanBody line = fromMaybe (unLoc line)
                    (stripPrefix ePrompt (dropWhile isSpace (unLoc line)))

    (cleanedE, result_) =
        case break ( (==) [":}"] . take 1 . words . cleanBody) xs of
            (body , endLine : rest) ->
                (unlines (eRest : map cleanBody body ++
                                [dropWhile isSpace (cleanBody endLine)]),
                        unindent (takeWhile isSpace (unLoc endLine)) rest)
            _ -> (eRest, unindent prefix xs)


tryStripPrefix :: String -> String -> String
tryStripPrefix prefix ys = fromMaybe ys $ stripPrefix prefix ys

mkExpectedLine :: String -> ExpectedLine
mkExpectedLine x = case x of
    "<BLANKLINE>" -> ExpectedLine [LineChunk ""]
    "..." -> WildCardLine
    _ -> ExpectedLine $ mkLineChunks x

mkLineChunks :: String -> [LineChunk]
mkLineChunks = finish . foldr go (0, [], [])
  where
    mkChunk :: String -> [LineChunk]
    mkChunk "" = []
    mkChunk x  = [LineChunk x]

    go :: Char -> (Int, String, [LineChunk]) -> (Int, String, [LineChunk])
    go '.' (count, acc, res) = if count == 2
          then (0, "", WildCardChunk : mkChunk acc ++ res)
          else (count + 1, acc, res)
    go c   (count, acc, res) = if count > 0
          then (0, c : replicate count '.' ++ acc, res)
          else (0, c : acc, res)
    finish (count, acc, res) = mkChunk (replicate count '.' ++ acc) ++ res


-- | Remove leading and trailing whitespace.
strip :: String -> String
strip = dropWhile isSpace . reverse . dropWhile isSpace . reverse