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

module Language.Fiddle.Compiler.Expansion (expandAst, expansionPhase) where

import Control.Monad.Identity (Identity (..))
import Control.Monad.State (get, modify, put)
import Data.List (intercalate)
import Data.Text (Text)
import Language.Fiddle.Ast
import Language.Fiddle.Compiler
import Language.Fiddle.Compiler.Qualification ()
import Language.Fiddle.Types

import qualified Data.Char as Char
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Text as Text

type M = Compile State

type Annot = Commented SourceSpan

type CurrentStage = ImportsResolved

newtype Path = Path [PathExpression]

newtype PathExpression = PathExpression String

joinPath :: Path -> String
joinPath (Path l) = intercalate "#" $ reverse (map (\(PathExpression s) -> s) l)

expandAst :: FiddleUnit CurrentStage I Annot -> Compile () (FiddleUnit Expanded I Annot)
expandAst = fmap snd . subCompile (State [] []) . advanceStage (Path mempty)

expansionPhase :: CompilationPhase CurrentStage Expanded
expansionPhase = CompilationPhase (pure ()) (\_ _ -> return ([], Just ())) (\_ _ -> expandAst)

-- Shorthand for Identity
type I = Identity

newtype Linkage = Linkage Text deriving (Show)

data State
  = State
      -- Anonymous object type bodies that need to be re-linked
      ![(Linkage, ObjTypeBody Expanded I Annot)]
      -- Anonymous enum bodies that need to be re-linked
      ![(Linkage, AnonymousBitsType Expanded I Annot)]

instance CompilationStage CurrentStage where
  type StageAfter CurrentStage = Expanded
  type StageMonad CurrentStage = M
  type StageState CurrentStage = Path
  type StageFunctor CurrentStage = Identity
  type StageAnnotation CurrentStage = Annot

deriving instance AdvanceStage CurrentStage ObjTypeBody

deriving instance AdvanceStage CurrentStage DeferredRegisterBody

deriving instance AdvanceStage CurrentStage RegisterBody

deriving instance AdvanceStage CurrentStage AnonymousBitsType

deriving instance AdvanceStage CurrentStage ImportStatement

deriving instance AdvanceStage CurrentStage BitType

deriving instance AdvanceStage CurrentStage EnumBody

deriving instance AdvanceStage CurrentStage EnumConstantDecl

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

instance AdvanceStage CurrentStage RegisterBitsDecl where
  modifyState t =
    return
      . case t of
        DefinedBits {definedBitsIdent = i} -> pushId i
        _ -> id

instance AdvanceStage CurrentStage PackageBody where
  advanceStage p (PackageBody decls a) =
    PackageBody <$> reconfigureFiddleDecls p decls <*> pure a

instance AdvanceStage CurrentStage ObjTypeDecl where
  modifyState t =
    return
      . case t of
        TypeSubStructure {subStructureName = (Just n)} -> pushId n
        RegisterDecl {regIdent = (Just n)} -> pushId n
        _ -> id

instance AdvanceStage CurrentStage FiddleDecl where
  modifyState t =
    return
      . case t of
        PackageDecl {packageName = n} -> pushName n
        BitsDecl {bitsIdent = i} -> pushId i
        ObjTypeDecl {objTypeIdent = i} -> pushId i
        ObjectDecl {objectIdent = i} -> pushId i
        _ -> id

instance AdvanceStage CurrentStage FiddleUnit where
  advanceStage path (FiddleUnit _ decls a) =
    FiddleUnit () <$> reconfigureFiddleDecls path decls <*> pure a

instance AdvanceStage CurrentStage Expression where
  advanceStage _ = \case
    (Var i a) -> return $ Var i a
    (LitNum t a) -> LitNum <$> parseNum (unCommented a) t <*> pure a

instance AdvanceStage CurrentStage RegisterBitsTypeRef where
  advanceStage path = \case
    RegisterBitsArray typeref expr annot ->
      RegisterBitsArray
        <$> advanceStage path typeref
        <*> advanceStage path expr
        <*> pure annot
    RegisterBitsReference q name annot ->
      return $ RegisterBitsReference q name annot
    RegisterBitsJustBits expr annot ->
      RegisterBitsJustBits
        <$> advanceStage path expr
        <*> pure annot
    RegisterBitsAnonymousType _ anonType annot -> do
      ident <-
        internAnonymousBitsType path
          =<< advanceStage path anonType
      return $ RegisterBitsReference (pure ()) (identToName ident) annot

