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

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

import Control.Monad.Identity (Identity (..))
import Control.Monad.State (get, gets, modify, put)
import qualified Data.Char as Char
import Data.List (intercalate)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Type.Bool
import Debug.Trace
import GHC.TypeLits
import Language.Fiddle.Ast
import Language.Fiddle.Compiler
import Language.Fiddle.Types
import Text.Printf (printf)

type Annot = Commented SourceSpan

newtype Path = Path [PathExpression]

newtype PathExpression = PathExpression String

type M = Compile State

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

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

-- 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 Parsed where
  type StageAfter Parsed = Expanded
  type StageMonad Parsed = M
  type StageState Parsed = Path
  type StageFunctor Parsed = Identity
  type StageAnnotation Parsed = Annot

deriving instance AdvanceStage Parsed ObjTypeBody

deriving instance AdvanceStage Parsed DeferredRegisterBody

deriving instance AdvanceStage Parsed RegisterBody

deriving instance AdvanceStage Parsed AnonymousBitsType

deriving instance AdvanceStage Parsed BitType

deriving instance AdvanceStage Parsed EnumBody

deriving instance AdvanceStage Parsed EnumConstantDecl

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

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

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

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

instance AdvanceStage Parsed 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 Parsed FiddleUnit where
  advanceStage path (FiddleUnit decls a) =
    FiddleUnit <$> reconfigureFiddleDecls path decls <*> pure a

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

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

instance AdvanceStage Parsed ObjType where
  advanceStage path = \case
    (AnonymousObjType _ (Identity body) annot) -> do
      body' <- advanceStage path body
      identifier <- internObjType path body'
      return (ReferencedObjType (identToName identifier) annot)
    (ReferencedObjType name annot) ->
      return $ ReferencedObjType 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 Parsed 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
        (Identifier linkage (annot objTypeBody))
        (pure objTypeBody)
        (annot objTypeBody)

    resolveAnonymousBitsType (Linkage linkage, AnonymousEnumBody expr body a) =
      BitsDecl (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