-- | Qualification compilation phase.
--
-- The qualification phase is responsible for resolving all type references in
-- the AST to their fully-qualified counterparts. This process involves
-- replacing unqualified references with their fully-qualified names and
-- attaching the necessary metadata to each reference. This enriched information
-- is then available for use in later stages of the compilation pipeline.
--
-- In this phase, symbol resolution statements (such as 'using' statements) are
-- removed, as they become unnecessary once references are fully qualified.
module Language.Fiddle.Compiler.Qualification (qualificationPhase) where

import Control.Monad.RWS (MonadWriter (tell))
import Control.Monad.State
import Data.Foldable (foldlM)
import Data.List (intercalate)
import Data.List.NonEmpty (NonEmpty (..), toList)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Maybe (mapMaybe)
import Data.Word
import Language.Fiddle.Ast
import Language.Fiddle.Compiler
import Language.Fiddle.Compiler.ConsistencyCheck ()
import Language.Fiddle.Internal.Scopes
import Language.Fiddle.Internal.UnitInterface as UnitInterface
import Language.Fiddle.Types
import Text.Printf (printf)

type S = Expanded

newtype GlobalState = GlobalState
  { unitInterface :: UnitInterface
  }

data LocalState = LocalState
  { currentScopePath :: ScopePath String,
    ephemeralScope :: Scope String (Metadata, ExportedDecl)
  }

type F = Either [Diagnostic]

type A = Commented SourceSpan

type M = Compile GlobalState

instance CompilationStage Expanded where
  type StageAfter Expanded = Qualified
  type StageMonad Expanded = M
  type StageState Expanded = LocalState
  type StageFunctor Expanded = F
  type StageAnnotation Expanded = A

qualificationPhase :: CompilationPhase Expanded Qualified
qualificationPhase =
  pureCompilationPhase $ \t -> do
    raw <-
      fmap snd $
        subCompile (GlobalState mempty) $
          advanceStage
            (LocalState mempty mempty)
            (soakA t)

    squeezeDiagnostics raw

deriving instance AdvanceStage S ObjTypeBody

deriving instance AdvanceStage S DeferredRegisterBody

deriving instance AdvanceStage S RegisterBody

deriving instance AdvanceStage S AnonymousBitsType

deriving instance AdvanceStage S ImportStatement

deriving instance AdvanceStage S BitType

deriving instance AdvanceStage S EnumBody

deriving instance AdvanceStage S EnumConstantDecl

deriving instance AdvanceStage S RegisterBitsDecl

deriving instance AdvanceStage S ObjTypeDecl

deriving instance AdvanceStage S Expression

instance AdvanceStage S RegisterBitsTypeRef where
  advanceStage localState = \case
    RegisterBitsArray a b c ->
      RegisterBitsArray
        <$> advanceStage localState a
        <*> advanceStage localState b
        <*> pure c
    RegisterBitsJustBits a b ->
      RegisterBitsJustBits
        <$> advanceStage localState a
        <*> pure b
    RegisterBitsReference _ name a -> do
      v <- fmap snd <$> resolveName name localState
      return $ RegisterBitsReference v name a

instance AdvanceStage S ObjType where
  advanceStage localState = \case
    ArrayObjType a b c ->
      ArrayObjType
        <$> advanceStage localState a
        <*> advanceStage localState b
        <*> pure c
    ReferencedObjType _ name a -> do
      v <- fmap snd <$> resolveName name localState
      return $ ReferencedObjType v name a

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

instance AdvanceStage S PackageBody where
  advanceStage localState (PackageBody decls a) =
    PackageBody <$> advanceFiddleDecls localState decls <*> pure a

instance AdvanceStage S FiddleUnit where
  advanceStage localState (FiddleUnit () decls a) =
    FiddleUnit () <$> advanceFiddleDecls localState decls <*> pure a

modifyEphemeralScope ::
  ( Scope String (Metadata, ExportedDecl) -> Scope String (Metadata, ExportedDecl)
  ) ->
  LocalState ->
  LocalState
modifyEphemeralScope fn ls@LocalState {ephemeralScope = es} =
  ls {ephemeralScope = fn es}

