module SimplifySyntax (
  SimpExp(..)
  , SelectorAndVal(..)
  , SimpAlt(..)
  , SimpDecl(..)
  , SimpPat(..)
  , stringToSimpDecl
  , qNameToString
  , nameToString
  , customParseDecl
  , hsDeclToSimpDecl
  , formatString
  ) where

import Data.List(foldl')
import Data.Maybe(catMaybes, isJust)

import qualified Language.Haskell.Exts as Exts

import TranslateCore(nTupleSectionString, nTupleString, nListString)

-- TODO use a data constructor for the special case instead of using string
-- matching for tempvars.
-- There is a special case in Icons.hs/makeLabelledPort to exclude " tempvar"
tempVarPrefix :: String
tempVarPrefix = " tempvar"

-- A simplified Haskell syntax tree

-- rhs is now SimpExp

-- A simplified Haskell expression.
data SimpExp l =
  SeName l String
  | SeLit l (Exts.Literal l)
  | SeApp l
    (SimpExp l)  -- function
    (SimpExp l)  -- argument
  | SeLambda l [SimpPat l] (SimpExp l)
  | SeLet l [SimpDecl l] (SimpExp l)
  | SeCase l (SimpExp l) [SimpAlt l]
  | SeMultiIf l [SelectorAndVal l]
  deriving (Show, Eq)

data SelectorAndVal l = SelectorAndVal {
  svSelector :: SimpExp l
  , svVal :: SimpExp l
  }
  deriving (Show, Eq)

data SimpAlt l = SimpAlt {
  saPat :: SimpPat l
  , saVal :: SimpExp l
  }
  deriving (Show, Eq)

data SimpDecl l =
  -- These don't have decl lists, since only lets have decl lists
  SdPatBind l (SimpPat l) (SimpExp l)
  | SdTypeSig l [Exts.Name l] (Exts.Type l)
  -- TODO Add a visual representation of data declarations
  | SdCatchAll (Exts.Decl l)
  deriving (Show, Eq)

data SimpPat l =
  SpVar l (Exts.Name l)
  | SpLit l (Exts.Sign l) (Exts.Literal l)
  | SpApp l (Exts.QName l) [SimpPat l]
  | SpAsPat l (Exts.Name l) (SimpPat l)
  | SpWildCard l
  deriving (Show, Eq)

-- Helper functions

strToQName :: l -> String -> Exts.QName l
strToQName l = Exts.UnQual l . Exts.Ident l

makeVarExp :: l -> String -> Exts.Exp l
makeVarExp l = Exts.Var l . strToQName l

makePatVar :: l -> String -> Exts.Pat l
makePatVar l = Exts.PVar l . Exts.Ident l

qOpToExp :: Exts.QOp l -> Exts.Exp l
qOpToExp (Exts.QVarOp l n) = Exts.Var l n
qOpToExp (Exts.QConOp l n) = Exts.Con l n

nameToString :: Exts.Name l -> String
nameToString (Exts.Ident _ s) = s
nameToString (Exts.Symbol _ s) = s

qNameToString :: Show l => Exts.QName l -> String
qNameToString qName = case qName of
  Exts.Qual _ (Exts.ModuleName _ modName) name
    -> modName ++ "." ++ nameToString name
  Exts.UnQual _ name -> nameToString name
  Exts.Special _ constructor -> case constructor of
    Exts.UnitCon _ -> "()"
    Exts.ListCon _ -> "[]"
    Exts.FunCon _ -> "(->)"
    Exts.TupleCon _ _ n -> nTupleString n
    Exts.Cons _ -> "(:)"
    -- unboxed singleton tuple constructor
    Exts.UnboxedSingleCon _ -> "(# #)"
    _ -> error $ "Unsupported syntax in qNameToSrting: " <> show qName

simpExpToRhs :: Show l => l -> SimpExp l -> Exts.Rhs l
simpExpToRhs l e = Exts.UnGuardedRhs l (simpExpToHsExp e)

--

hsPatToSimpPat :: Show a => Exts.Pat a -> SimpPat a
hsPatToSimpPat p = case p of
  Exts.PVar l n -> SpVar l n
  Exts.PLit l sign lit -> SpLit l sign lit
  Exts.PInfixApp l p1 qName p2 -> hsPatToSimpPat (Exts.PApp l qName [p1, p2])
  Exts.PApp l name patts -> SpApp l name (fmap hsPatToSimpPat patts)
  Exts.PTuple l _ patts -> SpApp
                           l
                           ((strToQName l . nTupleString . length) patts)
                           (fmap hsPatToSimpPat patts)
  Exts.PParen _ pat -> hsPatToSimpPat pat
  Exts.PAsPat l name pat -> SpAsPat l name (hsPatToSimpPat pat)
  Exts.PWildCard l -> SpWildCard l
  Exts.PList l patts -> SpApp
                        l
                        ((strToQName l . nListString . length) patts)
                        (fmap hsPatToSimpPat patts)
  _ -> error $  "Unsupported syntax in hsPatToSimpPat: " <> show p

simpPatToHsPat :: Show a => SimpPat a -> Exts.Pat a
simpPatToHsPat pat = case pat of
  SpVar l n -> Exts.PVar l n
  SpLit l s lit -> Exts.PLit l s lit
  SpApp l n pats -> Exts.PApp l n (fmap simpPatToHsPat pats)
  SpAsPat l n p -> Exts.PAsPat l n (simpPatToHsPat p)
  SpWildCard l -> Exts.PWildCard l

whereToLet :: Show a => a -> Exts.Rhs a -> Maybe (Exts.Binds a) -> SimpExp a
whereToLet l rhs maybeBinds = val
  where
    rhsExp = hsRhsToExp rhs
    val = case maybeBinds of
      Nothing -> rhsExp
      Just binds -> SeLet l (hsBindsToDecls binds) rhsExp

matchToSimpDecl :: Show a => Exts.Match a -> SimpDecl a
matchToSimpDecl (Exts.Match l name patterns rhs maybeWhereBinds)
  = SdPatBind
    l
    (SpVar l name)
    (SeLambda l
     (fmap hsPatToSimpPat patterns)
     (whereToLet l rhs maybeWhereBinds))
matchToSimpDecl m = error $ "Unsupported syntax in matchToSimpDecl: " <> show m

-- Only used by matchesToCase
matchToAlt :: Show l => Exts.Match l -> Exts.Alt l
matchToAlt (Exts.Match l _ mtaPats rhs binds)
  = Exts.Alt l altPattern rhs binds
  where
    altPattern = case mtaPats of
      [onePat] -> onePat
      _ -> Exts.PTuple l Exts.Boxed mtaPats
matchToAlt match = error $ "Unsupported syntax in matchToAlt: " <> show match

matchesToFunBind :: Show a => a -> [Exts.Match a] -> SimpDecl a
matchesToFunBind l matches = matchToSimpDecl $ case matches of
  [] -> error $ "Empty matches in matchesToFunBind. Label is :" <> show l
  [match] -> match
  (Exts.Match srcLoc funName pats _ _ : _)
    -> Exts.Match srcLoc funName tempPats rhs Nothing
    where
      -- There is a special case in Icons.hs/makeLabelledPort to exclude " tempvar"
      tempStrings = fmap (\x -> tempVarPrefix ++ show x) [0..(length pats - 1)]
      tempPats = fmap (makePatVar srcLoc) tempStrings
      tempVars = fmap (makeVarExp srcLoc) tempStrings
      tuple = Exts.Tuple srcLoc Exts.Boxed tempVars
      caseExp = case tempVars of
        [oneTempVar] -> Exts.Case srcLoc oneTempVar alts
        _ -> Exts.Case srcLoc tuple alts
      rhs = Exts.UnGuardedRhs srcLoc caseExp
      alts = fmap matchToAlt matches
  _ -> error $ "Unsupported syntax in matchesToFunBind: " <> show matches


hsDeclToSimpDecl :: Show a => Exts.Decl a -> SimpDecl a
hsDeclToSimpDecl decl = case decl of
  Exts.TypeSig l names typeForNames -> SdTypeSig l names typeForNames
  Exts.FunBind l matches -> matchesToFunBind l matches
  Exts.PatBind l pat rhs maybeBinds -> SdPatBind l (hsPatToSimpPat pat) expr
    where
      expr = whereToLet l rhs maybeBinds
  d -> SdCatchAll d

simpDeclToHsDecl :: Show a => SimpDecl a -> Exts.Decl a
simpDeclToHsDecl decl = case decl of
  SdPatBind l pat e
    -> Exts.PatBind l (simpPatToHsPat pat) (simpExpToRhs l e) Nothing
  SdTypeSig l names typeForNames -> Exts.TypeSig l names typeForNames
  SdCatchAll d -> d

hsBindsToDecls :: Show a => Exts.Binds a -> [SimpDecl a]
hsBindsToDecls binds = case binds of
  Exts.BDecls _ decls -> fmap hsDeclToSimpDecl decls
  _ -> error $ "Unsupported syntax in hsBindsToDecls: " ++ show binds

simpDeclsToHsBinds :: Show a => a -> [SimpDecl a] -> Exts.Binds a
simpDeclsToHsBinds l decls = Exts.BDecls l (fmap simpDeclToHsDecl decls)

guardedRhsToSelectorAndVal :: Show a => Exts.GuardedRhs a -> SelectorAndVal a
guardedRhsToSelectorAndVal rhs = case rhs of
  Exts.GuardedRhs _ [s] valExp -> SelectorAndVal{svSelector=stmtToExp s
                                                , svVal=hsExpToSimpExp valExp}
  _ -> error $ "Unsupported syntax in guardedRhsToSelectorAndVal: " ++ show rhs
  where
    stmtToExp stmt = case stmt of
      Exts.Qualifier _ e -> hsExpToSimpExp e
      _ -> error
           $ "Unsupported syntax in stmtToExp: " ++ show stmt

selAndValToGuardedRhs :: Show a => a -> SelectorAndVal a -> Exts.GuardedRhs a
selAndValToGuardedRhs l selAndVal = Exts.GuardedRhs
  l
  [Exts.Qualifier l (simpExpToHsExp $ svSelector selAndVal)]
  (simpExpToHsExp $ svVal selAndVal)

hsRhsToExp :: Show a => Exts.Rhs a -> SimpExp a
hsRhsToExp rhs = case rhs of
  Exts.UnGuardedRhs _ e -> hsExpToSimpExp e
  Exts.GuardedRhss l rhss
    -> SeMultiIf l (fmap guardedRhsToSelectorAndVal rhss)

hsAltToSimpAlt :: Show a => Exts.Alt a -> SimpAlt a
hsAltToSimpAlt (Exts.Alt l pat rhs maybeBinds)
  = SimpAlt{saPat=hsPatToSimpPat pat, saVal=whereToLet l rhs maybeBinds}

simpAltToHsAlt :: Show a => a -> SimpAlt a -> Exts.Alt a
simpAltToHsAlt l (SimpAlt pat e)
  = Exts.Alt l (simpPatToHsPat pat) (simpExpToRhs l e) Nothing

ifToGuard :: a -> SimpExp a -> SimpExp a -> SimpExp a -> SimpExp a
ifToGuard l e1 e2 e3
  = SeMultiIf l [SelectorAndVal{svSelector=e1, svVal=e2}
                , SelectorAndVal{svSelector=otherwiseExp, svVal=e3}]
  where
    otherwiseExp = SeName l "otherwise"

simplifyExp :: SimpExp l -> SimpExp l
simplifyExp e = case e of
  -- Reduce applications of function compositions (e.g. (f . g) x -> f (g x))
  SeApp l2 (SeApp l1 (SeApp _ (SeName _ ".") f1) f2) arg
    -> SeApp l1 f1 $ simplifyExp (SeApp l2 f2 arg)
  SeApp l (SeApp _ (SeName _ "$") exp1) exp2
    -> SeApp l exp1 exp2
  SeApp l1 (SeName l2 "<$>") arg
    -> SeApp l1 (SeName l2 "fmap") arg
  x -> x

deListifyApp :: Show l => l -> Exts.Exp l -> [Exts.Exp l] -> Exts.Exp l
deListifyApp l = foldl' (Exts.App l)

rewriteTupleSection :: Show l => l -> [Maybe (Exts.Exp l)] -> Exts.Exp l
rewriteTupleSection l mExprs = deListifyApp
                               l
                               (makeVarExp l $ nTupleSectionString expIsJustList)
                               exprs
  where
    exprs = catMaybes mExprs
    expIsJustList = fmap isJust mExprs

-- Rewrite a right section as a lambda.
-- TODO Simplify this type of lambda to use unused ports.
rewriteRightSection :: Show l => l -> Exts.QOp l -> Exts.Exp l -> Exts.Exp l
rewriteRightSection l op expr = Exts.Lambda l [tempPat] appExpr
  where
    tempStr = tempVarPrefix <> "0"
    tempPat = makePatVar l tempStr
    tempVar = makeVarExp l tempStr
    appExpr = Exts.App l (Exts.App l (qOpToExp op) tempVar) expr

desugarDo :: Show l => [Exts.Stmt l] -> Exts.Exp l
desugarDo stmts = case stmts of
  [Exts.Qualifier _ e] -> e
  (Exts.Qualifier l e : stmtsTail)
    -> Exts.App l (Exts.App l (makeVarExp l ">>") e) (desugarDo stmtsTail)
  (Exts.Generator l pat e : stmtsTail)
    -> Exts.App l (Exts.App l (makeVarExp l ">>=") e)
                  (Exts.Lambda l [pat] (desugarDo stmtsTail))
  (Exts.LetStmt l binds : stmtsTail) -> Exts.Let l binds (desugarDo stmtsTail)
  _ -> error $ "Unsupported syntax in degugarDo: " <> show stmts

desugarEnums :: Show l => l -> String -> [Exts.Exp l] -> SimpExp l
desugarEnums l funcName exprs = hsExpToSimpExp $ deListifyApp l
                                (makeVarExp l funcName)
                                exprs

hsExpToSimpExp :: Show a => Exts.Exp a -> SimpExp a
hsExpToSimpExp x = simplifyExp $ case x of
  Exts.Var l n -> SeName l (qNameToString n)
  Exts.Con l n -> SeName l (qNameToString n)
  Exts.Lit l n -> SeLit l n
  Exts.InfixApp l e1 op e2 ->
    hsExpToSimpExp $ Exts.App l (Exts.App l (qOpToExp op) e1) e2
  Exts.App l f arg -> SeApp l (hsExpToSimpExp f) (hsExpToSimpExp arg)
  Exts.NegApp l e -> hsExpToSimpExp $ Exts.App l (makeVarExp l "negate") e
  Exts.Lambda l patterns e
    -> SeLambda l (fmap hsPatToSimpPat patterns) (hsExpToSimpExp e)
  Exts.Let l bs e -> SeLet l (hsBindsToDecls bs) (hsExpToSimpExp e)
  Exts.If l e1 e2 e3
    -> ifToGuard l (hsExpToSimpExp e1) (hsExpToSimpExp e2) (hsExpToSimpExp e3)
  Exts.Case l e alts -> SeCase l (hsExpToSimpExp e) (fmap hsAltToSimpAlt alts)
  Exts.Paren _ e -> hsExpToSimpExp e
  Exts.List l exprs -> hsExpToSimpExp $ deListifyApp
                       l
                       (makeVarExp l $ nListString $ length exprs)
                       exprs
  Exts.Tuple l _ exprs -> hsExpToSimpExp $ deListifyApp
                       l
                       (makeVarExp l $ nTupleString $ length exprs)
                       exprs
  Exts.TupleSection l _ mExprs -> hsExpToSimpExp $ rewriteTupleSection l mExprs
  Exts.LeftSection l expr op -> hsExpToSimpExp $ Exts.App l (qOpToExp op) expr
  Exts.RightSection l op expr -> hsExpToSimpExp $ rewriteRightSection l op expr
  Exts.Do _ stmts -> hsExpToSimpExp $ desugarDo stmts
  Exts.EnumFrom l e -> desugarEnums l "enumFrom" [e]
  Exts.EnumFromTo l e1 e2 -> desugarEnums l "enumFromTo" [e1, e2]
  Exts.EnumFromThen l e1 e2 -> desugarEnums l "enumFromThen" [e1, e2]
  Exts.EnumFromThenTo l e1 e2 e3 -> desugarEnums l "enumFromThenTo" [e1, e2, e3]
  Exts.MultiIf l rhss -> SeMultiIf l (fmap guardedRhsToSelectorAndVal rhss)
  _ -> error $ "Unsupported syntax in hsExpToSimpExp: " ++ show x

simpExpToHsExp :: Show a => SimpExp a -> Exts.Exp a
simpExpToHsExp x = case x of
  -- TODO Sometimes SeName comes from Exts.Con
  --
  -- Put names in parens in case it's an operator
  SeName l str -> Exts.Paren l (Exts.Var l (strToQName l str))
  -- SeName l str -> (Exts.Var l (strToQName l str))
  SeLit l lit -> Exts.Lit l lit
  SeApp l e1 e2 -> Exts.App l (simpExpToHsExp e1) (simpExpToHsExp e2)
  SeLambda l pats e
    -> Exts.Lambda l (fmap simpPatToHsPat pats) (simpExpToHsExp e)
  SeLet l decls e -> Exts.Let l (simpDeclsToHsBinds l decls) (simpExpToHsExp e)
  SeCase l e alts
    -> Exts.Case l (simpExpToHsExp e) $ fmap (simpAltToHsAlt l) alts
  SeMultiIf l selsAndVal
    -> Exts.MultiIf l (fmap (selAndValToGuardedRhs l) selsAndVal)

-- Parsing

customParseMode :: Exts.ParseMode
customParseMode = Exts.defaultParseMode
  {Exts.extensions =
   [Exts.EnableExtension Exts.MultiParamTypeClasses,
    Exts.EnableExtension Exts.FlexibleContexts,
    Exts.EnableExtension Exts.TupleSections,
    Exts.EnableExtension Exts.MultiWayIf
   ]
  }

customParseDecl :: String -> Exts.Decl Exts.SrcSpanInfo
customParseDecl = Exts.fromParseResult . Exts.parseDeclWithMode customParseMode

stringToSimpDecl :: String -> SimpDecl Exts.SrcSpanInfo
stringToSimpDecl = hsDeclToSimpDecl . customParseDecl

formatString :: String -> Exts.Decl Exts.SrcSpanInfo
formatString = simpDeclToHsDecl . hsDeclToSimpDecl . customParseDecl