-- | Qualification compilation pha se.
--
-- 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, toList)
import Data.List (intercalate)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Maybe (isNothing, mapMaybe)
import qualified Data.Text
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.Internal.UnitNumbers
import Language.Fiddle.Types
import Text.Printf (printf)

type S = Expanded

data GlobalState = GlobalState
  { unitInterface :: UnitInterface,
    uniqueCounter :: Int
  }

data LocalState = LocalState
  { currentScopePath :: ScopePath String,
    -- | Current qualified path, used for building metadata.
    currentQualifiedPath :: QualifiedPath (),
    ephemeralScope :: Scope String (Metadata, ExportedDecl)
  }

type F = Either [Diagnostic]

type A = Commented SourceSpan

type M = Compile GlobalState

uniqueString :: String -> M String
uniqueString prefix = do
  cnt <- gets uniqueCounter
  modify $ \g -> g {uniqueCounter = cnt + 1}
  return $ "_" ++ prefix ++ show cnt

uniqueIdentifier :: String -> a -> M (Identifier F a)
uniqueIdentifier prefix a = (\s -> Identifier (Data.Text.pack s) a) <$> uniqueString prefix

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

pushPackage :: NonEmpty String -> LocalState -> (QualifiedPath String, LocalState)
pushPackage pk ls =
  let q = currentQualifiedPath ls
   in ( qualifyPackage pk,
        ls
          { currentScopePath = pushScope pk (currentScopePath ls),
            currentQualifiedPath =
              q
                { packagePart = packagePart q ++ NonEmpty.toList pk
                }
          }
      )
  where
    qualifyPackage :: NonEmpty String -> QualifiedPath String
    qualifyPackage (NonEmpty.reverse -> (l :| (reverse -> h))) =
      let q = currentQualifiedPath ls
       in q {packagePart = packagePart q ++ h, basenamePart = l}

pushObject :: NonEmpty String -> LocalState -> (QualifiedPath String, LocalState)
pushObject obj ls =
  let q = currentQualifiedPath ls
   in ( qualifyObject obj,
        ls
          { currentScopePath = pushScope obj (currentScopePath ls),
            currentQualifiedPath =
              q
                { objectPart = objectPart q ++ NonEmpty.toList obj
                }
          }
      )
  where
    qualifyObject :: NonEmpty String -> QualifiedPath String
    qualifyObject (NonEmpty.reverse -> (l :| (reverse -> h))) =
      let q = currentQualifiedPath ls
       in q {objectPart = objectPart q ++ h, basenamePart = l}

pushRegister :: String -> LocalState -> (QualifiedPath String, LocalState)
pushRegister regName ls =
  let q = currentQualifiedPath ls
   in ( fmap (const regName) q,
        ls
          { currentScopePath = pushScope (NonEmpty.singleton regName) (currentScopePath ls),
            currentQualifiedPath =
              q
                { registerPart = registerPart q ++ [regName]
                }
          }
      )

qualificationPhase :: CompilationPhase Expanded Qualified
qualificationPhase =
  let initialQualifiedPath = QualifiedPath [] [] [] ()
   in pureCompilationPhase $ \t -> do
        raw <-
          fmap snd $
            subCompile (GlobalState mempty 0) $
              advanceStage
                (LocalState mempty initialQualifiedPath mempty)
                (soakA t)

        squeezeDiagnostics raw

instance
  StageConvertible
    Expanded
    (When False String)
    (When True String)
  where
  convertInStage _ _ _ _ = Present <$> uniqueString "reserved"

instance AdvanceStage S (ConstExpression u) where
  advanceStage ls (ConstExpression (RightV exp) a) =
    case exp of
      Var var _ -> do
        (ExportedLocationDecl _ v) <- resolveOrFail =<< resolveName var ls
        return $ ConstExpression (LeftV $ fromIntegral v) a
      LitNum (RightV v) _ -> return $ ConstExpression (LeftV v) a

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