modifyCurrentScopePath ::
  (ScopePath String -> ScopePath String) ->
  LocalState ->
  LocalState
modifyCurrentScopePath fn ls@LocalState {currentScopePath = cs} =
  ls {currentScopePath = fn cs}

resolveIdent :: (ExportableDecl d, Functor f) => Identifier f A -> LocalState -> M (F ([String], d))
resolveIdent i = resolveSymbol (annot i) [identToString i]

resolveName :: (ExportableDecl d, Functor f) => Name f A -> LocalState -> M (F ([String], d))
resolveName n = resolveSymbol (annot n) (toList $ nameToList n)

resolveSymbol :: (ExportableDecl d) => A -> [String] -> LocalState -> M (F ([String], d))
resolveSymbol a (p : ps) (LocalState {ephemeralScope = ephemeralScope, currentScopePath = currentPath}) = do
  GlobalState {unitInterface = UnitInterface {rootScope = rootScope}} <- get

  let matches =
        concatMap
          ( mapMaybe (\(p, (m, e)) -> (p,) . (m,) <$> fromExportedDecl e)
              . lookupScopeWithPath currentPath (p :| ps)
          )
          [rootScope, ephemeralScope]

  return $
    case matches of
      [(p, (_, e))] -> Right (toList p, e)
      [] ->
        Left
          [ Diagnostic
              Error
              ( printf "Could not resolve symbol %s" (intercalate "." (p : ps))
              )
              (unCommented a)
          ]
      (_ : _ : _) -> do
        Left
          [ Diagnostic
              Error
              ( printf
                  "Ambiguous occurance of %s"
                  (intercalate "." (p : ps))
              )
              (unCommented a)
          ]
resolveSymbol a _ _ =
  return $ Left [Diagnostic Error "Empty path provided (this is a bug)" (unCommented a)]

advanceFiddleDecls ::
  LocalState ->
  [Directed FiddleDecl S F A] ->
  M [Directed FiddleDecl Qualified F A]
