{-# LANGUAGE Rank2Types, FlexibleInstances, TypeSynonymInstances,
             FlexibleContexts, MultiParamTypeClasses #-}

-- | XML picklers based on hexpat which are almost source code compatible with
-- HXT.
module Text.XML.Expat.Pickle (
        PU_(PU),
        Node,
        PU,
        pickleXML,
        unpickleXML,
        XmlPickler(..),
        pickleTree,
        unpickleTree,
        -- * Pickler primitives
        xpReadShow, -- nmrp3/drdozer
        xpText0,
        xpText,
        xpElem,
        xpAttr,
        xpOption,
        xpPair,
        xpTriple,
        xp4Tuple,
        xp5Tuple,
        xpList,
        xpWrap,
        xpWrapMaybe,
        xpWrapEither, -- nmrp3/drdozer
        xpAllAttrs,
        xpAlt,
        xpUnit,
        -- * Classes for abstracting parts of the tree
        Stringable,
        Nodeable,
        Attrable
    ) where

import Text.XML.Expat.Tree
import Text.XML.Expat.Format
import Data.Maybe
import Data.Either
import Data.List
import Data.Char
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Codec.Binary.UTF8.String as U8


-- | A two-way pickler/unpickler that pickles a ''t'' to type ''a''.  See
-- /Text.XML.Expat.Tree/ for the tree structure as used with ''t''.  ''t'' can be
-- /Node/, /[Node]/, /[(String,String)]/, /(String,String)/ or /String/.
--
-- A /PU_/ can be composed using the pickler primitives defined in this module.
data PU_ t a = PU {
        unpickleTree_ :: t -> Either String a,
        pickleTree_   :: a -> t -> t
    }

-- | In the most common case, where the part of the tree you're pickling/unpickling
-- is of type /Node/, you can use /PU/ and maintain source code compatibility with HXT. In
-- other cases you will need to use PU_, which will break compatibility.
type PU a = PU_ Node a

toByteString :: String -> B.ByteString
toByteString = B.pack . map (fromIntegral . ord)

-- | Pickle a Haskell data structure to XML text. Outputs a strict ByteString.
pickleXML :: Maybe Encoding -> PU_ Node a -> a -> B.ByteString
pickleXML mEnc pu value = toByteString $ formatDoc mEnc $ pickleTree pu value

-- | Unpickle XML text to a Haskell data structure.  Takes a lazy ByteString.
unpickleXML :: PU_ Node a -> Maybe Encoding -> BL.ByteString -> Either String a
unpickleXML pu enc xml = do
    case parse enc xml of
        Just tree -> unpickleTree pu tree
        Nothing   -> Left "XML parse failed"

-- | We build the attributes and tags lists backwards, then reverse them afterwards
-- for speed.
reverseElement :: Node -> Node
reverseElement (Element eName eAttrs eChildren) = Element eName (reverse eAttrs) (reverse eChildren)
reverseElement other = other

pickleTree :: PU a -> a -> Node
pickleTree pu value =
    let Element _ _ elems = pickleTree_ pu value (Element "" [] [])
    in  case elems of
            elem:_ -> elem
            _      -> error "No top-level element"

unpickleTree :: PU a -> Node -> Either String a
unpickleTree pu tree = unpickleTree_ pu (Element "" [] [tree])

-- | Takes one more argument than the HXT version of XmlPickler, which is the
-- type of the part of the tree, like /PU_/ does.
class XmlPickler t a where
    xpickle :: PU_ t a

class Show t => Stringable t where
    getString :: t -> String
    putString :: String -> t -> t

instance Stringable String where
    getString = id
    putString = \_ -> id

instance Stringable (String, String) where
    getString (name, value) = value
    putString value (name, _) = (name, value)
    
instance Stringable Node where
    getString (Element _ _ eChildren) = getString eChildren
    getString (Text str) = str
    putString value (Element eName eAttrs eChildren) = Element eName eAttrs (putString value eChildren)
    putString _ (Text _) = error "Can't put string into an existing text node"

-- | Tree parts that can be treated as strings.
instance Stringable [Node] where
    getString = concatMap ex
        where
            ex (Text str) = str
            ex _ = []
    putString "" nodes = nodes
    putString txt nodes = Text txt:nodes

-- | Tree parts that can be treated as nodes.
class Show t => Nodeable t where
    getNodes    :: t -> [Node]
    getSubnodes :: t -> [Node]
    addChild    :: Node -> t -> t

instance Nodeable [Node] where
    getNodes = id
    getSubnodes = id
    addChild = (:)

instance Nodeable Node where
    getNodes elt@(Element _ _ _) = [elt]
    getNodes _                   = error "No nodes to be found inside tag text"
    getSubnodes (Element _ _ children) = children
    getSubnodes _                = error "No subnodes to be found inside tag text"
    addChild value elt@(Element name attrs children) = Element name attrs (value `addChild` children)
    addChild _ _               = error "Can't append node to text tag"

-- | Tree parts that can be treated as attributes.
class Show t => Attrable t where
    getAttrs    :: t -> [(String,String)]
    putAttrs    :: [(String,String)] -> t -> t

instance Attrable Node where
    getAttrs (Element _ attrs _) = attrs
    getAttrs _                   = error "No attributes to be found inside tag text"
    putAttrs attrs (Element name _ children) = Element name attrs children
    putAttrs _ _                 = error "Can't put attributes into tag text"
    
instance Attrable [(String,String)] where
    getAttrs = id
    putAttrs attrs _ = attrs

getAttributes (Element _ eAttrs _) = eAttrs
getAttributes (Text _)             = error "No attributes to be found inside tag text"
setAttribute (name, value) (Element eName eAttrs eChildren) =
        Element eName eAttrs' eChildren
    where
        eAttrs' = (name, value):eAttrs
setAttribute _ (Text _)            = error "No attributes to be found inside tag text"

-- | Convert XML text \<-\> String. Handles empty strings.
xpText0 :: Stringable t => PU_ t String
xpText0 = PU {
        unpickleTree_ = Right . U8.decodeString . getString,
        pickleTree_  = putString . U8.encodeString
    }

-- | Convert XML text \<-\> String. Empty strings result in unpickle failure.
xpText :: Stringable t => PU_ t String
xpText = PU {
        unpickleTree_ = \t ->
            case getString t of
                "" -> Left "empty text"
                txt -> Right (U8.decodeString txt),
        pickleTree_  = putString . U8.encodeString
    }

maybeRead :: Read a => String -> Maybe a
maybeRead s = case reads s of
    [(x, "")] -> Just x
    _         -> Nothing

instance Stringable s => XmlPickler s Int where
    xpickle = xpReadShow

instance Stringable s => XmlPickler s Integer where
    xpickle = xpReadShow

-- | Convert an XML string \<-\> a type that implements Read and Show.
xpReadShow :: (Stringable s, Read n, Show n) => PU_ s n
xpReadShow = PU {
            unpickleTree_ = \t ->
                case maybeRead (getString t) of
                    Just val -> Right val
                    Nothing  -> Left "bad numeric value",
            pickleTree_ = \n -> putString (show n)
        }

instance XmlPickler Node a => XmlPickler Node [a] where
    xpickle = xpList xpickle

-- | Create/parse an XML element of the specified name.  Fails if an element of
-- this name can't be found at this point in the tree.  This implementation unpickles
-- elements of different names in any order, while HXT's xpElem will fail if the
-- XML order doesn't match the Haskell code.
xpElem :: Nodeable t => String -> PU_ Node a -> PU_ t a
xpElem name pu = PU {
        unpickleTree_ = \t ->
            let nodes = getSubnodes t
                doElem elt@(Element eName _ _) | eName == name =
                    Just $ unpickleTree_ pu elt
                doElem _ = Nothing
                mVals   = map doElem nodes
            in  case catMaybes mVals of
                    Right val:_ -> Right val
                    Left err:_  -> Left $ "in <"++name++">, "++err
                    []    -> Left $ "can't find element <"++name++">",
        pickleTree_ = \value nodes ->
            reverseElement (pickleTree_ pu value (Element name [] []))
                `addChild` nodes
    }

-- | Create/parse an XML attribute of the specified name.  Fails if the attribute
-- can't be found at this point in the tree.
xpAttr :: String -> PU_ (String, String) a -> PU_ Node a
xpAttr name pu = PU {
        unpickleTree_ = \t ->
            let attrs = getAttributes t
                doAttr attr@(aName, value) | aName == name =
                    case unpickleTree_ pu attr of
                        Right val -> Just val
                        Left _    -> Nothing
                doAttr _ = Nothing
                mVals = map doAttr attrs
            in  case catMaybes mVals of
                    val:_ -> Right val
                    []    -> Left $ "can't find attribute '"++name++"'",
        pickleTree_ = \value attrs ->
            pickleTree_ pu value (name, "") `setAttribute` attrs
    }

-- | Convert XML text <-> a Maybe type. During unpickling, Nothing is returned
-- if there's a failure during the unpickling of the first argument.
xpOption :: PU_ t a -> PU_ t (Maybe a)
xpOption pu = PU {
        unpickleTree_ = \t ->
            case unpickleTree_ pu t of
                Right val -> Right (Just val)
                Left _    -> Right Nothing,
        pickleTree_ = \mValue t ->
            case mValue of
                Just value -> pickleTree_ pu value t
                Nothing    -> t
    }

-- | Convert XML text \<-\> a 2-tuple using the two arguments.
xpPair :: PU_ t a -> PU_ t b -> PU_ t (a,b)
xpPair pua pub = PU {
        unpickleTree_ = \t ->
            case (unpickleTree_ pua t, unpickleTree_ pub t) of
                (Right a, Right b) -> Right (a,b)
                (Left err, _) -> Left err
                (_, Left err) -> Left err,
        pickleTree_ = \(a, b) t ->
            pickleTree_ pub b $ pickleTree_ pua a t
    }

-- | Convert XML text \<-\> a 3-tuple using the three arguments.
xpTriple :: PU_ t a -> PU_ t b -> PU_ t c -> PU_ t (a,b,c)
xpTriple pua pub puc = PU {
        unpickleTree_ = \t ->
            case (unpickleTree_ pua t, unpickleTree_ pub t, unpickleTree_ puc t) of
                (Right a, Right b, Right c) -> Right (a,b,c)
                (Left err, _, _) -> Left err
                (_, Left err, _) -> Left err
                (_, _, Left err) -> Left err,
        pickleTree_ = \(a, b, c) t ->
            pickleTree_ puc c $ pickleTree_ pub b $ pickleTree_ pua a t
    }

-- | Convert XML text \<-\> a 4-tuple using the four arguments.
xp4Tuple :: PU_ t a -> PU_ t b -> PU_ t c -> PU_ t d -> PU_ t (a,b,c,d)
xp4Tuple pua pub puc pud = PU {
        unpickleTree_ = \t ->
            case (unpickleTree_ pua t, unpickleTree_ pub t, unpickleTree_ puc t, unpickleTree_ pud t) of
                (Right a, Right b, Right c, Right d) -> Right (a,b,c,d)
                (Left err, _, _, _) -> Left err
                (_, Left err, _, _) -> Left err
                (_, _, Left err, _) -> Left err
                (_, _, _, Left err) -> Left err,
        pickleTree_ = \(a, b, c, d) t ->
            pickleTree_ pud d $ pickleTree_ puc c $ pickleTree_ pub b $ pickleTree_ pua a t
    }

-- | Convert XML text \<-\> a 5-tuple using the five arguments.
xp5Tuple :: PU_ t a -> PU_ t b -> PU_ t c -> PU_ t d -> PU_ t e -> PU_ t (a,b,c,d, e)
xp5Tuple pua pub puc pud pue = PU {
        unpickleTree_ = \t ->
            case (unpickleTree_ pua t, unpickleTree_ pub t, unpickleTree_ puc t, unpickleTree_ pud t, unpickleTree_ pue t) of
                (Right a, Right b, Right c, Right d, Right e) -> Right (a,b,c,d,e)
                (Left err, _, _, _, _) -> Left err
                (_, Left err, _, _, _) -> Left err
                (_, _, Left err, _, _) -> Left err
                (_, _, _, Left err, _) -> Left err
                (_, _, _, _, Left err) -> Left err,
        pickleTree_ = \(a, b, c, d, e) t ->
            pickleTree_ pue e $ pickleTree_ pud d $ pickleTree_ puc c $ pickleTree_ pub b $ pickleTree_ pua a t
    }

-- | Convert XML text \<-\> a list of elements. During unpickling, failure of the
-- argument unpickler is the end-of-list condition (and it isn't a failure).
xpList :: PU_ Node a -> PU_ Node [a]
xpList pu = PU {
        unpickleTree_ = \t ->
            let nodes = getSubnodes t
                munge [] out = out
                munge elts@(Element _ _ _:rest) out =
                    case unpickleTree_ pu (Element "" [] [head elts]) of
                        Right val -> munge rest (val:out)
                        Left _    -> out
                munge (Text _:rest) out = munge rest out
            in  Right $ reverse $ munge nodes [],
        pickleTree_ = \list t ->
            foldr (\elt t -> pickleTree_ pu elt t) t (reverse list)
    }

-- | Apply a lens to convert the type of your data structure into something that
-- the pickler primitives can handle (such as tuples).
xpWrap :: (a -> b, b -> a) -> PU_ t a -> PU_ t b
xpWrap (a2b, b2a) pua = PU {
        unpickleTree_ = \t -> case unpickleTree_ pua t of
            Right val -> Right (a2b val)
            Left err  -> Left err,
        pickleTree_ = \value t -> pickleTree_ pua (b2a value) t
    }

-- Like xpWrap, but removes Just (and treats Nothing as a failure) during unpickling.
xpWrapMaybe :: (a -> Maybe b, b -> a) -> PU_ t a -> PU_ t b
xpWrapMaybe (a2b, b2a) pua = PU {
        unpickleTree_ = \t -> case unpickleTree_ pua t of
            Right val ->
                case a2b val of
                    Just val' -> Right val'
                    Nothing   -> Left "xpWrapMaybe can't encode Nothing value"
            Left err  -> Left err,
        pickleTree_ = \value t -> pickleTree_ pua (b2a value) t
    }

-- Like xpWrap, except it removes Right (and treats Left as a failure) during unpickling.
xpWrapEither :: (a -> Either String b, b -> a) -> PU_ t a -> PU_ t b
xpWrapEither (a2b, b2a) pua = PU {
        unpickleTree_ = \t -> case unpickleTree_ pua t of
            Right val -> a2b val
            Left err  -> Left err,
        pickleTree_ = \value t -> pickleTree_ pua (b2a value) t
    }

-- Convert an attribute list in the XML tree into [(String, String)]. (Does not
-- exist in HXT.)
xpAllAttrs :: Attrable t => PU_ t [(String, String)]
xpAllAttrs = PU {
        unpickleTree_ = \t -> Right (getAttrs t),
        pickleTree_  = putAttrs
    }

-- | Allow alternative picklers. Selector function is used during pickling, but
-- unpickling is done by trying each list element in order until one succeeds.
xpAlt :: (a -> Int)  -- ^ Selector
      -> [PU_ t a]
      -> PU_ t a
xpAlt selector picklers = PU {
        unpickleTree_ = \t ->
            let tryAll [] = Left "all xpAlt unpickles failed"
                tryAll (x:xs) =
                    case unpickleTree_ x t of
                        Right val -> Right val
                        Left err  -> tryAll xs
            in  tryAll picklers,
        pickleTree_ = \value t -> pickleTree_ (picklers !! (selector value)) value t
    }

-- | Convert nothing \<-\> (). Does not output or consume any XML text. 
xpUnit :: PU_ t ()
xpUnit = PU {
        unpickleTree_ = \t -> Right (),
        pickleTree_ = \_ t -> t
    }