{-|
Module      : Idris.CaseSplit
Description : Module to provide case split functionality.

License     : BSD3
Maintainer  : The Idris Community.

Given a pattern clause and a variable 'n', elaborate the clause and find the
type of 'n'.

Make new pattern clauses by replacing 'n' with all the possibly constructors
applied to '_', and replacing all other variables with '_' in order to
resolve other dependencies.

Finally, merge the generated patterns with the original, by matching.
Always take the "more specific" argument when there is a discrepancy, i.e.
names over '_', patterns over names, etc.
-}

{-# LANGUAGE FlexibleContexts, PatternGuards #-}

module Idris.CaseSplit(
    splitOnLine, replaceSplits
  , getClause, getProofClause
  , mkWith
  , nameMissing
  , getUniq, nameRoot
  ) where

import Idris.AbsSyntax
import Idris.Core.Evaluate
import Idris.Core.TT
import Idris.Delaborate
import Idris.Elab.Term
import Idris.Elab.Value
import Idris.ElabDecls
import Idris.Error
import Idris.Output
import Idris.Parser

import Control.Monad
import Control.Monad.State.Strict
import Data.Char
import Data.List (isPrefixOf, isSuffixOf)

-- | Given a variable to split, and a term application, return a list
-- of variable updates, paired with a flag to say whether the given
-- update typechecks (False = impossible) if the flag is 'False' the
-- splits should be output with the 'impossible' flag, otherwise they
-- should be output as normal
split :: Name -> PTerm -> Idris (Bool, [[(Name, PTerm)]])
split n t'
   = do ist <- getIState
        -- Make sure all the names in the term are accessible
        mapM_ (\n -> setAccessibility n Public) (allNamesIn t')
        -- ETyDecl rather then ELHS because there'll be explicit type
        -- matching
        (tm, ty, pats) <- elabValBind (recinfo (fileFC "casesplit")) ETyDecl True (addImplPat ist t')
        -- ASSUMPTION: tm is in normal form after elabValBind, so we don't
        -- need to do anything special to find out what family each argument
        -- is in
        logElab 4 ("Elaborated:\n" ++ show tm ++ " : " ++ show ty ++ "\n" ++ show pats)
--         iputStrLn (show (delab ist tm) ++ " : " ++ show (delab ist ty))
--         iputStrLn (show pats)
        let t = mergeUserImpl (addImplPat ist t') (delabDirect ist tm)
        let ctxt = tt_ctxt ist
        case lookup n pats of
             Nothing -> ifail $ show n ++ " is not a pattern variable"
             Just ty ->
                do let splits = findPats ist ty
                   logElab 1 ("New patterns " ++ showSep ", "
                         (map showTmImpls splits))
                   let newPats_in = zipWith (replaceVar ctxt n) splits (repeat t)
                   logElab 4 ("Working from " ++ showTmImpls t)
                   logElab 4 ("Trying " ++ showSep "\n"
                               (map (showTmImpls) newPats_in))
                   newPats_in <- mapM elabNewPat newPats_in
                   case anyValid [] [] newPats_in of
                        Left fails -> do
                           let fails' = mergeAllPats ist n t fails
                           return (False, (map snd fails'))
                        Right newPats -> do
                           logElab 3 ("Original:\n" ++ show t)
                           logElab 3 ("Split:\n" ++
                                      (showSep "\n" (map show newPats)))
                           logElab 3 "----"
                           let newPats' = mergeAllPats ist n t newPats
                           logElab 1 ("Name updates " ++ showSep "\n"
                                 (map (\ (p, u) -> show u ++ " " ++ show p) newPats'))
                           return (True, (map snd newPats'))
   where
     anyValid ok bad [] = if null ok then Left (reverse bad)
                                     else Right (reverse ok)
     anyValid ok bad ((tc, p) : ps)
         | tc = anyValid (p : ok) bad ps
         | otherwise = anyValid ok (p : bad) ps

data MergeState = MS { namemap :: [(Name, Name)],
                       invented :: [(Name, Name)],
                       explicit :: [Name],
                       updates :: [(Name, PTerm)] }

addUpdate :: Name -> PTerm -> State MergeState ()
addUpdate n tm = do ms <- get
                    put (ms { updates = ((n, stripNS tm) : updates ms) } )

inventName :: IState -> Maybe Name -> Name -> State MergeState Name
inventName ist ty n =
    do ms <- get
       let supp = case ty of
                       Nothing -> []
                       Just t -> getNameHints ist t
       let nsupp = case n of
                        MN i n | not (tnull n) && thead n == '_'
                               -> mkSupply (supp ++ varlist)
                        MN i n -> mkSupply (UN n : supp ++ varlist)
                        UN n | thead n == '_'
                               -> mkSupply (supp ++ varlist)
                        x -> mkSupply (x : supp)
       let badnames = map snd (namemap ms) ++ map snd (invented ms) ++
                      explicit ms
       case lookup n (invented ms) of
          Just n' -> return n'
          Nothing ->
             do let n' = uniqueNameFrom nsupp badnames
                put (ms { invented = (n, n') : invented ms })
                return n'

mkSupply :: [Name] -> [Name]
mkSupply ns = mkSupply' ns (map nextName ns)
  where mkSupply' xs ns' = xs ++ mkSupply ns'

varlist :: [Name]
varlist = map (sUN . (:[])) "xyzwstuv" -- EB's personal preference :)

stripNS :: PTerm -> PTerm
stripNS tm = mapPT dens tm where
    dens (PRef fc hls n) = PRef fc hls (nsroot n)
    dens t = t

mergeAllPats :: IState -> Name -> PTerm -> [PTerm] -> [(PTerm, [(Name, PTerm)])]
mergeAllPats ist cv t [] = []
mergeAllPats ist cv t (p : ps)
    = let (p', MS _ _ _ u) = runState (mergePat ist t p Nothing)
                                      (MS [] [] (filter (/=cv) (patvars t)) [])
          ps' = mergeAllPats ist cv t ps in
          ((p', u) : ps')
  where patvars (PRef _ _ n) = [n]
        patvars (PApp _ _ as) = concatMap (patvars . getTm) as
        patvars (PPatvar _ n) = [n]
        patvars _ = []

mergePat :: IState -> PTerm -> PTerm -> Maybe Name -> State MergeState PTerm
mergePat ist orig new t =
    do -- collect user names for name map, by matching user pattern against
       -- the generated pattern
       case matchClause ist orig new of
            Left _ -> return ()
            Right ns -> mapM_ addNameMap ns
       mergePat' ist orig new t
  where
    addNameMap (n, PRef fc _ n') = do ms <- get
                                      put (ms { namemap = ((n', n) : namemap ms) })
    addNameMap _ = return ()

-- | If any names are unified, make sure they stay unified. Always
-- prefer user provided name (first pattern)
mergePat' ist (PPatvar fc n) new t
  = mergePat' ist (PRef fc [] n) new t
mergePat' ist old (PPatvar fc n) t
  = mergePat' ist old (PRef fc [] n) t
mergePat' ist orig@(PRef fc _ n) new@(PRef _ _ n') t
  | isDConName n' (tt_ctxt ist) = do addUpdate n new
                                     return new
  | otherwise
    = do ms <- get
         case lookup n' (namemap ms) of
              Just x -> do addUpdate n (PRef fc [] x)
                           return (PRef fc [] x)
              Nothing -> do put (ms { namemap = ((n', n) : namemap ms) })
                            return (PRef fc [] n)
mergePat' ist (PApp _ _ args) (PApp fc f args') t
      = do newArgs <- zipWithM mergeArg args (zip args' (argTys ist f))
           return (PApp fc f newArgs)
   where mergeArg x (y, t)
              = do tm' <- mergePat' ist (getTm x) (getTm y) t
                   case x of
                        (PImp _ _ _ _ _) ->
                             return (y { machine_inf = machine_inf x,
                                         getTm = tm' })
                        _ -> return (y { getTm = tm' })
mergePat' ist (PRef fc _ n) tm ty = do tm <- tidy ist tm ty
                                       addUpdate n tm
                                       return tm
mergePat' ist x y t = return y

mergeUserImpl :: PTerm -> PTerm -> PTerm
mergeUserImpl x y = x

argTys :: IState -> PTerm -> [Maybe Name]
argTys ist (PRef fc hls n)
    = case lookupTy n (tt_ctxt ist) of
           [ty] -> let ty' = normalise (tt_ctxt ist) [] ty in
                       map (tyName . snd) (getArgTys ty') ++ repeat Nothing
           _ -> repeat Nothing
  where tyName (Bind _ (Pi _ _ _ _) _) = Just (sUN "->")
        tyName t | (P _ d _, [_, ty]) <- unApply t,
                   d == sUN "Delayed" = tyName ty
                 | (P _ n _, _) <- unApply t = Just n
                 | otherwise = Nothing
argTys _ _ = repeat Nothing

tidy :: IState -> PTerm -> Maybe Name -> State MergeState PTerm
tidy ist orig@(PRef fc hls n) ty
     = do ms <- get
          case lookup n (namemap ms) of
               Just x -> return (PRef fc [] x)
               Nothing -> case n of
                               (UN _) -> return orig
                               _ -> do n' <- inventName ist ty n
                                       return (PRef fc [] n')
tidy ist (PApp fc f args) ty
     = do args' <- zipWithM tidyArg args (argTys ist f)
          return (PApp fc f args')
    where tidyArg x ty' = do tm' <- tidy ist (getTm x) ty'
                             return (x { getTm = tm' })
tidy ist tm ty = return tm


-- mapPT tidyVar tm
--   where tidyVar (PRef _ _) = Placeholder
--         tidyVar t = t

elabNewPat :: PTerm -> Idris (Bool, PTerm)
elabNewPat t = idrisCatch (do (tm, ty) <- elabVal (recinfo (fileFC "casesplit")) ELHS t
                              i <- getIState
                              return (True, delabDirect i tm))
                          (\e -> do i <- getIState
                                    logElab 5 $ "Not a valid split:\n" ++ showTmImpls t ++ "\n"
                                                     ++ pshow i e
                                    return (False, t))

findPats :: IState -> Type -> [PTerm]
findPats ist t | (P _ n _, _) <- unApply t
    = case lookupCtxt n (idris_datatypes ist) of
           [ti] -> map genPat (con_names ti)
           _ -> [Placeholder]
    where genPat n = case lookupCtxt n (idris_implicits ist) of
                        [args] -> PApp emptyFC (PRef emptyFC [] n)
                                         (map toPlaceholder args)
                        _ -> error $ "Can't happen (genPat) " ++ show n
          toPlaceholder tm = tm { getTm = Placeholder }
findPats ist t = [Placeholder]

replaceVar :: Context -> Name -> PTerm -> PTerm -> PTerm
replaceVar ctxt n t (PApp fc f pats) = PApp fc f (map substArg pats)
  where subst :: PTerm -> PTerm
        subst orig@(PPatvar _ v) | v == n = t
                                 | otherwise = Placeholder
        subst orig@(PRef _ _ v) | v == n = t
                                | isDConName v ctxt = orig
        subst (PRef _ _ _) = Placeholder
        subst (PApp fc (PRef _ _ t) pats)
            | isTConName t ctxt = Placeholder -- infer types
        subst (PApp fc f pats) = PApp fc f (map substArg pats)
        subst x = x

        substArg arg = arg { getTm = subst (getTm arg) }

replaceVar ctxt n t pat = pat

splitOnLine :: Int         -- ^ line number
            -> Name        -- ^ variable
            -> FilePath    -- ^ name of file
            -> Idris (Bool, [[(Name, PTerm)]])
splitOnLine l n fn = do
    cl <- getInternalApp fn l
    logElab 3 ("Working with " ++ showTmImpls cl)
    tms <- split n cl
    return tms

replaceSplits :: String -> [[(Name, PTerm)]] -> Bool -> Idris [String]
replaceSplits l ups impossible
    = do ist <- getIState
         updateRHSs 1 (map (rep ist (expandBraces l)) ups)
  where
    rep ist str [] = str ++ "\n"
    rep ist str ((n, tm) : ups)
        = rep ist (updatePat False (show n) (nshow (resugar ist tm)) str) ups

    updateRHSs i [] = return []
    updateRHSs i (x : xs)
       | impossible = do xs' <- updateRHSs i xs
                         return (setImpossible False x : xs')
       | otherwise = do (x', i') <- updateRHS (null xs) i x
                        xs' <- updateRHSs i' xs
                        return (x' : xs')

    updateRHS last i ('?':'=':xs) = do (xs', i') <- updateRHS last i xs
                                       return ("?=" ++ xs', i')
    updateRHS last i ('?':xs)
        = do let (nm, rest_in) = span (not . (\x -> isSpace x || x == ')'
                                                              || x == '(')) xs
             let rest = if last then rest_in else
                           case span (not . (=='\n')) rest_in of
                                (_, restnl) -> restnl
             (nm', i') <- getUniq nm i
             return ('?':nm' ++ rest, i')
    updateRHS last i (x : xs) = do (xs', i') <- updateRHS last i xs
                                   return (x : xs', i')
    updateRHS last i [] = return ("", i)

    setImpossible brace ('}':xs) = '}' : setImpossible False xs
    setImpossible brace ('{':xs) = '{' : setImpossible True xs
    setImpossible False ('=':xs) = "impossible\n"
    setImpossible brace (x : xs) = x : setImpossible brace xs
    setImpossible brace [] = ""

    -- TMP HACK: If there are Nats, we don't want to show as numerals since
    -- this isn't supported in a pattern, so special case here
    nshow (PRef _ _ (UN z)) | z == txt "Z" = "Z"
    nshow (PApp _ (PRef _ _ (UN s)) [x]) | s == txt "S" =
               "(S " ++ addBrackets (nshow (getTm x)) ++ ")"
    nshow t = show t

    -- if there's any {n} replace with {n=n}
    -- but don't replace it in comments
    expandBraces ('{' : '-' : xs) = '{' : '-' : xs
    expandBraces ('{' : xs)
        = let (brace, (_:rest)) = span (/= '}') xs in
              if (not ('=' `elem` brace))
                 then ('{' : brace ++ " = " ++ brace ++ "}") ++
                         expandBraces rest
                 else ('{' : brace ++ "}") ++ expandBraces rest
    expandBraces (x : xs) = x : expandBraces xs
    expandBraces [] = []

    updatePat start n tm [] = []
    updatePat start n tm ('{':rest) =
        let (space, rest') = span isSpace rest in
            '{' : space ++ updatePat False n tm rest'
    updatePat start n tm done@('?':rest) = done
    updatePat True n tm xs@(c:rest) | length xs > length n
        = let (before, after@(next:_)) = splitAt (length n) xs in
              if (before == n && not (isAlphaNum next))
                 then addBrackets tm ++ updatePat False n tm after
                 else c : updatePat (not (isAlphaNum c)) n tm rest
    updatePat start n tm (c:rest) = c : updatePat (not ((isAlphaNum c) || c == '_')) n tm rest

    addBrackets tm | ' ' `elem` tm
                   , not (isPrefixOf "(" tm && isSuffixOf ")" tm)
                       = "(" ++ tm ++ ")"
                   | otherwise = tm


getUniq :: (Show t, Num t) => [Char] -> t -> Idris ([Char], t)
getUniq nm i
       = do ist <- getIState
            let n = nameRoot [] nm ++ "_" ++ show i
            case lookupTy (sUN n) (tt_ctxt ist) of
                 [] -> return (n, i+1)
                 _ -> getUniq nm (i+1)

nameRoot acc nm | all isDigit nm = showSep "_" acc
nameRoot acc nm =
        case span (/='_') nm of
             (before, ('_' : after)) -> nameRoot (acc ++ [before]) after
             _ -> showSep "_" (acc ++ [nm])

-- Show a name for use in pattern definitions on the lhs
showLHSName :: Name -> String
showLHSName n = let fn = show n in
                    if any (flip elem opChars) fn
                       then "(" ++ fn ++ ")"
                       else fn

-- Show a name for the purpose of generating a metavariable name on the rhs
showRHSName :: Name -> String
showRHSName n = let fn = show n in
                    if any (flip elem opChars) fn
                       then "op"
                       else fn

getClause :: Int      -- ^ line number that the type is declared on
          -> Name     -- ^ Function name
          -> Name     -- ^ User given name
          -> FilePath -- ^ Source file name
          -> Idris String
getClause l _ un fp
    = do i <- getIState
         indentClause <- getIndentClause
--         runIO . traceIO $ "indentClause = " ++ show indentClause
         case lookupCtxt un (idris_interfaces i) of
              [c] -> return (mkInterfaceBodies indentClause i (interface_methods c))
              _ -> do ty_in <- getInternalApp fp l
                      let ty = case ty_in of
                                    PTyped _ t -> t
                                    x -> x
                      ist <- get
                      let app = mkApp ist ty []
                      return (showLHSName un ++ " " ++ app ++ "= ?"
                                      ++ showRHSName un ++ "_rhs")
   where mkApp :: IState -> PTerm -> [Name] -> String
         mkApp i (PPi (Exp _ _ False _) (MN _ _) _ ty sc) used
               = let n = getNameFrom i used ty in
                     show n ++ " " ++ mkApp i sc (n : used)
         mkApp i (PPi (Exp _ _ False _) (UN n) _ ty sc) used
            | thead n == '_'
               = let n = getNameFrom i used ty in
                     show n ++ " " ++ mkApp i sc (n : used)
         mkApp i (PPi (Exp _ _ False _) n _ _ sc) used
               = show n ++ " " ++ mkApp i sc (n : used)
         mkApp i (PPi _ _ _ _ sc) used = mkApp i sc used
         mkApp i _ _ = ""

         getNameFrom i used (PPi _ _ _ _ _)
              = uniqueNameFrom (mkSupply [sUN "f", sUN "g"]) used
         getNameFrom i used (PApp fc f as) = getNameFrom i used f
         getNameFrom i used (PRef fc _ f)
            = case getNameHints i f of
                   [] -> uniqueNameFrom (mkSupply [sUN "x", sUN "y",
                                                   sUN "z"]) used
                   ns -> uniqueNameFrom (mkSupply ns) used
         getNameFrom i used _ = uniqueNameFrom (mkSupply [sUN "x", sUN "y",
                                                          sUN "z"]) used

         -- write method declarations, indent with `indent` spaces.
         mkInterfaceBodies :: Int -> IState -> [(Name, (Bool, FnOpts, PTerm))] -> String
         mkInterfaceBodies indent i ns
             = showSep "\n"
                  (zipWith (\(n, (_, _, ty)) m -> replicate indent ' ' ++
                            def (show (nsroot n)) ++ " "
                                 ++ mkApp i ty []
                                 ++ "= ?"
                                 ++ showRHSName un ++ "_rhs_" ++ show m) ns [1..])

         def n@(x:xs) | not (isAlphaNum x) = "(" ++ n ++ ")"
         def n = n

getProofClause :: Int      -- ^ line number that the type is declared
               -> Name     -- ^ Function name
               -> FilePath -- ^ Source file name
               -> Idris String
getProofClause l fn fp
                  = do ty_in <- getInternalApp fp l
                       let ty = case ty_in of
                                     PTyped n t -> t
                                     x -> x
                       return (mkApp ty ++ " = ?" ++ showRHSName fn ++ "_rhs")
   where mkApp (PPi _ _ _ _ sc) = mkApp sc
         mkApp rt = "(" ++ show rt ++ ") <== " ++ show fn

-- Purely syntactic - turn a pattern match clause into a with and a new
-- match clause

mkWith :: String -> Name -> Int -> String
mkWith str n indent = let str' = replaceRHS str "with (_)"
                      in str' ++ "\n" ++ newpat str

   where replaceRHS [] str = str
         replaceRHS ('?':'=': rest) str = str
         replaceRHS ('=': rest) str
              | not ('=' `elem` rest) = str
         replaceRHS (x : rest) str = x : replaceRHS rest str

         newpat ('>':patstr) = '>':newpat patstr
         newpat patstr = replicate indent ' ' ++
           replaceRHS patstr "| with_pat = ?" ++ showRHSName n ++ "_rhs"

-- Replace _ with names in missing clauses

nameMissing :: [PTerm] -> Idris [PTerm]
nameMissing ps = do ist <- get
                    newPats <- mapM nm ps
                    let newPats' = mergeAllPats ist (sUN "_") (base (head ps))
                                                newPats
                    return (map fst newPats')
  where
    base (PApp fc f args) = PApp fc f (map (fmap (const (PRef fc [] (sUN "_")))) args)
    base t = t

    nm ptm = do mptm <- elabNewPat ptm
                case mptm of
                     (False, _) -> return ptm
                     (True, ptm') -> return ptm'