{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE InstanceSigs #-}
module Language.Java.Classfile.Extract (Extract(), bytes, runExtract, expectRaw, expectEqual) where

import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy as ByteString
import Control.Applicative (Alternative (empty, (<|>)))

data Extract a = Extract (Continuation a)

type Continuation a = ByteString -> Reply a

data Reply a
  = Done ByteString a -- rest, result
  | Fail
  deriving (Functor)

instance Functor Extract where
  fmap :: (a -> b) -> Extract a -> Extract b
  fmap f (Extract cont) = Extract $ \ input -> f <$> cont input

instance Applicative Extract where
  pure :: a -> Extract a
  pure x = Extract $ \ rest -> Done rest x
  (<*>) :: Extract (a -> b) -> Extract a -> Extract b
  (<*>) (Extract computeF) (Extract computeX) = Extract $ \ input -> case computeF input of
    Done rest f -> case computeX rest of
      Done rest' x -> Done rest' (f x)
      Fail -> Fail
    Fail -> Fail

instance Monad Extract where
  (>>=) :: Extract a -> (a -> Extract b) -> Extract b
  (>>=) (Extract computeA) f = Extract $ \ input -> case computeA input of
    Done rest a -> let
      (Extract computeB) = f a
      in computeB rest
    Fail -> Fail

instance Alternative Extract where
  empty :: Extract a
  empty = Extract $ const Fail
  (<|>) :: Extract a -> Extract a -> Extract a
  (<|>) (Extract left) (Extract right) = Extract $ \ input -> case left input of
    Fail -> right input
    t -> t

-- | Get a specified count of bytes. Fail if there are not enough bytes available.

bytes :: Word -> Extract ByteString
bytes count = Extract $ \ input -> let
  count' = fromIntegral count
  (bs, rest) = ByteString.splitAt count' input
  in if ByteString.length bs /= count'
    then Fail
    else Done rest bs

expectRaw :: ByteString -> Extract ByteString
expectRaw e = Extract $ \ input -> let
     (actual, rest) = ByteString.splitAt (ByteString.length e) input
     in case actual == e of
       True -> Done rest actual
       False -> Fail 

expectEqual :: Eq a => Extract a -> a -> Extract a
expectEqual (Extract computeActual) expected = Extract $ \ input -> case computeActual input of
  Done rest actual -> if actual == expected
    then Done rest actual
    else Fail
  Fail -> Fail

{- It seems I cannot define a lawful monad instance
instance Monad Extract where
  (>>=) :: Extract a -> (a -> Extract b) -> Extract b
  (>>=) (Extract computeA) f = Extract $ \ input -> case computeA input of
    Done rest a -> _
    Fail -> Fail
-}

runExtract :: ByteString -> Extract b -> Maybe (ByteString, b)
runExtract string (Extract computation) = case computation string of
  Fail -> Nothing
  Done rest x -> Just (rest, x)