{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE IncoherentInstances #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE UndecidableInstances #-}

module Language.Fiddle.Compiler.ConsistencyCheck (consistencyCheckPhase) where

import Control.Monad (forM_, unless, when)
import Control.Monad.RWS (MonadWriter (tell))
import Control.Monad.Trans.Writer (Writer, execWriter)
import Data.Foldable (foldlM, toList)
import Data.Functor.Identity
import qualified Data.IntMap as IntMap
import Data.List (intercalate)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Typeable
import GHC.TypeError as TypeError
import Language.Fiddle.Ast
import Language.Fiddle.Compiler
import Language.Fiddle.Internal.UnitInterface as UnitInterface
import Language.Fiddle.Internal.UnitNumbers
import Language.Fiddle.Types
import Text.Printf (printf)
import Prelude hiding (unzip)

type S = Qualified

type S' = Checked

type F = Identity

type A = Commented SourceSpan

type M = Compile ()

pattern QMdP :: t -> Identity (When True t)
pattern QMdP t = Identity (Present t)

instance CompilationStage Checked where
  type StageAfter Checked = TypeError (TypeError.Text "No stage after Checked")
  type StageMonad Checked = M
  type StageState Checked = ()
  type StageFunctor Checked = Identity
  type StageAnnotation Checked = A

instance CompilationStage S where
  type StageAfter S = S'
  type StageMonad S = M
  type StageState S = ()
  type StageFunctor S = F
  type StageAnnotation S = A

consistencyCheckPhase :: CompilationPhase S S'
consistencyCheckPhase = pureCompilationPhase $ advanceStage ()

instance AdvanceStage S ObjTypeBody where
  advanceStage () objTypeBody = snd <$> advanceObjTypeBody objTypeBody 0

deriving instance AdvanceStage S AnonymousBitsType

deriving instance AdvanceStage S ImportStatement

instance AdvanceStage S BitType where
  customAdvanceStage t _ = do
    case t of
      (EnumBitType sz (Identity body) _) -> do
        checkEnumConsistency sz body
      _ -> return ()
    return Nothing

deriving instance AdvanceStage S EnumBody

deriving instance AdvanceStage S EnumConstantDecl

deriving instance AdvanceStage S PackageBody

deriving instance AdvanceStage S FiddleDecl

deriving instance AdvanceStage S (ConstExpression u)

instance AdvanceStage S FiddleUnit where
  advanceStage () fu@(FiddleUnit _ decls a) =
    FiddleUnit (Present $ getUnitInterface fu) <$> mapM (advanceStage ()) decls <*> pure a
    where
      getUnitInterface = execWriter . walk_ doWalk

      doWalk :: forall t'. (Walk t', Typeable t') => t' F A -> Writer UnitInterface ()
      doWalk t =
        case () of
          ()
            | (Just (PackageDecl {packageQualificationMetadata = (QMdP d)})) <-
                castTS t ->
                tell (UnitInterface.singleton d)
            | (Just (LocationDecl {locationQualificationMetadata = (QMdP d)})) <-
                castTS t ->
                tell (UnitInterface.singleton d)
            | (Just (BitsDecl {bitsQualificationMetadata = (QMdP d)})) <-
                castTS t ->
                tell (UnitInterface.singleton d)
            | (Just (ObjTypeDecl {objTypeQualificationMetadata = (QMdP d)})) <-
                castTS t ->
                tell (UnitInterface.singleton d)
            | (Just (ObjectDecl {objectQualificationMetadata = (QMdP d)})) <-
                castTS t ->
                tell (UnitInterface.singleton d)
            | (Just (ImportStatement {importInterface = ii})) <-
                castTS t ->
                tell (UnitInterface mempty (dependencies (unwrap ii)))
          _ -> return ()

      castTS ::
        (Typeable t', Typeable t, Typeable f, Typeable a) =>
        t' f a ->
        Maybe (t S f a)
      castTS = cast

deriving instance AdvanceStage S (Expression u)

deriving instance AdvanceStage S RegisterBitsTypeRef

deriving instance AdvanceStage S ObjType

deriving instance (AdvanceStage S t) => AdvanceStage S (Directed t)

advanceObjTypeBody :: ObjTypeBody S F A -> N Bytes -> M (N Bytes, ObjTypeBody S' F A)
advanceObjTypeBody (ObjTypeBody us decls a) startOffset = do
  (decls', _) <- advanceDecls

  calcSize <- case us of
    Union {} -> do
      checkJagged decls'
      return $ maximum (map fst decls')
    Struct {} -> return $ sum (map fst decls')

  return (calcSize, ObjTypeBody us (reverse $ map snd decls') a)
  where
    advanceDecls :: M ([(N Bytes, Directed ObjTypeDecl S' F A)], N Bytes)
    advanceDecls = do
      foldlM
        ( \(ret, offset) d ->
            let advanceOffset :: N Bytes -> N Bytes -> N Bytes
                advanceOffset = case us of
                  Union {} -> const
                  Struct {} -> (+)
                doReturn ::
                  (Monad m) =>
                  ObjTypeDecl S' F A ->
                  N Bytes ->
                  m ([(N Bytes, Directed ObjTypeDecl S' F A)], N Bytes)
                doReturn x size = return ((size, mapDirected (const x) d) : ret, advanceOffset offset size)
             in case undirected d of
                  e@AssertPosStatement {assertExpr = expr} -> do
                    assertedPos <- expressionToIntM expr
                    checkPositionAssertion (annot e) assertedPos offset
                    return (ret, offset)
                  (RegisterDecl qmeta mod ident size Nothing a) -> do
                    let declaredSize = regSzToBits (getLeft size)
                    reifiedSizeBytes <- checkBitsSizeMod8 a declaredSize

                    let span = Present (FieldSpan offset reifiedSizeBytes)
                        qmeta' = fmap (\q -> q {regSpan = span}) qmeta
                    doReturn (RegisterDecl qmeta' mod ident (changeRight size) Nothing a)
                      =<< checkBitsSizeMod8 a declaredSize
                  (RegisterDecl qmeta mod ident size (Just body) a) -> do
                    let declaredSize = regSzToBits (getLeft size)

                    (actualSize, body') <- advanceRegisterBody 0 body
                    checkSizeMismatch a declaredSize actualSize
                    reifiedSizeBytes <- checkBitsSizeMod8 a declaredSize

                    let span = Present (FieldSpan offset reifiedSizeBytes)
                        qmeta' = fmap (\q -> q {regSpan = span}) qmeta

                    doReturn
                      ( RegisterDecl
                          qmeta'
                          mod
                          ident
                          (changeRight size)
                          (Just body')
                          a
                      )
                      reifiedSizeBytes
                  (TypeSubStructure (Identity body) name a) -> do
                    (size, body') <- advanceObjTypeBody body offset
                    doReturn (TypeSubStructure (Identity body') name a) size
        )
        (([], startOffset) :: ([(N Bytes, Directed ObjTypeDecl S' F A)], N Bytes))
        decls

    advanceAndGetSize :: Expression u S F A -> M (Expression u S' F A, N u)
    advanceAndGetSize e = (,) <$> advanceStage () e <*> expressionToIntM e

pattern RegisterBodyPattern :: BodyType F A -> [Directed RegisterBitsDecl s F A] -> A -> A -> RegisterBody s F A
pattern RegisterBodyPattern u decls a b = RegisterBody u (Identity (DeferredRegisterBody decls b)) a

-- registerBodyPattern u decls a b = RegisterBody u (Identity (DeferredRegisterBody decls a)) a

advanceRegisterBody :: N Bits -> RegisterBody S F A -> M (N Bits, RegisterBody S' F A)
-- Handle the case where it's a union.
advanceRegisterBody
  startOffset
  (RegisterBodyPattern us (NonEmpty.nonEmpty -> Just decls) a b) = do
    (structSize, reverse -> decls') <-
      foldlM
        ( \(offset, ret) d -> do
            (sz, t) <- advanceDecl offset (undirected d)
            let advanceOffset off sz =
                  case us of
                    Union {} -> off
                    Struct {} -> off + sz
            return (advanceOffset offset sz, (sz, mapDirected (const t) d) : ret)
        )
        (startOffset, [])
        decls
    calcSize <- case us of
      Union {} -> do
        checkJagged decls'
        return $ maximum (map fst decls')
      Struct {} -> return structSize

    return (calcSize, RegisterBodyPattern us (map snd $ toList decls') a b)

-- Handle the case where there's no decls.
advanceRegisterBody _ (RegisterBodyPattern u _ a b) =
  return (0, RegisterBodyPattern u [] a b)
advanceRegisterBody _ RegisterBody {} = error "GHC not smart enuf"

checkJagged :: (Annotated t) => [(N u, t f A)] -> Compile s ()
checkJagged decls = do
  let expectedSize = maximum (fmap fst decls)
  forM_ decls $ \(sz, annot -> a) ->
    when (sz /= expectedSize) $
      emitDiagnosticWarning
        ( printf
            "[JaggedUnion] - All elements of a union should be the same size. \
            \ this element is size %d, expected size %d. Maybe bundle this with \
            \ reserved(%d)?"
            sz
            expectedSize
            (expectedSize - sz)
        )
        a

advanceDecl :: N Bits -> RegisterBitsDecl S F A -> M (N Bits, RegisterBitsDecl S' F A)
advanceDecl offset = \case
  ReservedBits expr an -> do
    sz <- expressionToIntM expr
    (sz,)
      <$> ( ReservedBits
              <$> advanceStage () expr
              <*> pure an
          )
  DefinedBits qmeta mod ident typ annot -> do
    size <- bitsTypeSize typ
    let span = Present (FieldSpan offset size)
        qmeta' = fmap (\q -> q {bitsSpan = span}) qmeta

    (size,)
      <$> (DefinedBits qmeta' mod ident <$> advanceStage () typ <*> pure annot)
  BitsSubStructure subBody subName ann -> do
    (sz, body') <- advanceRegisterBody offset subBody
    return (sz, BitsSubStructure body' subName ann)

bitsTypeSize :: RegisterBitsTypeRef S F A -> M (N Bits)
bitsTypeSize (RegisterBitsArray tr nExpr _) = do
  sz <- bitsTypeSize tr
  return (sz .*. trueValue nExpr)
bitsTypeSize
  RegisterBitsReference
    { bitsRefQualificationMetadata =
        QMdP (ExportedBitsDecl {exportedBitsDeclSizeBits = sz})
    } = return sz
bitsTypeSize (RegisterBitsReference {}) = error "should be exhaustive"
bitsTypeSize (RegisterBitsJustBits expr _) = return $ trueValue expr

checkSizeMismatch :: (NamedUnit u) => A -> N u -> N u -> Compile s ()
checkSizeMismatch _ a b | a == b = return ()
checkSizeMismatch pos declaredSize calculatedSize =
  emitDiagnosticError
    ( printf
        "Size assertion failed. Declared size %s, calculated %s"
        (unitName declaredSize)
        (unitName calculatedSize)
    )
    pos

checkPositionAssertion :: A -> N u -> N u -> Compile s ()
checkPositionAssertion _ a b | a == b = return ()
checkPositionAssertion pos declaredPosition calculatedPostion =
  emitDiagnosticError
    ( printf
        "Position assertion failed. Asserted 0x%x, calculated 0x%x"
        declaredPosition
        calculatedPostion
    )
    pos

expressionToIntM ::
  (stage .< Expanded ~ False) =>
  Expression u stage f A ->
  Compile s (N u)
expressionToIntM expr =
  resolveOrFail $
    either
      ( \reason -> Left [Diagnostic Error reason (unCommented $ annot expr)]
      )
      return
      (expressionToInt expr)

checkBitsSizeMod8 :: A -> N Bits -> M (N Bytes)
checkBitsSizeMod8 a w = do
  let (x, rem) = bitsToBytes w
  when (rem /= 0) $
    emitDiagnosticError
      (printf "Register size %d is not a multiple of 8. Please add padding to this register." w)
      a
  return x

checkEnumConsistency :: Expression Bits S F A -> EnumBody S F A -> M ()
checkEnumConsistency expr enumBody@(EnumBody {enumConsts = constants}) = do
  declaredSize <- expressionToIntM expr

  -- If the declared size is less than or equal to 4, we'll enforce that the
  -- enum is packed. This is to make sure the user has covered all bases.
  when (declaredSize <= 4) $ do
    imap <-
      foldlM
        ( \imap (undirected -> enumConst) -> do
            number <- case enumConst of
              EnumConstantDecl _ expr _ -> return $ trueValue expr
              EnumConstantReserved expr _ -> expressionToIntM expr

            when (number >= 2 ^ declaredSize) $
              emitDiagnosticError
                ( printf
                    "Enum constant too large. Max allowed %d\n"
                    ((2 :: Int) ^ declaredSize)
                )
                (annot enumConst)

            return $ IntMap.insert (fromIntegral number) True imap
        )
        IntMap.empty
        constants

    let missing =
          filter (not . (`IntMap.member` imap)) [0 .. 2 ^ declaredSize - 1]

    unless (null missing) $
      emitDiagnosticWarning
        ( printf
            "Missing enum constants %s. Small enums should be fully \
            \ populated. Use 'reserved' if needed."
            (intercalate ", " (map show missing))
        )
        (annot enumBody)