advanceFiddleDecls localState decls = fmap (reverse . fst) $ do
  foldlM
    ( \(declsRet, localState' :: LocalState) unsqeezedd -> do
        d <- case squeeze unsqeezedd of
          Left diags -> tell diags >> compilationFailure
          Right x -> return x
        case unsqeezedd of
          (Directed directives t dann) ->
            let doReturn ::
                  FiddleDecl Qualified F A ->
                  M ([Directed FiddleDecl Qualified F A], LocalState)
                doReturn v = return (Directed directives v dann : declsRet, localState')
                doReturnWith s v = return (Directed directives v dann : declsRet, s)
                qualify = qualifyPath (currentScopePath localState')
                metadata = directiveToMetadata d
             in case t of
                  UsingDecl {usingName = name} ->
                    return (declsRet, modifyCurrentScopePath (addUsingPath (nameToList name)) localState')
                  OptionDecl key value ann -> doReturn $ OptionDecl key value ann
                  ImportDecl st@(ImportStatement {importInterface = interface}) a ->
                    let localState'' = modifyEphemeralScope (<> rootScope interface) localState'
                     in doReturnWith localState''
                          =<< ImportDecl
                            <$> advanceStage localState'' st
                            <*> pure a
                  PackageDecl _ name body ann ->
                    let qualifiedName = qualify (nameToList name)
                        localState'' = modifyCurrentScopePath (pushScope (nameToList name)) localState'
                        decl = ExportedPackageDecl (metadata qualifiedName)
                     in do
                          insertDecl decl
                          doReturn
                            =<< PackageDecl
                              (pure decl)
                              name
                              <$> mapM (advanceStage localState'') body
                              <*> pure ann
                  LocationDecl _ ident expr ann ->
                    let qualifiedName = qualify (NonEmpty.singleton (identToString ident))
                     in do
                          exprValue <- expressionToIntM expr
                          let decl =
                                ExportedLocationDecl
                                  (metadata qualifiedName)
                                  exprValue
                          insertDecl decl
                          doReturn
                            =<< LocationDecl
                              (pure decl)
                              ident
                              <$> advanceStage localState' expr
                              <*> pure ann
                  BitsDecl _ ident typ ann ->
                    let qualifiedName = qualify (NonEmpty.singleton (identToString ident))
                     in do
                          sizeBits <- getBitTypeDeclaredSize typ
                          let decl =
                                ExportedBitsDecl
                                  (metadata qualifiedName)
                                  sizeBits
                          insertDecl decl
                          doReturn
                            =<< BitsDecl
                              (pure decl)
                              ident
                              <$> advanceStage localState' typ
                              <*> pure ann
                  ObjTypeDecl _ ident body ann ->
                    let qualifiedName = qualify (NonEmpty.singleton (identToString ident))
                     in do
                          typeSize <- calculateTypeSize =<< resolveOrFail body
                          let decl =
                                ExportedTypeDecl
                                  (metadata qualifiedName)
                                  typeSize
                          insertDecl decl
                          doReturn
                            =<< ObjTypeDecl
                              (pure decl)
                              ident
                              <$> mapM (advanceStage localState') body
                              <*> pure ann
                  ObjectDecl _ ident loc typ ann ->
                    let qualifiedName = qualify (NonEmpty.singleton (identToString ident))
                     in do
                          location <- resolveLocationExpression localState' loc
                          exportedType <- objTypeToExport localState' typ
                          let decl =
                                ExportedObjectDecl
                                  (metadata qualifiedName)
                                  location
                                  exportedType
                          insertDecl decl
                          doReturn
                            =<< ObjectDecl
                              (pure decl)
                              ident
                              <$> advanceStage localState' loc
                              <*> advanceStage localState' typ
                              <*> pure ann
    )
    ([], localState)
    decls

insertDecl :: (ExportableDecl d) => d -> M ()
insertDecl decl =
  modify $ \(GlobalState ui) -> GlobalState (UnitInterface.insert decl ui)

objTypeToExport :: LocalState -> ObjType Expanded F A -> M ReferencedObjectType
objTypeToExport ls = \case
  ArrayObjType {arraySize = size, arrayObjType = objType} ->
    ArrayObjectType
      <$> objTypeToExport ls objType
      <*> expressionToIntM size
  ReferencedObjType {refName = n} -> do
    (full, _ :: ExportedTypeDecl) <- resolveOrFail =<< resolveName n ls
    return $ ReferencedObjectType (intercalate "." full)

calculateTypeSize :: ObjTypeBody Expanded F A -> M Word32
calculateTypeSize (ObjTypeBody bodyType decls _) =
  ( case bodyType of
      Union {} -> maximum
      Struct {} -> sum
  )
    <$> mapM calculateDeclSize decls
  where
    calculateDeclSize :: Directed ObjTypeDecl Expanded F A -> M Word32
    calculateDeclSize (undirected -> decl) =
      case decl of
        AssertPosStatement {} -> return 0
        RegisterDecl {regSize = size} -> expressionToIntM size
        ReservedDecl {reservedExpr = size} -> expressionToIntM size
        TypeSubStructure {subStructureBody = b} -> calculateTypeSize =<< resolveOrFail b

getBitTypeDeclaredSize :: BitType Expanded F A -> M Word32
getBitTypeDeclaredSize = \case
  (EnumBitType declaredSize _ _) -> expressionToIntM declaredSize
  (RawBits declaredSize _) -> expressionToIntM declaredSize

resolveLocationExpression ::
  (Integral i, Integral (NumberType stage)) =>
  LocalState ->
  Expression stage F A ->
  M i
resolveLocationExpression ls (Var var _) = do
  (_, ExportedLocationDecl _ v) <- resolveOrFail =<< resolveName var ls
  return (fromIntegral v)
resolveLocationExpression _ e = expressionToIntM e

expressionToIntM ::
  (Integral i, Integral (NumberType stage)) =>
  Expression stage f A ->
  M i
expressionToIntM expr =
  resolveOrFail $
    either
      ( \reason -> Left [Diagnostic Error reason (unCommented $ annot expr)]
      )
      return
      (expressionToInt expr)