module MiniC.Core
  ( Term(..)
  , CType(..)
  , Statement(..)
  , Expr(..)
  ) where

import Control.Monad (void, unless)
import Control.Monad.Catch (Exception, MonadThrow(throwM))
import Data.Foldable (for_)
import Data.Text (Text)

import Momo.CoreSyntax (CoreSyntax(..))
import Momo.CoreTyping (CoreTyping(..))
import Momo.ModSyntax qualified as Mod
import Momo.Env (Env)
import Momo.Env qualified as Env
import Momo.Ident (Ident)
import Momo.Path (Path)
import Momo.Subst (Subst)
import Momo.Subst qualified as Subst

data Term
  = VarDecl CType
  | FunDef [(Ident, CType)] CType Statement
  deriving (Eq, Ord, Show)

data CType
  = CVoid
  | CInt
  | CFloat
  | CPointer CType
  | CFunction [CType] CType
  | CTypename Path
  deriving (Eq, Ord, Show)

data Statement
  = Expr Expr
  | If Expr Statement Statement
  | For Expr Expr Expr Statement
  | Return Expr
  | Block [(Ident, CType)] [Statement]
  deriving (Eq, Ord, Show)

data Expr
  = IntConst Int
  | FloatConst Float
  | Variable Path
  | Apply Expr [Expr]
  | Assign Expr Expr
  | UnaryOp Text Expr
  | BinaryOp Text Expr Expr
  | Cast Expr CType
  deriving (Eq, Ord, Show)

instance CoreSyntax Term where
  type Val Term = CType
  substVal = flip substType

  type Def Term = CType
  substDef = flip substType

  type Kind Term = ()
  substKind () _subst = ()

substType :: Subst -> CType -> CType
substType subst = \case
  CTypename path ->
    CTypename (Subst.path path subst)

  CPointer cty ->
    CPointer (substType subst cty)

  CFunction args res ->
    CFunction
      (map (substType subst) args)
      (substType subst res)

  cty ->
    cty

data TypeError
  = TypeMismatch { expected :: CType, got :: CType }
  | OpTypeMismatch { op :: Text }
  deriving (Eq, Ord, Show)

instance Exception TypeError

data ExprError
  = VariableNotFound { path :: Path }
  | WrongNumberOfArguments
  | ApplicationOfANonFunction
  | NonLValueAssignment
  deriving (Eq, Ord, Show)

instance Exception ExprError

instance CoreTyping Term where
  typeTerm env = \case
    VarDecl ty -> do
      checkValType env ty
      pure ty

    FunDef params tyRes body -> do
      checkValType env tyRes
      env' <- addVariables env params
      checkStatement env' tyRes body
      pure $ CFunction (map snd params) tyRes

  kindDefType env def =
    void $ checkValType env def

  checkValType env = \case
    CTypename path ->
      case Env.findType path env of
        Nothing ->
          error "boo!"
        Just _found ->
          pure ()

    CPointer cty ->
      checkValType env cty

    CFunction args res -> do
      for_ args $
        checkValType env

      checkValType env res

    _noTypes ->
      pure ()

  checkKind _env () =
    pure ()

  valTypeMatch env cty1 cty2 =
    case (cty1, cty2) of
      (CVoid, CVoid) ->
        True

      (CInt, CInt) ->
        True

      (CFloat, CFloat) ->
        True

      (CFunction args1 res1, CFunction args2 res2)
        | length args1 == length args2 ->
            and $
              zipWith
                (valTypeMatch env)
                (res1 : args1)
                (res2 : args2)

      (CTypename path1, CTypename path2)
        | path1 == path2 ->
            True
        | otherwise ->
            let
              -- XXX: conflating non-manifest types with not-found paths, boo
              manifest1 = Env.findType path1 env >>= Mod.manifest
              manifest2 = Env.findType path2 env >>= Mod.manifest
            in
              case (manifest1, manifest2) of
                (Just def, _) ->
                  valTypeMatch env def cty2
                (_, Just def) ->
                  valTypeMatch env cty1 def
                _noMatch ->
                  False

      (CTypename path1, _) ->
        let
          -- XXX: boo
          manifest = Env.findType path1 env >>= Mod.manifest
        in
          case manifest of
            Nothing ->
              False
            Just def ->
              valTypeMatch env def cty2

      (_, CTypename path2) ->
        let
          -- XXX: boo
          manifest = Env.findType path2 env >>= Mod.manifest
        in
          case manifest of
            Nothing ->
              False
            Just def ->
              valTypeMatch env cty1 def

      _noMatch ->
        False

  defTypeEquiv env () =
    valTypeMatch env

  kindMatch _env () () =
    True

  deftypeOfPath path _kind =
    CTypename path

checkStatement
  :: MonadThrow m
  => Env Term
  -> CType
  -> Statement
  -> m ()
