{-# LANGUAGE NamedFieldPuns, TupleSections #-}
module CheckMatch
    ( checkBody
    , Pat (..)
    , CanonicalPat
    ) where

import Control.Applicative ((<$>), (<*>))
import Control.Monad.Except (throwError)
import Control.Monad (liftM)
import Control.Monad (forM)

import Data.Maybe (mapMaybe, catMaybes)
import qualified Data.Set as Set
import qualified Data.Map as Map

import qualified AST.Expression.General as General
import qualified AST.Expression.Canonical as Canonical
import qualified AST.Module as Module
import qualified AST.Pattern as Pattern
import AST.Pattern (CanonicalPattern)
import qualified Reporting.Annotation as Annotation
import qualified Reporting.Region as Region
import qualified AST.Variable as Var
import qualified Reporting.Error.CheckMatch as Error
import Reporting.Warning.CheckMatch

fromPattern :: CanonicalPattern -> CanonicalPat
fromPattern (Annotation.A _region p) = case p of
    Pattern.Data ctor ps -> Data ctor (map fromPattern ps)
    Pattern.Record fs    -> Record fs
    Pattern.Alias s p    -> Alias s (fromPattern p)
    Pattern.Var s        -> Var s
    Pattern.Anything     -> Anything
    Pattern.Literal l    -> Literal l

concatMapM :: Monad m => (a -> m [b]) -> [a] -> m [b]
concatMapM f = liftM concat . mapM f

mapMaybeM :: Monad m => (a -> m (Maybe b)) -> [a] -> m [b]
mapMaybeM f = liftM catMaybes . mapM f

maybeCons :: Maybe a -> [a] -> [a]
maybeCons = maybe id (:)

type M = Either (Annotation.Located Error.Error)
checkBody
  :: Module.Interfaces
  -> Module.CanonicalBody
  -> M [Annotation.Located Warning]
checkBody interfaces = checkExpr . Module.program where
    checkExpr :: Canonical.Expr -> M [Annotation.Located Warning]
    checkExpr (Annotation.A r e) = case e of
        General.ExplicitList es  -> concatMapM checkExpr es
        General.Binop _ e e'     -> (++) <$> checkExpr e <*> checkExpr e'
        General.Lambda p e       -> maybeCons <$> checkMatch r [p] <*> checkExpr e
        General.App e e'         -> (++) <$> checkExpr e <*> checkExpr e'
        General.MultiIf branches -> concatMapM (\(b,e) -> (++) <$> checkExpr b <*> checkExpr e) branches

        General.Let defs e -> fmap concat $ sequence [ps, rhssResults, checkExpr e]
            where
            ps          = mapMaybeM (\(Canonical.Definition p _ _) -> checkMatch r [p]) defs
            rhssResults = concatMapM (\(Canonical.Definition _ e _) -> checkExpr e) defs

        General.Case e branches -> do
            eResults      <- checkExpr e
            result        <- checkMatch r ps
            branchResults <- concatMapM (checkExpr . snd) branches
            return (eResults ++ maybeCons result branchResults)
            where ps = map fst branches

        General.Data _ctor es  -> concatMapM checkExpr es
        General.Access e _f    -> checkExpr e
        General.Remove e _f    -> checkExpr e
        General.Insert e _f e' -> (++) <$> checkExpr e <*> checkExpr e'
        General.Modify e _     -> checkExpr e
        General.Record fields  -> concatMapM (checkExpr . snd) fields
        General.Port _pi       -> return []
        _                      -> return []

    ctorToAdt :: Var.Canonical -> Maybe (Module.AdtInfo String)
    ctorToAdt = \(Var.Canonical {Var.home,Var.name}) -> Map.lookup name =<< Map.lookup home m where
        m :: Map.Map Var.Home (Map.Map String {- ctor name -} (Module.AdtInfo String))
        m = builtins `Map.union` fromModules where
          builtins = Map.empty -- TODO

          fromModules = Map.mapKeysMonotonic Var.Module (Map.map Module.iAdts interfaces)

    ctorToAdtErr :: Region.Region -> Var.Canonical -> M (Module.AdtInfo String)
    ctorToAdtErr region c = case ctorToAdt c of
        Nothing  -> throwError (Annotation.A region (Error.ConstructorTypeNotFound c))
        Just adt -> return adt

    complement :: Region.Region -> CanonicalPat -> M [CanonicalPat]
    complement region p = case p of
        Data ctor ps -> do
            let ctorName = Var.name ctor
                ctorHome = Var.home ctor
            (_tyVars, ctors) <- ctorToAdtErr region ctor
            -- A pattern is in the complement of this one if it uses one of the
            -- other constructors of this type...
            let otherCtorPatterns =
                    map (\(ctor',args) ->
                        Data (Var.Canonical {Var.name = ctor', Var.home = ctorHome})
                            (replicate (length args) Anything))
                    $ filter (\(ctor', _) -> ctor' /= ctorName) ctors
            -- or if it uses this constructor and one of its argument patterns is
            -- in the complement of an argument pattern here.
                numArgs        = length ps
                sandwich pat i = replicate i Anything ++ pat : replicate (numArgs - i - 1) Anything

            argComplement <- fmap concat . forM (zip [0..] ps) $ \(i, argP) ->
                map (\argP' -> Data ctor (sandwich argP' i)) <$> (complement region argP)

            return (otherCtorPatterns ++ argComplement)

        Record _fs     -> return []
        Alias _s p     -> complement region p
        Var _s         -> return []
        Anything       -> return []
        Literal l      -> return [AnythingBut (Set.singleton l)]
        AnythingBut ls -> return (map Literal (Set.toList ls))

    intersection :: CanonicalPat -> CanonicalPat -> Maybe CanonicalPat
    intersection p1 p2 = case (p1, p2) of
        (Alias _ p1', _) -> intersection p1' p2
        (_, Alias _ p2') -> intersection p1 p2'
        (Anything, _)    -> Just p2
        (Var _, _)       -> Just p2
        (_, Anything)    -> Just p1
        (_, Var _)       -> Just p1

        (Data ctor1 args1, Data ctor2 args2) ->
            if ctor1 /= ctor2
            then Nothing
            -- TODO: Assumes that args must be the same length. I think that's ok.
            else fmap (Data ctor1) (sequence (zipWith intersection args1 args2))

        -- TODO Assuming the fields are identical.
        (Record _fs, Record _)             -> Just p1
        (Literal l1, Literal l2)           -> if l1 == l2 then Just p1 else Nothing
        (AnythingBut ls, Literal l)        -> if l `Set.member` ls then Nothing else Just p2
        (Literal l, AnythingBut ls)        -> if l `Set.member` ls then Nothing else Just p1
        (AnythingBut ls1, AnythingBut ls2) -> Just (AnythingBut (Set.union ls1 ls2))
        _                                  -> Nothing

    checkMatch :: Region.Region -> [CanonicalPattern] -> M (Maybe (Annotation.Located Warning))
    checkMatch region = go [Anything] where
        go :: [CanonicalPat] -> [CanonicalPattern] -> M (Maybe (Annotation.Located Warning))
        go []       []      = return Nothing
        go qs@(_:_) []      = return (Just (Annotation.A region (Inexhaustive qs)))
        go unhandled (p:ps) =
            let pp                    = fromPattern p
                Annotation.A region _ = p
            in
            if all (\q -> intersection q pp == Nothing) unhandled
            then return (Just (Annotation.A region (Redundant p)))
            else do
                pComp <- complement region pp
                go (concatMap (\q -> mapMaybe (intersection q) pComp) unhandled) ps