{-# LANGUAGE OverloadedStrings #-}
module Main(main) where

import GHC.Driver.Session (getDynFlags)
import GHC.Driver.Monad (liftIO)
import GHC.Paths (libdir)
import GHC (runGhc, DynFlags)
import GHC.Utils.Outputable (Outputable)

import Test.Tasty
import Test.Tasty.HUnit

import GHC.SourceGen
import GhcVersion

data TestCase a = String :~ a

infixr 0 :~

testCases :: Outputable a => DynFlags -> String -> [TestCase a] -> TestTree
testCases dflags name cases = testGroup name $ map run cases
  where
    run (expected :~ x) =
        testCase (takeWhile (/='\n') expected) $ expected @=? showPpr dflags x

testTypes :: DynFlags ->  String -> [TestCase HsType'] -> TestTree
testTypes = testCases

testExprs :: DynFlags ->  String -> [TestCase HsExpr'] -> TestTree
testExprs = testCases

testDecls :: DynFlags ->  String -> [TestCase HsDecl'] -> TestTree
testDecls = testCases

testPats :: DynFlags ->  String -> [TestCase Pat'] -> TestTree
testPats = testCases

testModule :: DynFlags -> String -> [TestCase HsModule'] -> TestTree
testModule = testCases

main :: IO ()
main = runGhc (Just libdir) $ do
    dflags <- getDynFlags
    liftIO $ defaultMain $ testGroup "Tests"
        [ typesTest dflags
        , exprsTest dflags
        , declsTest dflags
        , patsTest dflags
        , modulesTest dflags
        ]

typesTest, exprsTest, declsTest, patsTest, modulesTest :: DynFlags -> TestTree
typesTest dflags = testGroup "Type"
    [ test "var"
        [ "A" :~ var "A"
        , "x" :~ var "x"
        , "A.x" :~ var "A.x"
        , "x" :~ var (unqual "x")
        , "A.x" :~ var (qual "A" "x")
        ]
    , test "app"
        [ "A x" :~ var "A" @@ var "x"
        , "(+) x" :~ var "+" @@ var "x"
        , "A (B x)" :~ var "A" @@ par (var "B" @@ var "x")
        , "A (B x)" :~ var "A" @@ par (var "B" @@ var "x")
        , "A ((B x))" :~ var "A" @@ par (par (var "B" @@ var "x"))
        , "A x (B y z)" :~ var "A" @@ var "x" @@ (var "B" @@ var "y" @@ var "z")
        , "A w (B x y) Z"
            :~ var "A" @@ var "w" @@ (var "B" @@ var "x" @@ var "y") @@ var "Z"
        ]
    , test "op"
        [ "x + y" :~ op (var "x") "+" (var "y")
        , "x `add` y" :~ op (var "x") "add" (var "y")
        , "x * (y + z)" :~ op (var "x") "*" (op (var "y") "+" (var "z"))
        , "(x * y) + z" :~ op (op (var "x") "*" (var "y")) "+" (var "z")
        , "x `mult` (y `add` z)" :~ op (var "x") "mult" (op (var "y") "add" (var "z"))
        , "A x * (B y + C z)" :~ op (var "A" @@ var "x") "*"
                                    (op (var "B" @@ var "y") "+" (var "C" @@ var "z"))
        , "(f . g) x" :~ op (var "f") "." (var "g") @@ var "x"
        ]
    , test "function"
        [ "a -> b" :~ var "a" --> var "b"
        , "a -> b -> c" :~ var "a" --> var "b" --> var "c"
        , "a -> b -> c" :~ var "a" --> (var "b" --> var "c")
        , "(a -> b) -> c" :~ (var "a" --> var "b") --> var "c"
        -- These tests also check that ==> and --> have compatible precedences:
        , "A a => a -> b" :~ [var "A" @@ var "a"] ==> var "a" --> var "b"
        , "(A a, B b) => a -> b" :~
            [var "A" @@ var "a", var "B" @@ var "b"] ==> var "a" --> var "b"
        -- It appears to be correct to *not* wrap `A => c` in parentheses;
        -- GHC still parses it as a function between two HsQualTy.
        , "(A => b) -> A => c" :~
            ([var "A"] ==> var "b") --> ([var "A"] ==> var "c")
        , "(A => b) -> A => c" :~
            ([var "A"] ==> var "b") --> [var "A"] ==> var "c"
        ]
    , test "literals"
        [ "\"abc\"" :~ stringTy "abc"
        , "123" :~ numTy 123
        ]
    , test "unit"
        [ "()" :~ unit
        , "(# #)" :~ unboxedTuple []
        ]
   , test "list"
        [ "[x]" :~ listTy (var "x")
        , "'[]" :~ listPromotedTy []
        , "'[x]" :~ listPromotedTy [var "x"]
        , "'[y, z]" :~ listPromotedTy [var "y", var "z"]
        ]
    , test "tuple"
        [ "(a, b)" :~ tuple [(var "a"), (var "b")]
        , "(# a, b #)" :~ unboxedTuple [(var "a"), (var "b")]
        , "'(a, b)" :~ tuplePromotedTy [(var "a"), (var "b")]
        ]
    , test "tyPromotedVar"
        -- For some reason, older GHC pretty-printed an extra space.
        [ ifGhc88 "'Abc" " 'Abc" :~ tyPromotedVar "Abc"
        , ifGhc88 "T 'Abc" "T  'Abc" :~ var "T" @@ tyPromotedVar "Abc"
        ]
    ]
  where
    test = testTypes dflags

exprsTest dflags = testGroup "Expr"
    [ test "var"
        [ "A" :~ var "A"
        , "x" :~ var "x"
        , "A.x" :~ var "A.x"
        , "x" :~ var (unqual "x")
        , "A.x" :~ var (qual "A" "x")
        ]
    , test "app"
        [ "A x" :~ var "A" @@ var "x"
        , "(+) x" :~ var "+" @@ var "x"
        , "(Prelude.+) x" :~ var "Prelude.+" @@ var "x"
        , "A (B x)" :~ var "A" @@ (var "B" @@ var "x")
        , "A (B x)" :~ var "A" @@ par (var "B" @@ var "x")
        , "A ((B x))" :~ var "A" @@ par (par (var "B" @@ var "x"))
        , "A x (B y z)" :~ var "A" @@ var "x" @@ (var "B" @@ var "y" @@ var "z")
        , "A w (B x y) Z"
            :~ var "A" @@ var "w" @@ (var "B" @@ var "x" @@ var "y") @@ var "Z"
        , "A 3" :~ var "A" @@ int 3
        , "A (-3)" :~ var "A" @@ int (-3)
        , "A 3.0" :~ var "A" @@ frac 3.0
        , "A (-3.0)" :~ var "A" @@ frac (-3.0)
        , "A 'x'" :~ var "A" @@ char 'x'
        , "A \"xyz\"" :~ var "A" @@ string "xyz"
        , "(\\ x -> x) (\\ x -> x)" :~
            let f = lambda [bvar "x"] (var "x")
            in f @@ f
        , "f x @t" :~ tyApp (var "f" @@ var "x") (var "t")
        , "f (x @t)" :~ var "f" @@ (tyApp (var "x") (var "t"))
        ]
    , test "op"
        [ "x + y" :~ op (var "x") "+" (var "y")
        , "x Prelude.+ y" :~ op (var "x") "Prelude.+" (var "y")
        , "x `add` y" :~ op (var "x") "add" (var "y")
        , "x * (y + z)" :~ op (var "x") "*" (op (var "y") "+" (var "z"))
        , "(x * y) + z" :~ op (op (var "x") "*" (var "y")) "+" (var "z")
        , "x `mult` (y `add` z)" :~ op (var "x") "mult" (op (var "y") "add" (var "z"))
        , "A x * (B y + C z)" :~ op (var "A" @@ var "x") "*"
                                    (op (var "B" @@ var "y") "+" (var "C" @@ var "z"))
        , "(f . g) x" :~ op (var "f") "." (var "g") @@ var "x"
        , "(\\ x -> x) . (\\ x -> x)" :~
            let f = lambda [bvar "x"] (var "x")
            in op f "." f
        , "x @s + y @t" :~
                op (var "x" `tyApp` var "s") "+" (var "y" `tyApp` var "t")
        ]
    , test "period-op"
        [ "(Prelude..) x" :~ var "Prelude.." @@ var "x"
        , "x Prelude.. y" :~ op (var "x") "Prelude.." (var "y")
        ]
    , test ":@@:"
        -- TODO: GHC puts extra space here.
        [ "  e :: t" :~ var "e" @::@ var "t" ]
    , test "unit"
        [ "()" :~ unit ]
    , test "list"
        [ "[]" :~ list []
        , "[]" :~ nil
        , "[x]" :~ list [var "x"]
        , "[y, z]" :~ list [var "y", var "z"]
        , "(:)" :~ cons
        , "(:) x y" :~ cons @@ var "x" @@ var "y"
        ]
    , test "tyApp"
        [ "x @a" :~ tyApp (var "x") (var "a")
        , "x @a @b" :~ tyApp (tyApp (var "x") (var "a")) (var "b")
        , "x @a b" :~ tyApp (var "x") (var "a") @@ var "b"
        , "x @(a b)" :~ tyApp (var "x") (var "a" @@ var "b")
        , "x @(a + b)" :~ tyApp (var "x") (op (var "a") "+" (var "b"))
        , "f x @t" :~ (var "f" @@ var "x") `tyApp` var "t"
        , "f (x @t)" :~ var "f" @@ (var "x" `tyApp` var "t")
        ]
    , test "recordConE"
        [ "A {}" :~ recordConE "A" []
        , "A {x = 1, y = 2}" :~ recordConE "A" [("x", int 1), ("y", int 2)]
        ]
    , test "recordUpd"
        [ "r {b = x, c = y}"
            :~ recordUpd (var "r") [("b", var "x"), ("c", var "y")]
        , "(f x) {b = x}"
            :~ recordUpd (var "f" @@ var "x") [("b", var "x")]
        , "f x {b = x}"
            :~ var "f" @@ recordUpd (var "x") [("b", var "x")]
        , "(x + y) {b = x}"
            :~ recordUpd (op (var "x") "+" (var "y")) [("b", var "x")]
        , "x + y {b = x}"
            :~ op (var "x") "+" (recordUpd (var "y") [("b", var "x")])
        ]
    , test "let"
        [ "let x = 1 in x" :~ let' [valBind "x" $ int 1] (var "x")
        , "let f x = 1 in f" :~
            let' [ funBind "f" $ match [bvar "x"] $ int 1] (var "f")
        , "let f (A x) = 1 in f" :~
            let' [ funBind "f" $ match [conP "A" [bvar "x"]] $ int 1] (var "f")
        ]
    , test "do"
        -- TODO: add more tests.
        [ "do (let x = 1 in x)" :~ do' [stmt $ let' [valBind "x" (int 1)] (var "x")]
        ]
    , test "arithSeq"
        [ "[a .. ]" :~ from (var "a")
        , "[a, b .. ]" :~ fromThen (var "a") (var "b")
        , "[a .. b]" :~ fromTo (var "a") (var "b")
        , "[a, b .. c]" :~ fromThenTo (var "a") (var "b") (var "c")
        ]
    , test "list comprehension"
        [ "[x | x <- [1 .. 10]]" :~
            listComp (bvar "x") [bvar "x" <-- fromTo (int 1) (int 10)]
        , "[x + y | x <- [1 .. 10], y <- [11 .. 20], even y]" :~
            listComp (op (bvar "x") "+" (bvar "y"))
                     [ bvar "x" <-- fromTo (int 1) (int 10)
                     , bvar "y" <-- fromTo (int 11) (int 20)
                     , stmt $ var "even" @@ bvar "y"
                     ]
        ]
    ]
  where
    test = testExprs dflags

declsTest dflags = testGroup "Decls"
    [ test "patBind"
        [ "x = x" :~ patBind (bvar "x") (var "x")
        , "(x, y) = (y, x)" :~ patBind (tuple [bvar "x", bvar "y"])
                                    (tuple [var "y", var "x"])
        , "(x, y)\n  | test = (1, 2)\n  | otherwise = (2, 3)" :~
            patBindGRHSs (tuple [bvar "x", bvar "y"])
                $ guardedRhs
                    [ var "test" `guard` tuple [int 1, int 2]
                        , var "otherwise" `guard` tuple [int 2, int 3]
                    ]
        , "z | Just y <- x, y = ()" :~
            patBindGRHSs (bvar "z")
                $ guardedRhs
                    [guards
                        [ conP "Just" [bvar "y"] <-- var "x"
                        , stmt (var "y")
                        ]
                        unit
                    ]
        ]
    , test "valBind"
        [ "x = y" :~ valBindGRHSs "x" $ rhs $ var "y"
        , "x = y" :~ valBind "x" $ var "y"
        , "x | test = 1\n  | otherwise = 2" :~
            valBindGRHSs "x"
            $ guardedRhs
                [ var "test" `guard` int 1
                , var "otherwise" `guard` int 2
                ]
        , "x = (+)" :~ valBind "x" $ var "+"
        ]
    , test "funBind"
        [ "not True = False\nnot False = True" :~
             funBinds "not"
                [ match [bvar "True"] (var "False")
                , match [bvar "False"] (var "True")
                ]
        , "True && True = True\nTrue && False = False" :~
             funBindsWithFixity Nothing "&&"
                [ match [bvar "True", bvar "True"] (var "True")
                , match [bvar "True", bvar "False"] (var "False")
                ]
        , "not x\n  | x = False\n  | otherwise = True" :~
            funBind "not"
                $ matchGRHSs [bvar "x"] $ guardedRhs
                    [ guard (var "x") (var "False")
                    , guard (var "otherwise") (var "True")
                    ]
        , "f (A x) = 1" :~ funBind "f" $ match [conP "A" [bvar "x"]] (int 1)
        ]
    , test "tyFamInst"
        [ "type instance Elt String = Char"
            :~ tyFamInst "Elt" [var "String"] (var "Char")
        , "instance Container String where\n  type Elt String = Char"
            :~ instance' (var "Container" @@ var "String")
                [tyFamInst "Elt" [var "String"] (var "Char")]
        ]
    , test "patSynSigs"
        [ "pattern F, G :: T" :~ patSynSigs ["F", "G"] $ var "T"
        , "pattern F :: T" :~ patSynSig "F" $ var "T"
        ]
    , test "patSynBind"
        [ "pattern F a b = G b a"
            :~ patSynBind "F" ["a", "b"] $ conP "G" [bvar "b", bvar "a"]
        ]
   , test "dataDecl"
        [ "data Either a b\n  = Left a | Right b\n  deriving Show"
            :~ data' "Either" [bvar "a", bvar "b"]
               [ prefixCon "Left" [field $ var "a"]
               , prefixCon "Right" [field $ var "b"]
               ] $ [deriving' [var "Show"]]
        , "data Either a (b :: Type)\n  = Left a | Right b\n  deriving Show"
            :~ data' "Either" [bvar "a", kindedVar "b" (var "Type")]
               [ prefixCon "Left" [field $ var "a"]
               , prefixCon "Right" [field $ var "b"]
               ] $ [deriving' [var "Show"]]
        ]
    , test "newtypeDecl"
        [ "newtype Const a b\n  = Const a\n  deriving Show"
            :~ newtype' "Const" [bvar "a", bvar "b"]
               (prefixCon "Const" [field $ var "a"])
               $ [deriving' [var "Show"]]
        , "newtype Const a (b :: Type)\n  = Const a\n  deriving Show"
            :~ newtype' "Const" [bvar "a", kindedVar "b" (var "Type")]
               (prefixCon "Const" [field $ var "a"])
               [deriving' [var "Show"]]
        ]
    , test "standaloneDeriving"
        [ "deriving instance Show Int"
            :~ standaloneDeriving (var "Show" @@ var "Int")
        , "deriving instance Show a => Show (Maybe a)"
            :~ standaloneDeriving ([var "Show" @@ var "a"] ==> var "Show" @@ (var "Maybe" @@ var "a"))
        , "deriving stock instance Show Int"
            :~ standaloneDerivingStock (var "Show" @@ var "Int")
        , "deriving newtype instance Show a => Show (Identity a)"
            :~ standaloneDerivingNewtype ([var "Show" @@ var "a"] ==> var "Show" @@ (var "Identity" @@ var "a"))
        , "deriving anyclass instance Show Person"
            :~ standaloneDerivingAnyclass (var "Show" @@ var "Person")
        ]
    ]
  where
    test = testDecls dflags

patsTest dflags = testGroup "Pats"
    [ test "app"
        [ "A x y" :~ conP "A" [bvar "x", bvar "y"]
        , "(:) x y" :~ conP ":" [bvar "x", bvar "y"]
        , "(Prelude.:) x" :~ conP "Prelude.:" [bvar "x"]
        , "A (B x)" :~ conP "A" [conP "B" [bvar "x"]]
        , "A (B x)" :~ conP "A" [par $ conP "B" [bvar "x"]]
        , "A ((B x))" :~ conP "A" [par $ par $ conP "B" [bvar "x"]]
        , "A x (B y z)" :~ conP "A" [bvar "x", conP "B" [bvar "y", bvar "z"]]
        , "A w (B x y) Z"
            :~ conP "A" [bvar "w", conP "B" [bvar "x", bvar "y"], conP "Z" []]
        , "A 3" :~ conP "A" [int 3]
        , "A (-3)" :~ conP "A" [int (-3)]
        , "A 3.0" :~ conP "A" [frac 3.0]
        , "A (-3.0)" :~ conP "A" [frac (-3.0)]
        , "A 'x'" :~ conP "A" [char 'x']
        , "A \"xyz\"" :~ conP "A" [string "xyz"]
        , "A B {x = C}"
            :~ conP "A" [recordConP "B" [("x", conP "C" [])]]
        ]
    , test "asP"
        [ "x@B" :~ asP "x" $ conP "B" []
        , "x@(B y)" :~ asP "x" $ conP "B" [bvar "y"]
        , "x@_" :~ asP "x" wildP
        ]
    , test "strictP"
        [ "!x" :~ strictP $ bvar "x"
        , "!B" :~ strictP $ conP "B" []
        , "!(B y)" :~ strictP $ conP "B" [bvar "y"]
        , "!_" :~ strictP wildP
        ]
    , test "lazyP"
        [ "~x" :~ lazyP $ bvar "x"
        , "~B" :~ lazyP $ conP "B" []
        , "~(B y)" :~ lazyP $ conP "B" [bvar "y"]
        , "~_" :~ lazyP wildP
        ]
    , test "sigPat"
        [ "x :: A" :~ sigP (bvar "x") (bvar "A")
        , "A x :: A x" :~ sigP (conP "A" [bvar "x"]) (bvar "A" @@ bvar "x")
        ]
    , test "recordConP"
        [ "A {x = Y}" :~ recordConP "A" [("x", conP "Y" [])]
        ]
    ]
  where
    test = testPats dflags

-- TODO: Add more test cases from pprint_examples.hs.
modulesTest dflags = testGroup "Modules"
    [ test "import"
        [ "import M" :~
            module' Nothing Nothing [import' "M"] []
        , "import {-# SOURCE #-} M" :~
            module' Nothing Nothing
                [source $ import' "M"] []
        ]
    ]
  where
    test = testModule dflags