checkStatement env tyRet = \case
  Expr e ->
    void $ typeExpr env e

  If cond ifso ifnot -> do
    checkTypeExpr env cond CInt
    checkStatement env tyRet ifso
    checkStatement env tyRet ifnot

  For initExpr test incr body -> do
    void $ typeExpr env initExpr
    checkTypeExpr env test CInt
    void $ typeExpr env incr
    checkStatement env tyRet body

  Return e ->
    checkTypeExpr env e tyRet

  Block decls stmts -> do
    newEnv <- addVariables env decls
    for_ stmts $
      checkStatement newEnv tyRet

typeExpr
  :: MonadThrow m
  => Env Term
  -> Expr
  -> m CType
typeExpr env = \case
  IntConst{} ->
    pure CInt

  FloatConst{} ->
    pure CFloat

  Variable path ->
    case Env.findValue path env of
      Nothing ->
        throwM VariableNotFound{path}
      Just cty ->
        pure cty

  Apply funct args ->
    typeExpr env funct >>= \case
      CFunction tyArgs tyRes -> do
        unless (length args == length tyArgs) $
          throwM WrongNumberOfArguments
        for_ (zip args tyArgs) $
          uncurry $ checkTypeExpr env
        pure tyRes
      _ ->
        throwM ApplicationOfANonFunction

  Assign arg1 arg2 -> do
    case arg1 of
      Variable{} ->
        pure ()
      UnaryOp "*" _arg ->
        pure ()
      _ ->
        throwM NonLValueAssignment
    cty <- typeExpr env arg1
    checkTypeExpr env arg2 cty
    pure cty

  UnaryOp op arg -> do
    argType <- typeExpr env arg
    case (op, scrape env argType) of
      ("*", CPointer cty) ->
        pure cty
      ("-", CInt) ->
        pure CInt
      ("-", CFloat) ->
        pure CFloat
      ("!", CInt) ->
        pure CInt
      (_op, _type) ->
        throwM OpTypeMismatch{op}

  BinaryOp op arg1 arg2 -> do
    argType1 <- typeExpr env arg1
    argType2 <- typeExpr env arg2
    if
      | op `elem` ["+", "-"] ->
          case (argType1, argType2) of
            (CPointer _pty, CInt) ->
              pure argType1
            (CInt, CInt) ->
              pure argType1
            (CInt, CFloat) ->
              pure argType2
            (CFloat, CInt) ->
              pure argType1
            (CFloat, CFloat) ->
              pure argType1
            _types ->
              throwM OpTypeMismatch{op}
      | op `elem` ["*", "/"] ->
          case (argType1, argType2) of
            (CInt, CInt) ->
              pure argType1
            (CInt, CFloat) ->
              pure argType2
            (CFloat, CInt) ->
              pure argType1
            (CFloat, CFloat) ->
              pure argType1
            _types ->
              throwM OpTypeMismatch{op}
      | op `elem` ["==", "!=", "<", "<=", ">", ">="] ->
          case (argType1, argType2) of
            (CPointer pty1, CPointer pty2)
              | valTypeMatch env pty1 pty2 ->
                  pure CInt
            (CInt, CInt) ->
              pure argType1
            (CInt, CFloat) ->
              pure argType1
            (CFloat, CInt) ->
              pure argType2
            (CFloat, CFloat) ->
              pure CInt
            _types ->
              throwM OpTypeMismatch{op}
      | True ->
          error $ "Unknown op: " <> show op

  Cast expr cty -> do
    void $ typeExpr env expr
    checkValType env cty
    pure cty

checkTypeExpr
  :: MonadThrow m
  => Env Term
  -> Expr
  -> CType
  -> m ()
checkTypeExpr env expr expectedType = do
  exprType <- typeExpr env expr
  case (scrape env expectedType, scrape env exprType) of
    -- XXX: happy debugging
    (CInt, CFloat) ->
      pure ()
    (CFloat, CInt) ->
      pure ()
    (CPointer CVoid, CPointer _cty) ->
      pure ()
    (CPointer _cty, CPointer CVoid) ->
      pure ()
    (cty1, cty2) ->
      unless (cty1 == cty2) $
        throwM TypeMismatch
          { expected = cty1
          , got      = cty2
          }

scrape :: Env Term -> CType -> CType
scrape env cty =
  case cty of
    CTypename path ->
      case Env.findType path env of
        Just found ->
          case found.manifest of
            Just ty' ->
              scrape env ty'
            Nothing ->
              cty

        _ignore ->
          cty

    _ignore ->
      cty

addVariables
  :: MonadThrow m
  => Env Term
  -> [(Ident, CType)]
  -> m (Env Term)
addVariables env = \case
  [] ->
    pure env
  (ident, ty) : next -> do
    checkValType env ty
    addVariables (Env.addValue ident ty env) next