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

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

import Control.Monad (forM, forM_, unless, when)
import Control.Monad.Identity (Identity (Identity))
import Control.Monad.RWS (MonadState (get, put), MonadWriter (tell), gets, modify')
import Data.Foldable (Foldable (toList), foldlM)
import Data.Functor.Identity
import qualified Data.IntMap as IntMap
import Data.Kind (Type)
import Data.List (inits, intercalate)
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NonEmpty
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, mapMaybe)
import qualified Data.Set as Set
import qualified Data.Text as Text
import Data.Void
import Data.Word (Word32)
import GHC.TypeError as TypeError
import GHC.TypeLits
import Language.Fiddle.Ast
import Language.Fiddle.Compiler
import Language.Fiddle.Internal.Scopes
import Language.Fiddle.Internal.UnitInterface as UnitInterface
import Language.Fiddle.Types (Comment (DocComment), Commented (Commented, unCommented), SourceSpan)
import Text.Printf (printf)
import Prelude hiding (unzip)

data GlobalState = GlobalState
  { globalScope :: Scope String (Either SizeBits SizeBytes),
    fileDependencies :: [FilePath],
    unitInterface :: UnitInterface
  }

newtype LocalState = LocalState (ScopePath String)

type CurrentStage = Qualified

type I = Identity

type Annot = Commented SourceSpan

type SizeBits = Word32

type SizeBytes = Word32

consistencyCheckPhase :: CompilationPhase CurrentStage Checked
consistencyCheckPhase = pureCompilationPhase checkConsistency

checkConsistency ::
  FiddleUnit CurrentStage I Annot ->
  Compile () (FiddleUnit Checked I Annot)
checkConsistency =
  fmap snd
    . subCompile (GlobalState mempty mempty mempty)
    . advanceStage (LocalState mempty)

instance CompilationStage Checked where
  type StageAfter Checked = TypeError (TypeError.Text "No stage after Checked")
  type StageMonad Checked = Compile GlobalState
  type StageState Checked = LocalState
  type StageFunctor Checked = Identity
  type StageAnnotation Checked = Commented SourceSpan

instance CompilationStage CurrentStage where
  type StageAfter CurrentStage = Checked
  type StageMonad CurrentStage = Compile GlobalState
  type StageState CurrentStage = LocalState
  type StageFunctor CurrentStage = Identity
  type StageAnnotation CurrentStage = Commented SourceSpan

instance AdvanceStage CurrentStage FiddleUnit where
  advanceStage localState (FiddleUnit _ decls a) = do
    decls' <- mapM (advanceStage localState) decls
    intf <- gets unitInterface
    return $ FiddleUnit intf decls' a

-- advanceStage localState (FiddleUnit decls _ annot) = do

--   decls' <- mapM (advanceStage localState) decls

deriving instance AdvanceStage CurrentStage Expression

deriving instance AdvanceStage CurrentStage ObjType

deriving instance AdvanceStage CurrentStage DeferredRegisterBody

deriving instance AdvanceStage CurrentStage RegisterBitsDecl

deriving instance AdvanceStage CurrentStage RegisterBitsTypeRef

deriving instance AdvanceStage CurrentStage AnonymousBitsType

deriving instance AdvanceStage CurrentStage BitType

deriving instance AdvanceStage CurrentStage EnumBody

deriving instance AdvanceStage CurrentStage EnumConstantDecl

deriving instance AdvanceStage CurrentStage PackageBody

instance AdvanceStage CurrentStage ImportStatement where
  modifyState
    ( ImportStatement
        { importInterface =
            ( UnitInterface
                { rootScope = unitScope,
                  dependencies = importDependencies
                }
              )
        }
      )
    ls = do
      modify'
        ( \s@GlobalState
             { globalScope = globalScope,
               unitInterface = unitInterface
             } ->
              s
                { globalScope =
                    unitInterfaceScopeToGlobalScope unitScope <> globalScope,
                  unitInterface =
                    unitInterface
                      { dependencies =
                          importDependencies ++ dependencies unitInterface
                      }
                }
        )
      return ls
      where
        unitInterfaceScopeToGlobalScope =
          fmap
            ( \(_, exportedValue) -> case exportedValue of
                ExportedBitsType sz -> Left sz
                ExportedObjType sz -> Right sz
            )

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

instance AdvanceStage CurrentStage RegisterBody where
  advanceStage s body = fst <$> registerBodyToStage3 s body

instance AdvanceStage CurrentStage ObjTypeBody where
  advanceStage s body = fst <$> objTypeBodyToStage3 s body 0

deriving instance AdvanceStage CurrentStage FiddleDecl

