{-# 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 qualified Data.Char as Char
import Data.List (intercalate)
import Data.List.NonEmpty (NonEmpty, (<|))
import qualified Data.List.NonEmpty as NonEmpty
import Data.Text (Text)
import qualified Data.Text as Text
import Language.Fiddle.Ast
import Language.Fiddle.Compiler
import Language.Fiddle.Compiler.Qualification ()
import Language.Fiddle.Internal.UnitNumbers
import Language.Fiddle.Types

type M = Compile State

type Annot = Commented SourceSpan

type CurrentStage = ImportsResolved

type Path = [Text]

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

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

-- Shorthand for Identity
type I = Identity

newtype Linkage = Linkage (NonEmpty 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 (ConstExpression u)

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 = (Perhaps (Just n))} -> pushId n
        _ -> id

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

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

instance AdvanceStage CurrentStage (Expression u) where
  advanceStage _ = \case
    (Var i a) -> return $ Var i a
    (LitNum (LeftV t) a) ->
      LitNum . RightV <$> 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
      name <-
        internAnonymousBitsType path
          =<< advanceStage path anonType
      return $ RegisterBitsReference (Identity Vacant) name annot

instance AdvanceStage CurrentStage ObjType where
  advanceStage path = \case
    (AnonymousObjType _ (Identity body) annot) -> do
      body' <- advanceStage path body
      name <- internObjType path body'
      return (ReferencedObjType (Identity Vacant) name 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 (N u)
parseNum span txt = fmap NumberWithUnit $
  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 Int
    toNumWithRadix (removeUnders -> txt) (fromIntegral -> radix) =
      Text.foldl
        ( \mAcc x ->
            mAcc >>= (\acc -> (acc * radix +) <$> digitToInt x radix)
        )
        (Just 0)
        txt

    digitToInt :: Char -> Int -> Maybe Int
    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
        (Identity Vacant)
        (Name (fmap (\t -> Identifier t (annot objTypeBody)) (NonEmpty.reverse linkage)) (annot objTypeBody))
        (pure objTypeBody)
        (annot objTypeBody)

    resolveAnonymousBitsType (Linkage linkage, AnonymousEnumBody expr body a) =
      BitsDecl
        (Identity Vacant)
        (Name (fmap (\t -> Identifier t a) (NonEmpty.reverse 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 (Name I Annot)
internObjType [] _ = compilationFailure
internObjType (NonEmpty.fromList -> path) body =
  do
    modify $ \(State objTypeBodies a) ->
      State ((Linkage path, body) : objTypeBodies) a
    let a = annot body
     in return (Name (fmap (\t -> Identifier t a) (NonEmpty.reverse path)) a)

internAnonymousBitsType ::
  Path ->
  AnonymousBitsType Expanded I Annot ->
  M (Name I Annot)
internAnonymousBitsType [] _ = compilationFailure
internAnonymousBitsType (NonEmpty.fromList -> path) anonymousBitsType =
  do
    modify $ \(State a anonymousBitsTypes) ->
      State a ((Linkage path, anonymousBitsType) : anonymousBitsTypes)
    let a = annot anonymousBitsType
     in return (Name (fmap (\t -> Identifier t a) (NonEmpty.reverse path)) a)

pushId :: Identifier f a -> Path -> Path
pushId (Identifier str _) lst = str : lst

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