-- | Extracting documentation from instructions set.
module Michelson.Typed.Doc
  ( buildInstrDoc
  , modifyInstrDoc
  , cutInstrNonDoc
  ) where

import Control.Lens (at)
import Data.Default (def)
import Data.Typeable (cast)
import Prelude hiding (Ordering(..))

import Michelson.Doc
import Michelson.Typed.Instr
import Michelson.Typed.Util

someDefinitionDocItemToContractDoc :: SomeDocDefinitionItem -> State ContractDoc ()
someDefinitionDocItemToContractDoc sdi@(SomeDocDefinitionItem di) =
  modify $ flip mappend
    mempty
    { cdContents = mempty
    , cdDefinitions = docItemToBlock di
    , cdDefinitionsSet = one sdi
    , cdDefinitionIds = one $ case docItemRef di of
        DocItemRef docItemId -> docItemId
    }

someDocItemToContractDoc :: SomeDocItem -> State ContractDoc ()
someDocItemToContractDoc (SomeDocItem di) = do
  () <- case docItemRef di of
    DocItemNoRef ->
      modify (<> mempty{ cdContents = docItemToBlock di })
    DocItemRef{} ->
      someDefinitionDocItemToContractDoc (SomeDocDefinitionItem di)
  forM_ @_ @_ @() (docItemDependencies di) $ \(SomeDocDefinitionItem dep) ->
    case docItemRef dep of
      DocItemRef{} -> do
        -- Taking special treatment for possible cyclic dependencies.
        isPresent <- use $ cdDefinitionsSetL . at (SomeDocDefinitionItem dep)
        case isPresent of
          Just () -> pass
          Nothing -> someDocItemToContractDoc (SomeDocItem dep)

-- | Assemble contract documentation.
buildInstrDoc :: Instr inp out -> ContractDoc
buildInstrDoc = dfsFoldInstr dfsSettings $ \case
  Ext ext -> case ext of
    DOC_ITEM sdi ->
      execState (someDocItemToContractDoc sdi) mempty
    _ -> mempty
  _ -> mempty
  where
  dfsSettings :: DfsSettings ContractDoc
  dfsSettings = def
    { dsCtorEffectsApp = CtorEffectsApp
        { ceaName = "Building DocGroup"
        , ceaApplyEffects = \resChildren _ -> \case
            i@(DocGroup grouping _) ->
              (i, docGroupContent grouping resChildren)
            i -> (i, resChildren)
        }
    }

-- | Modify all documentation items recursively.
modifyInstrAllDoc
  :: (SomeDocItem -> SomeDocItem)
  -> Instr inp out
  -> Instr inp out
modifyInstrAllDoc mapper = dfsModifyInstr def $ \case
  Ext ext -> Ext $
    case ext of
      DOC_ITEM sdi -> DOC_ITEM (mapper sdi)
      i -> i
  i -> i

-- | Recursevly traverse an instruction and modify documentation items
-- matching given type.
modifyInstrDoc
  :: DocItem i
  => (i -> i)
  -> Instr inp out
  -> Instr inp out
modifyInstrDoc mapper = modifyInstrAllDoc untypedMapper
  where
  untypedMapper sdi@(SomeDocItem di) = maybe sdi (SomeDocItem . mapper) (cast di)

-- | Leave only instructions related to documentation.
--
-- Generated documentation for resulting instruction remains the same, but
-- semantics of instruction itself gets lost.
-- We have to pass optimizer here as an argument to avoid cyclic dependencies.
cutInstrNonDoc :: (forall i o. Instr i o -> Instr i o) -> Instr inp out -> Instr s s
cutInstrNonDoc optimize = optimize . dfsFoldInstr dfsSettings step
  where
  dfsSettings :: DfsSettings (Instr s s)
  dfsSettings = def
    { dsCtorEffectsApp = CtorEffectsApp
        { ceaName = "Wrap into DocGroup"
        , ceaApplyEffects = \resChildren _ -> \case
            i@(DocGroup g _) -> (i, DocGroup g resChildren)
            i -> (i, resChildren)
        }
    }
  step :: Instr inp out -> Instr s s
  step = \case
    Ext ext -> case ext of
      DOC_ITEM di -> Ext $ DOC_ITEM di
      _ -> Nop
    _ -> Nop