instance AdvanceStage CurrentStage (Directed FiddleDecl) where
  modifyState (Directed directives t _) s = case t of
    (BitsDecl _ id typ annotation) -> do
      typeSize <- getTypeSize typ
      insertTypeSize annotation s id typeSize
      return s
    (PackageDecl _ n _ _) -> do
      let strs = nameToList n
      let (LocalState scopePath) = s

      return $
        LocalState $
          scopePath {currentScope = strs ++ currentScope scopePath}
    (UsingDecl _ n _) ->
      let (LocalState scopePath) = s
       in return $
            LocalState $
              scopePath
                { usingPaths = nameToList n : usingPaths scopePath
                }
    _ -> return s

  customAdvanceStage (Directed directives t a) (LocalState scopePath) = case t of
    (ObjTypeDecl q ident (Identity body) annot) -> do
      (body', size) <- objTypeBodyToStage3 (LocalState scopePath) body 0

      let fullName =
            NonEmpty.prependList
              (currentScope scopePath)
              (NonEmpty.singleton (Text.unpack (identifierName ident)))

      ui <- gets unitInterface
      let ui' = insertIntoUnitInterface fullName ui annot (ExportedObjType size)
      modify' $ \gs -> gs {unitInterface = ui'}

      return $ Just $ Directed directives (ObjTypeDecl q ident (Identity body') annot) a
    _ -> return Nothing

objTypeBodyToStage3 ::
  LocalState ->
  ObjTypeBody CurrentStage I Annot ->
  Word32 ->
  Compile GlobalState (ObjTypeBody Checked I Annot, Word32)
objTypeBodyToStage3 st (ObjTypeBody bodyType decls a) startOff = do
  let isUnion = case bodyType of
        Union {} -> True
        _ -> False
  (cur, returned) <-
    foldlM
      ( \(cursor, returned) decl ->
          case undirected decl of
            RegisterDecl mMod mIdent expr mBody a -> do
              (s3RegisterBody, mCalculatedSize) <-
                fUnzip <$> mapM (registerBodyToStage3 st) mBody

              nExpr <- advanceStage st expr

              let s3 =
                    mapDirected
                      ( const $
                          RegisterDecl
                            mMod
                            mIdent
                            nExpr
                            s3RegisterBody
                            a
                      )
                      decl

              declaredSizeBits <- fromIntegral <$> exprToSize expr

              when ((declaredSizeBits `mod` 8) /= 0) $
                tell
                  [ Diagnostic
                      Error
                      "Register size is not a multiple of 8. Please pad register size to align with 8. "
                      (unCommented a)
                  ]

              forM_ mCalculatedSize $ \(fromIntegral -> calculatedSize) ->
                unless (calculatedSize == declaredSizeBits) $
                  let helpful =
                        if calculatedSize < declaredSizeBits
                          then
                            printf
                              "\nPerhaps you should add 'reserved(%d)' to the end of your register declaration?"
                              (declaredSizeBits - calculatedSize)
                          else ""
                   in tell
                        [ Diagnostic
                            Error
                            ( printf
                                "Calculated size %d does not match declared size %d.%s"
                                calculatedSize
                                declaredSizeBits
                                helpful
                            )
                            (unCommented a)
                        ]

              if isUnion
                then
                  checkUnion cursor (declaredSizeBits `div` 8) (s3 : returned) a
                else
                  return (cursor + declaredSizeBits `div` 8, s3 : returned)
            TypeSubStructure (Identity subBody) maybeIdent annot -> do
              (newBody, size) <-
                objTypeBodyToStage3
                  st
                  subBody
                  ( if isUnion then startOff else cursor
                  )
              let s3 = mapDirected (const $ TypeSubStructure (Identity newBody) maybeIdent annot) decl

              checkTypesSubStructure subBody maybeIdent annot
              if isUnion
                then
                  checkUnion cursor size (s3 : returned) a
                else
                  return (cursor + size, s3 : returned)
            ReservedDecl expr annot -> do
              size' <- fromIntegral <$> exprToSize expr
              when ((size' `mod` 8) /= 0) $
                tell
                  [ Diagnostic
                      Error
                      "Can only reserve a multiple of 8 bits in this context."
                      (unCommented a)
                  ]

              expr' <- advanceStage st expr
              let size = size' `div` 8
              let s3 = mapDirected (const $ ReservedDecl expr' annot) decl
              if isUnion
                then
                  checkUnion cursor size (s3 : returned) a
                else
                  return (cursor + size, s3 : returned)
            AssertPosStatement _ expr a -> do
              declaredPos <- fromIntegral <$> exprToSize expr

              let expectedPos = if isUnion then startOff else cursor + startOff

              when (expectedPos /= declaredPos) $ do
                tell
                  [ Diagnostic
                      Error
                      ( printf
                          "Position assertion failed. Asserted 0x%x, calculated 0x%x"
                          declaredPos
                          expectedPos
                      )
                      (unCommented a)
                  ]
              return (cursor, returned)
      )
      (0, [])
      decls

  return (ObjTypeBody bodyType (reverse returned) a, cur)
  where
    checkTypesSubStructure
      (ObjTypeBody bodyType decls _)
      maybeIdent
      annot =
        let emitWarning s = tell [Diagnostic Warning s (unCommented annot)]
         in case () of
              ()
                | [_] <- decls,
                  (Union {}) <- bodyType ->
                    emitWarning "Union with a single field. Should this be a struct?"
              ()
                | [_] <- decls,
                  (Struct {}) <- bodyType,
                  Nothing <- maybeIdent ->
                    emitWarning "Anonymous sub-struct with single field is superfluous."
              ()
                | [] <- decls ->
                    emitWarning
                      ( printf
                          "Empty sub-%s is superfluous."
                          ( case bodyType of
                              Union {} -> "union"
                              Struct {} -> "struct"
                          )
                      )
              _ -> return ()
    fUnzip xs = (fst <$> xs, snd <$> xs)
    pushApply :: Maybe (a, b) -> (Maybe a, Maybe b)
    pushApply (Just (a, b)) = (Just a, Just b)
    pushApply Nothing = (Nothing, Nothing)

registerBodyToStage3 ::
  LocalState ->
  RegisterBody CurrentStage I Annot ->
  Compile GlobalState (RegisterBody Checked I Annot, Word32)
registerBodyToStage3
  st
  (RegisterBody bodyType (Identity deferredRegisterBody) a') = do
    let isUnion = case bodyType of
          Union {} -> True
          _ -> False

    case deferredRegisterBody of
      DeferredRegisterBody decls a -> do
        (cur, returned) <-
          foldlM
            ( \(cursor, returned) decl ->
                case undirected decl of
                  ReservedBits expr a -> do
                    size <- fromIntegral <$> exprToSize expr
                    expr' <- advanceStage st expr
                    let s3 =
                          mapDirected
                            (const $ ReservedBits expr' a)
                            decl
                    if isUnion
                      then checkUnion cursor size (s3 : returned) a
                      else
                        return (cursor + size, s3 : returned)
                  BitsSubStructure registerBody maybeIdent annot -> do
                    checkBitsSubStructure registerBody maybeIdent annot

                    (newBody, subsize) <- registerBodyToStage3 st registerBody
                    let s3 =
                          mapDirected
                            (const $ BitsSubStructure newBody maybeIdent annot)
                            decl

                    if isUnion
                      then checkUnion cursor subsize (s3 : returned) a
                      else
                        return (cursor + subsize, s3 : returned)
                  DefinedBits modifier identifier typeref a -> do
                    (s3TypeRef, size) <- registerBitsTypeRefToStage3 st typeref
                    let s3 =
                          mapDirected
                            (const $ DefinedBits modifier identifier s3TypeRef a)
                            decl

                    if isUnion
                      then checkUnion cursor size (s3 : returned) a
                      else
                        return (cursor + size, s3 : returned)
            )
            (0, [])
            decls

        return
          ( RegisterBody
              bodyType
              (Identity (DeferredRegisterBody (reverse returned) a))
              a',
            cur
          )
    where
      checkBitsSubStructure
        (RegisterBody bodyType (Identity (DeferredRegisterBody decls _)) _)
        maybeIdent
        annot =
          let emitWarning s = tell [Diagnostic Warning s (unCommented annot)]
           in case () of
                ()
                  | [_] <- decls,
                    (Union {}) <- bodyType ->
                      emitWarning "Union with a single field. Should this be a struct?"
                ()
                  | [_] <- decls,
                    (Struct {}) <- bodyType,
                    Nothing <- maybeIdent ->
                      emitWarning "Anonymous sub-struct with single field is superfluous."
                ()
                  | [] <- decls ->
                      emitWarning
                        ( printf
                            "Empty sub-%s is superfluous."
                            ( case bodyType of
                                Union {} -> "union"
                                Struct {} -> "struct"
                            )
                        )
                _ -> return ()

registerBitsTypeRefToStage3 ::
  LocalState ->
  RegisterBitsTypeRef CurrentStage I Annot ->
  Compile GlobalState (RegisterBitsTypeRef Checked I Annot, Word32)
registerBitsTypeRefToStage3 localState = \case
  RegisterBitsArray ref expr a -> do
    (ref', size) <- registerBitsTypeRefToStage3 localState ref
    multiplier <- exprToSize expr
    expr' <- advanceStage localState expr
    return
      ( RegisterBitsArray ref' expr' a,
        size * fromIntegral multiplier
      )
  RegisterBitsReference q name a ->
    (RegisterBitsReference q name a,) <$> lookupTypeSize localState name
  RegisterBitsJustBits expr a -> do
    expr' <- advanceStage localState expr
    (RegisterBitsJustBits expr' a,)
      . fromIntegral
      <$> exprToSize expr

checkUnion :: Word32 -> Word32 -> b -> Commented SourceSpan -> Compile s (Word32, b)
checkUnion cursor subsize ret a = do
  when (cursor /= 0 && subsize /= cursor) $ do
    tell
      [ Diagnostic
          Warning
          ( printf
              "Jagged union found. Found size %d, expected %d.\n \
              \ Please wrap smaller fields in a struct with padding so all \
              \ fields are the same size?"
              subsize
              cursor
          )
          (unCommented a)
      ]
  return (max cursor subsize, ret)

exprToSize ::
  (NumberType stage ~ Integer) =>
  Expression stage I Annot ->
  Compile s Integer
exprToSize (LitNum num _) = return num
exprToSize e = do
  tell [Diagnostic Error "Variables not allowed" (unCommented $ annot e)]
  compilationFailure

lookupTypeSize :: LocalState -> Name I Annot -> Compile GlobalState SizeBits
lookupTypeSize (LocalState scopePath) (Name idents a) = do
  -- Convert the list of identifiers to a string path
  let path = fmap (\(Identifier s _) -> Text.unpack s) idents

  -- Get the current scope and perform the lookup
  results <- gets $ lookupScopeWithPath scopePath path . globalScope

  case results of
    -- Successfully resolved to a unique size
    [(_, Left sz)] -> return sz
    -- Multiple ambiguous results found
    matches@(_ : _ : _) -> do
      -- Generate a list of ambiguous paths for error reporting
      let ambiguousPaths =
            map
              ( \(resolvedPath, _) ->
                  intercalate "." (NonEmpty.toList resolvedPath)
              )
              matches
      tell
        [ Diagnostic
            Error
            ( printf
                "Ambiguous occurrence of '%s'. Multiple matches found:\n%s"
                (intercalate "." $ NonEmpty.toList path)
                (unlines ambiguousPaths) -- List all ambiguous paths
            )
            (unCommented a)
        ]
      compilationFailure

    -- No matches found
    [] -> do
      tell
        [ Diagnostic
            Error
            ( printf
                "Cannot resolve '%s'. No matching symbols found."
                (intercalate "." $ NonEmpty.toList path)
            )
            (unCommented a)
        ]
      compilationFailure

getTypeSize :: BitType CurrentStage I Annot -> Compile s SizeBits
getTypeSize (RawBits expr _) = fromIntegral <$> exprToSize expr
getTypeSize (EnumBitType expr (Identity (EnumBody constants _)) ann) = do
  declaredSize <- fromIntegral <$> exprToSize 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 _ -> exprToSize expr
              EnumConstantReserved expr _ -> exprToSize expr

            when (number >= 2 ^ declaredSize) $
              tell
                [ Diagnostic
                    Error
                    ( printf
                        "Enum constant too large. Max allowed %d\n"
                        ((2 :: Int) ^ declaredSize)
                    )
                    (unCommented (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) $
      tell
        [ Diagnostic
            Warning
            ( printf
                "Missing enum constants %s. Small enums should be fully \
                \ populated. Use 'reserved' if needed."
                (intercalate ", " (map show missing))
            )
            (unCommented ann)
        ]

  return declaredSize

diagnosticError :: String -> Annot -> Compile a ()
diagnosticError str a = tell [Diagnostic Error str (unCommented a)]

diagnosticInfo :: String -> Annot -> Compile a ()
diagnosticInfo str a = tell [Diagnostic Info str (unCommented a)]

insertIntoUnitInterface path ui (Commented comments srcspan) val =
  let docComments =
        mconcat
          ( mapMaybe
              ( \com -> do
                  (DocComment txt) <- Just com
                  return txt
              )
              comments
          )
   in ui
        { rootScope =
            insertScope
              path
              ( Metadata srcspan docComments [],
                val
              )
              (rootScope ui)
        }

insertTypeSize ::
  Annot ->
  LocalState ->
  Identifier f Annot ->
  SizeBits ->
  Compile GlobalState ()
insertTypeSize annot (LocalState scopePath) (Identifier s idannot) size = do
  modifyM $
    \state@GlobalState
       { globalScope = globalScope,
         unitInterface = unitInterface
       } ->
        let fullName =
              NonEmpty.prependList
                (currentScope scopePath)
                (NonEmpty.singleton (Text.unpack s))
         in case upsertScope fullName (Left size) globalScope of
              (Just _, _) -> do
                diagnosticError (printf "Duplicate type %s" s) idannot

                compilationFailure
              (Nothing, n) ->
                let unitInterface' =
                      insertIntoUnitInterface
                        fullName
                        unitInterface
                        annot
                        (ExportedBitsType size)
                 in return $
                      state
                        { globalScope = n,
                          unitInterface = unitInterface'
                        }
  where
    modifyM fn = do
      s <- get
      put =<< fn s