instance AdvanceStage S RegisterBitsDecl where
  advanceStage localState = \case
    ReservedBits expr an -> ReservedBits <$> advanceStage localState expr <*> pure an
    BitsSubStructure Vacant mod bod name an -> do
      let (path, localState') =
            maybe
              (fmap (const "") (currentQualifiedPath localState), localState)
              ( \ident -> pushRegister (identToString ident) localState
              )
              name
          qMeta =
            QBitsMetadata
              { bitsSpan = Vacant,
                bitsFullPath = path,
                bitsUnnamed = isNothing name
              }
      BitsSubStructure (Present qMeta) (guarantee (ModifierKeyword Rw an) mod)
        <$> advanceStage localState' bod
        <*> pure name
        <*> pure an
    DefinedBits _ mod ident typ an -> do
      let (path, _) = pushRegister (identToString ident) localState
          qMeta =
            QBitsMetadata
              { bitsSpan = Vacant,
                bitsFullPath = path,
                bitsUnnamed = False
              }
      DefinedBits (Present qMeta) (guarantee (ModifierKeyword Rw an) mod) ident
        <$> advanceStage localState typ
        <*> pure an

getProperRegSize :: (stage .< Expanded ~ False) => Expression Bits stage F A -> M RegSz
getProperRegSize expr = do
  v <- expressionToIntM expr
  case v of
    8 -> return RegSz8
    16 -> return RegSz16
    32 -> return RegSz32
    64 -> return RegSz64
    _ -> do
      emitDiagnosticError "Exotic register size." (annot expr)
      return RegSz32

instance AdvanceStage S ObjTypeDecl where
  advanceStage localState = \case
    AssertPosStatement d e a ->
      AssertPosStatement d <$> advanceStage localState e <*> pure a
    SkipToStatement d _ e a -> do
      ident' <- uniqueIdentifier "reg" a
      let (qualified, localState') = pushRegister (identToString ident') localState

      let qRegMeta =
            QRegMetadata
              { regSpan = Vacant,
                regIsPadding = True,
                regIsUnnamed = True,
                regFullPath = qualified
              }

      SkipToStatement d (Present qRegMeta) <$> advanceStage localState' e <*> pure a
    BufferDecl _ ident sz ann -> do
      ident' <- guaranteeM (uniqueIdentifier "reg" ann) ident

      let (qualified, localState') =
            pushRegister (identToString $ unwrap ident') localState
          -- Avoid pushing the anonymized name onto the stack.
          localState'' =
            if isNothing (toMaybe ident) then localState else localState'

      let qRegMeta =
            QRegMetadata
              { regSpan = Vacant,
                regIsPadding = False,
                regIsUnnamed = isNothing (toMaybe ident),
                regFullPath = qualified
              }

      BufferDecl
        (Present qRegMeta)
        ident'
        <$> advanceStage localState'' sz
        <*> pure ann
    RegisterDecl _ mod ident size Vacant bod ann -> do
      ident' <- guaranteeM (uniqueIdentifier "reg" ann) ident

      let (qualified, localState') =
            pushRegister (identToString $ unwrap ident') localState
          -- Avoid pushing the anonymized name onto the stack.
          localState'' =
            if isNothing (toMaybe ident) then localState else localState'

      let qRegMeta =
            QRegMetadata
              { regSpan = Vacant,
                regIsPadding = False,
                regIsUnnamed = isNothing (toMaybe ident),
                regFullPath = qualified
              }

      RegisterDecl
        (Present qRegMeta)
        (guarantee (ModifierKeyword Rw ann) mod)
        ident'
        <$> progressBackM getProperRegSize size
        <*> pure Vacant
        <*> mapM (advanceStage localState'') bod
        <*> pure ann
    ReservedDecl _ expr ann -> do
      ident <- uniqueIdentifier "reserved" ann

      let (q, _) = pushRegister (identToString ident) localState

      let qRegMeta =
            QRegMetadata
              { regSpan = Vacant,
                regIsPadding = True,
                regIsUnnamed = True,
                regFullPath = q
              }

      RegisterDecl
        (Present qRegMeta)
        (Guaranteed $ ModifierKeyword Pr ann)
        (Guaranteed ident)
        <$> fmap LeftV (getProperRegSize expr)
        <*> pure Vacant
        <*> pure Nothing
        <*> pure ann
    TypeSubStructure bod name an -> do
      let localState' =
            maybe
              localState
              (\n -> snd $ pushRegister (identToString n) localState)
              name
      TypeSubStructure
        <$> mapM (advanceStage localState') bod
        <*> pure name
        <*> pure an

deriving instance AdvanceStage S (Expression u)

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 Present <$> 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 Present <$> 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 v decls a) =
    FiddleUnit v <$> 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 d)
resolveIdent i = resolveSymbol (annot i) [identToString i]

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

resolveSymbol :: (ExportableDecl d) => A -> [String] -> LocalState -> M (F 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
      [(_, (_, e))] -> Right 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)]

qMd :: (Applicative f) => t -> f (QMd Qualified t)
qMd = pure . Present

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)

                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 (unwrap interface)) localState'
                     in doReturnWith localState''
                          =<< ImportDecl
                            <$> advanceStage localState'' st
                            <*> pure a
                  PackageDecl _ name body ann ->
                    let (qualifiedName, localState'') =
                          pushPackage
                            (nameToList name)
                            localState'
                        decl = ExportedPackageDecl (metadata qualifiedName)
                     in do
                          insertDecl decl
                          doReturn
                            =<< PackageDecl
                              (qMd decl)
                              name
                              <$> mapM (advanceStage localState'') body
                              <*> pure ann
                  LocationDecl _ ident expr ann ->
                    let qualifiedName =
                          fmap
                            (const $ identToString ident)
                            (currentQualifiedPath localState')
                     in do
                          expr' <- advanceStage localState' expr
                          let decl =
                                ExportedLocationDecl
                                  (metadata qualifiedName)
                                  (trueValue expr')
                          insertDecl decl
                          doReturn
                            =<< LocationDecl
                              (qMd decl)
                              ident
                              <$> advanceStage localState' expr
                              <*> pure ann
                  BitsDecl _ name@(Name ids _) typ ann ->
                    let (qualifiedName, _) =
                          pushObject
                            (fmap identToString ids)
                            localState'
                     in do
                          sizeBits <- getBitTypeDeclaredSize typ
                          let decl =
                                ExportedBitsDecl
                                  (metadata qualifiedName)
                                  sizeBits
                          insertDecl decl
                          doReturn
                            =<< BitsDecl
                              (qMd decl)
                              name
                              <$> advanceStage localState' typ
                              <*> pure ann
                  ObjTypeDecl _ name@(Name ids _) body ann ->
                    let (qualifiedName, localState'') =
                          pushObject
                            (fmap identToString ids)
                            localState'
                     in do
                          body' <- mapM (advanceStage localState'') body
                          typeSize <- calculateTypeSize =<< resolveOrFail body'
                          let decl =
                                ExportedTypeDecl
                                  (metadata qualifiedName)
                                  typeSize
                          insertDecl decl
                          doReturn $ ObjTypeDecl (qMd decl) name body' ann
                  ObjectDecl _ ident loc typ ann ->
                    let qualifiedName =
                          fmap
                            (const $ identToString ident)
                            (currentQualifiedPath localState')
                     in do
                          location <- resolveLocationExpression localState' loc
                          typ' <- advanceStage localState' typ
                          exportedType <- objTypeToExport localState' typ'
                          let decl =
                                ExportedObjectDecl
                                  (metadata qualifiedName)
                                  location
                                  exportedType
                          insertDecl decl
                          doReturn
                            =<< ObjectDecl
                              (qMd decl)
                              ident
                              <$> advanceStage localState' loc
                              <*> pure typ'
                              <*> pure ann
    )
    ([], localState)
    decls

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

objTypeToExport :: LocalState -> ObjType Qualified F A -> M ReferencedObjectType
objTypeToExport ls = \case
  ArrayObjType {arraySize = size, arrayObjType = objType} ->
    ArrayObjectType
      <$> objTypeToExport ls objType
      <*> pure (trueValue size)
  ReferencedObjType {refName = n} -> do
    (td :: ExportedTypeDecl) <- resolveOrFail =<< resolveName n ls
    return $ ReferencedObjectType (metadataFullyQualifiedPath $ getMetadata td)

calculateTypeSize :: ObjTypeBody Qualified F A -> M (N Bytes)
calculateTypeSize (ObjTypeBody bodyType decls _) = do
  (summed, maxxed) <- foldlM f (0, 0) decls
  return $
    case bodyType of
      Union {} -> maxxed
      Struct {} -> summed
  where
    f :: (N Bytes, N Bytes) -> Directed ObjTypeDecl Qualified F A -> M (N Bytes, N Bytes)
    f (pos, mx) decl = case undirected decl of
      AssertPosStatement {} -> return (pos, mx)
      RegisterDecl {regSize = size} -> do
        let v = fst $ bitsToBytes $ regSzToBits $ getLeft size
        return (pos + v, max mx v)
      SkipToStatement _ _ expr ann -> do
        let newLoc = getLeft (constExpression expr)
            sz = newLoc - pos
         in do
              if newLoc < pos
                then do
                  emitDiagnosticError "Skip to location already passed." ann
                  return (pos, mx)
                else
                  return (pos + sz, max mx sz)
      BufferDecl {bufSize = expr} ->
        let v = getLeft (constExpression expr)
         in return (pos + v, max mx v)
      TypeSubStructure {subStructureBody = b} -> do
        v <- calculateTypeSize =<< resolveOrFail b
        return (pos + v, max mx v)

-- calculateDeclSize :: Directed ObjTypeDecl Qualified F A -> M (N Bytes)
-- calculateDeclSize (undirected -> decl) =
--   case decl of
--     AssertPosStatement {} -> return 0
--     RegisterDecl {regSize = size} -> fst . bitsToBytes <$> expressionToIntM (getRight size)
--     ReservedDecl {reservedExpr = size} -> fst . bitsToBytes <$> expressionToIntM size
--     TypeSubStructure {subStructureBody = b} -> calculateTypeSize =<< resolveOrFail b

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

resolveLocationExpression ::
  (stage .< Expanded ~ False) =>
  LocalState ->
  Expression u stage F A ->
  M (N u)
resolveLocationExpression ls (Var var _) = do
  (ExportedLocationDecl _ v) <- resolveOrFail =<< resolveName var ls
  return (fromIntegral v)
resolveLocationExpression _ e = expressionToIntM e

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