{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE UndecidableInstances #-}

{-# OPTIONS -Wno-orphans#-}

module PrettyParseTests  where

import           Data.Algorithm.Diff
import           Data.Algorithm.DiffOutput
import           Data.Char
import           Data.Fix
import qualified Data.List.NonEmpty as NE
import           Data.Text (Text, pack)
import           Hedgehog
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import           Nix.Atoms
import           Nix.Expr
import           Nix.Parser
import           Nix.Pretty
import           Test.Tasty
import           Test.Tasty.Hedgehog
import           Text.Megaparsec (Pos, SourcePos, mkPos)
import           Text.PrettyPrint.ANSI.Leijen ((</>), text)
import qualified Text.PrettyPrint.ANSI.Leijen as P
import qualified Text.Show.Pretty as PS

asciiString :: MonadGen m => m String
asciiString = Gen.list (Range.linear 1 15) Gen.lower

asciiText :: Gen Text
asciiText = pack <$> asciiString

-- Might want to replace this instance with a constant value
genPos :: Gen Pos
genPos = mkPos <$> Gen.int (Range.linear 1 256)

genSourcePos :: Gen SourcePos
genSourcePos = SourcePos <$> asciiString <*> genPos <*> genPos

genKeyName :: Gen (NKeyName NExpr)
genKeyName = Gen.choice [ DynamicKey <$> genAntiquoted genString
                        , StaticKey <$> asciiText ]

genAntiquoted :: Gen a -> Gen (Antiquoted a NExpr)
genAntiquoted gen = Gen.choice
  [ Plain <$> gen
  , pure EscapedNewline
  , Antiquoted <$> genExpr
  ]

genBinding :: Gen (Binding NExpr)
genBinding = Gen.choice
  [ NamedVar <$> genAttrPath <*> genExpr <*> genSourcePos
  , Inherit <$> Gen.maybe genExpr
            <*> Gen.list (Range.linear 0 5) genKeyName
            <*> genSourcePos
  ]

genString :: Gen (NString NExpr)
genString = Gen.choice
  [ DoubleQuoted <$> Gen.list (Range.linear 0 5) (genAntiquoted asciiText)
  , Indented <$> Gen.int (Range.linear 0 10)
             <*> Gen.list (Range.linear 0 5) (genAntiquoted asciiText)
  ]

genAttrPath :: Gen (NAttrPath NExpr)
genAttrPath = (NE.:|) <$> genKeyName
                      <*> Gen.list (Range.linear 0 4) genKeyName

genParams :: Gen (Params NExpr)
genParams = Gen.choice
  [ Param    <$> asciiText
  , ParamSet <$> Gen.list (Range.linear 0 10) ((,) <$> asciiText
                                                   <*> Gen.maybe genExpr)
             <*> Gen.bool
             <*> Gen.choice [pure Nothing, Just <$> asciiText]
  ]

genAtom :: Gen NAtom
genAtom = Gen.choice
  [ NInt   <$> Gen.integral (Range.linear 0 1000)
  , NFloat <$> Gen.float (Range.linearFrac 0.0 1000.0)
  , NBool  <$> Gen.bool
  , pure NNull ]

-- This is written by hand so we can use `fairList` rather than the normal
-- list Arbitrary instance which makes the generator terminate. The
-- distribution is not scientifically chosen.
genExpr :: Gen NExpr
genExpr = Gen.sized $ \(Size n) ->
  Fix <$>
      if n < 2
      then Gen.choice
        [genConstant, genStr, genSym, genLiteralPath, genEnvPath ]
      else
        Gen.frequency
          [ ( 1, genConstant)
          , ( 1, genSym)
          , ( 4, Gen.resize (Size (n `div` 3)) genIf)
          , (10, genRecSet )
          , (20, genSet )
          , ( 5, genList )
          , ( 2, genUnary )
          , ( 2, Gen.resize (Size (n `div` 3)) genBinary )
          , ( 3, Gen.resize (Size (n `div` 3)) genSelect )
          , (20, Gen.resize (Size (n `div` 2)) genAbs )
          , ( 2, Gen.resize (Size (n `div` 2)) genHasAttr )
          , (10, Gen.resize (Size (n `div` 2)) genLet )
          , (10, Gen.resize (Size (n `div` 2)) genWith )
          , ( 1, Gen.resize (Size (n `div` 2)) genAssert)
          ]
 where
  genConstant    = NConstant    <$> genAtom
  genStr         = NStr         <$> genString
  genSym         = NSym         <$> asciiText
  genList        = NList        <$> fairList genExpr
  genSet         = NSet         <$> fairList genBinding
  genRecSet      = NRecSet      <$> fairList genBinding
  genLiteralPath = NLiteralPath . ("./" ++) <$> asciiString
  genEnvPath     = NEnvPath     <$> asciiString
  genUnary       = NUnary       <$> Gen.enumBounded <*> genExpr
  genBinary      = NBinary      <$> Gen.enumBounded <*> genExpr <*> genExpr
  genSelect      = NSelect      <$> genExpr <*> genAttrPath <*> Gen.maybe genExpr
  genHasAttr     = NHasAttr     <$> genExpr <*> genAttrPath
  genAbs         = NAbs         <$> genParams <*> genExpr
  genLet         = NLet         <$> fairList genBinding <*> genExpr
  genIf          = NIf          <$> genExpr <*> genExpr <*> genExpr
  genWith        = NWith        <$> genExpr <*> genExpr
  genAssert      = NAssert      <$> genExpr <*> genExpr

-- | Useful when there are recursive positions at each element of the list as
--   it divides the size by the length of the generated list.
fairList :: Gen a -> Gen [a]
fairList g = Gen.sized $ \s -> do
  k <- Gen.int (Range.linear 0 (unSize s))
  -- Use max here to avoid dividing by zero when there is the empty list
  Gen.resize (Size (unSize s `div` max 1 k)) $ Gen.list (Range.singleton k) g

equivUpToNormalization :: NExpr -> NExpr -> Bool
equivUpToNormalization x y = normalize x == normalize y

normalize :: NExpr -> NExpr
normalize = cata $ \case
  NConstant (NInt n)   | n < 0 -> Fix (NUnary NNeg (Fix (NConstant (NInt (negate n)))))
  NConstant (NFloat n) | n < 0 -> Fix (NUnary NNeg (Fix (NConstant (NFloat (negate n)))))

  NSet binds      -> Fix (NSet (map normBinding binds))
  NRecSet binds   -> Fix (NRecSet (map normBinding binds))
  NLet binds r    -> Fix (NLet (map normBinding binds) r)

  NAbs params r   -> Fix (NAbs (normParams params) r)

  r               -> Fix r

 where
  normBinding (NamedVar path r pos) = NamedVar (NE.map normKey path) r pos
  normBinding (Inherit mr names pos) = Inherit mr (map normKey names) pos

  normKey (DynamicKey quoted) = DynamicKey (normAntiquotedString quoted)
  normKey (StaticKey name) = StaticKey name

  normAntiquotedString :: Antiquoted (NString NExpr) NExpr
                       -> Antiquoted (NString NExpr) NExpr
  normAntiquotedString (Plain (DoubleQuoted [EscapedNewline])) =
      EscapedNewline
  normAntiquotedString (Plain (DoubleQuoted strs)) =
      let strs' = map normAntiquotedText strs
      in if strs == strs'
         then Plain (DoubleQuoted strs)
         else normAntiquotedString (Plain (DoubleQuoted strs'))
  normAntiquotedString r = r

  normAntiquotedText :: Antiquoted Text NExpr -> Antiquoted Text NExpr
  normAntiquotedText (Plain "\n")   = EscapedNewline
  normAntiquotedText (Plain "''\n") = EscapedNewline
  normAntiquotedText r = r

  normParams (ParamSet binds var (Just "")) = ParamSet binds var Nothing
  normParams r = r

-- | Test that parse . pretty == id up to attribute position information.
prop_prettyparse :: Monad m => NExpr -> PropertyT m ()
prop_prettyparse p = do
  let prog = show (pretty p)
  case parse (pack prog) of
    Failure s -> do
        footnote $ show $
            text "Parse failed:" </> text (show s)
              P.<$> P.indent 2 (pretty p)
        discard
    Success v
        | equivUpToNormalization p v -> success
        | otherwise -> do
          let pp = normalise prog
              pv = normalise (show (pretty v))
          footnote $ show $
                  text "----------------------------------------"
            P.<$> text "Expr before:" P.<$> P.indent 2 (text (PS.ppShow p))
            P.<$> text "----------------------------------------"
            P.<$> text "Expr after:"  P.<$> P.indent 2 (text (PS.ppShow v))
            P.<$> text "----------------------------------------"
            P.<$> text "Pretty before:" P.<$> P.indent 2 (text prog)
            P.<$> text "----------------------------------------"
            P.<$> text "Pretty after:"  P.<$> P.indent 2 (pretty v)
            P.<$> text "----------------------------------------"
            P.<$> text "Normalised before:" P.<$> P.indent 2 (text pp)
            P.<$> text "----------------------------------------"
            P.<$> text "Normalised after:"  P.<$> P.indent 2 (text pv)
            P.<$> text "========================================"
            P.<$> text "Normalised diff:"
            P.<$> text (ppDiff (diff pp pv))
            P.<$> text "========================================"
          assert (pp == pv)
  where
    pretty = prettyNix
    parse  = parseNixText

    normalise = unlines . map (reverse . dropWhile isSpace . reverse) . lines

    diff :: String -> String -> [Diff [String]]
    diff s1 s2 = getDiff (map (:[]) (lines s1)) (map (:[]) (lines s2))

tests :: TestLimit -> TestTree
tests n = testProperty "Pretty/Parse Property" $ withTests n $ property $ do
    x <- forAll genExpr
    prop_prettyparse x