{-# LANGUAGE NamedFieldPuns #-}
module Transducer (
      Transducer
    , fromList
    , run
  ) where

import Data.Map ((!), insert)
import Graph (Graph(..), Vertex(..), editVertex, follow, rewind, singleton, weave)
import Stream (Stream(..), merge)
import qualified Stream (empty)

type Transducer input output = Graph input [output]
newtype Stack a = Stack [a]

push :: a -> Stack a -> Stack a
push a (Stack s) = Stack (a:s)

content :: Stack a -> [a]
content (Stack s) = reverse s

empty :: Transducer input output
empty = singleton []

entry :: Ord input => Transducer input output -> ([input], output) -> Transducer input output
entry transducer (path, output) =
  rewind $ editVertex (weave transducer path) pushLabel
  where
    pushLabel vertex = vertex {label = output:(label vertex)}

loop :: Ord input => Transducer input output -> ([input], [input]) -> Transducer input output
loop transducer ([], loopPath) =
  case splitAt (length loopPath - 1) loopPath of
    (_, []) -> transducer
    (loopBegining, lastInput:_) ->
      let end = weave transducer loopBegining in
      rewind $
        editVertex end $ tieKnot lastInput
  where
    tieKnot input vertex =
      vertex {edges = insert input (focus transducer) $ edges vertex}
loop transducer (prefix, loopPath) = loop (weave transducer prefix) ([], loopPath)

fromList :: Ord input => [([input], Either [input] output)] -> Transducer input output
fromList = foldl add empty
  where
    add transducer (path, Left loopPath) = loop transducer (path, loopPath)
    add transducer (path, Right output) = entry transducer (path, output)

run :: (Ord input, Eq output) => Transducer input output -> Stream input -> Stream ([input], output)
run transducer = runWithStack initState
  where
    initState = (transducer, Stack [])
    emit stack stream output = ((content stack, output), runWithStack initState stream)
    runWithStack (state, stack) (Stream inputs) = foldl (\outputsStream (input, stream) ->
        case follow state input of
          Nothing -> Stream.empty
          Just newState@(Graph {vertices, focus}) ->
            let newStack = push input stack in
            let emitted = Stream $ emit newStack stream <$> label (vertices ! focus) in
            let continue = runWithStack (newState, newStack) stream in
            emitted `merge` continue `merge` outputsStream
      ) Stream.empty inputs