instance AdvanceStage CurrentStage ObjType where
  advanceStage path = \case
    (AnonymousObjType _ (Identity body) annot) -> do
      body' <- advanceStage path body
      identifier <- internObjType path body'
      return (ReferencedObjType (pure ()) (identToName identifier) annot)
    (ReferencedObjType q name annot) ->
      return $ ReferencedObjType q name annot
    (ArrayObjType objType expr a) ->
      ArrayObjType
        <$> advanceStage path objType
        <*> advanceStage path expr
        <*> pure a

parseNum :: SourceSpan -> Text -> Compile s Integer
parseNum span txt = fromMayberOrFail span "Unable to parse number" $
  case Text.unpack (Text.take 2 txt) of
    "0b" -> toNumWithRadix (Text.drop 2 txt) 2
    "0x" -> toNumWithRadix (Text.drop 2 txt) 16
    ('0' : _) -> toNumWithRadix (Text.tail txt) 8
    _ -> toNumWithRadix txt 10
  where
    removeUnders :: Text -> Text
    removeUnders = Text.replace (Text.pack "_") Text.empty

    toNumWithRadix :: Text -> Int -> Maybe Integer
    toNumWithRadix (removeUnders -> txt) (fromIntegral -> radix) =
      Text.foldl
        ( \mAcc x ->
            mAcc >>= (\acc -> (acc * radix +) <$> digitToInt x radix)
        )
        (Just 0)
        txt

    digitToInt :: Char -> Integer -> Maybe Integer
    digitToInt (Char.toLower -> ch) radix =
      let a
            | Char.isDigit ch = Just (Char.ord ch - Char.ord '0')
            | ch >= 'a' && ch <= 'f' = Just $ (Char.ord ch - Char.ord 'a') + 10
            | otherwise = Nothing
       in a
            >>= ( \a' ->
                    if a' >= fromIntegral radix
                      then Nothing
                      else Just (fromIntegral a')
                )

reconfigureFiddleDecls ::
  Path ->
  [Directed FiddleDecl CurrentStage I Annot] ->
  M [Directed FiddleDecl Expanded I Annot]
reconfigureFiddleDecls p decls = do
  lastState <- get
  put (State [] [])
  decls <- mapM (mapDirectedM $ advanceStage p) decls
  (State anonymousObjTypes anonymousBitsTypes) <- get
  put lastState

  return $
    map (asDirected . resolveAnonymousObjType) anonymousObjTypes
      ++ map (asDirected . resolveAnonymousBitsType) anonymousBitsTypes
      ++ decls
  where
    resolveAnonymousObjType (Linkage linkage, objTypeBody) =
      ObjTypeDecl
        (pure ())
        (Identifier linkage (annot objTypeBody))
        (pure objTypeBody)
        (annot objTypeBody)

    resolveAnonymousBitsType (Linkage linkage, AnonymousEnumBody expr body a) =
      BitsDecl (pure ()) (Identifier linkage a) (EnumBitType expr body a) a

identToName :: Identifier I a -> Name I a
identToName ident = Name (NonEmpty.singleton ident) (annot ident)

internObjType :: Path -> ObjTypeBody Expanded I Annot -> M (Identifier I Annot)
internObjType path body =
  let str = Text.pack $ joinPath path
   in do
        modify $ \(State objTypeBodies a) ->
          State ((Linkage str, body) : objTypeBodies) a
        return (Identifier str (annot body))

internAnonymousBitsType ::
  Path ->
  AnonymousBitsType Expanded I Annot ->
  M (Identifier I Annot)
internAnonymousBitsType path anonymousBitsType =
  let str = Text.pack $ joinPath path
   in do
        modify $ \(State a anonymousBitsTypes) ->
          State a ((Linkage str, anonymousBitsType) : anonymousBitsTypes)
        return (Identifier str (annot anonymousBitsType))

pushId :: Identifier f a -> Path -> Path
pushId (Identifier str _) (Path lst) =
  Path (PathExpression (Text.unpack str) : lst)

pushName :: Name f a -> Path -> Path
pushName (Name idents _) path =
  foldl (flip pushId) path idents