module Merkle
    exposing
        ( Tree
        , initialize
        , singleton
        , fromList
        , insert
        , insertFromList
        , contains
        , get
        , flatten
        , depth
        , isValid
        , toJson
        , fromJson
        )

{-| This is an implementation of the [Merkle
tree](https://en.wikipedia.org/wiki/Merkle_tree) data structure in Elm. It's
implemented as an immutable balanced binary hash tree, which guarantees for
logarithmic inserts. Default hash function is
[**SHA-2**56](https://en.wikipedia.org/wiki/SHA-2#Comparison_of_SHA_functions)
but others can be used. Hash functions are specified on initiation and can't be
changed afterwards which ensures data consistency.

# Trees

@docs Tree

# Creating Trees

@docs initialize, singleton, fromList

# Basics

@docs insert, insertFromList, get, depth, flatten, contains

# Check for Consistency

@docs isValid

# Persistence

@docs fromJson,toJson

-}

-- ELM-LANG LIBS

import Basics exposing (max)
import Bitwise exposing (and, shiftLeftBy, shiftRightZfBy)
import Json.Decode as JsonD
    exposing
        ( Decoder
        , andThen
        , decodeString
        , decodeValue
        , field
        , int
        , maybe
        , map
        , map2
        , map3
        , oneOf
        , string
        , value
        )
import Json.Encode as JsonE
    exposing
        ( Value
        , encode
        , int
        , null
        , object
        , string
        )
import List exposing (foldl)
import String exposing (left)


-- OTHER LIBS

import SHA exposing (sha256sum)


{-| Representation of the immutable balanced binary hash tree. You can create
Merkle trees of integers (`Merkle.Tree Int`) or strings (`Merkle.Tree String`)
or any other type of values supported by Elm (Ints, Floats, Bools, Strings,
Maybes, Lists, Arrays, Tuples, Json.Values or concrete records).
-}
type Tree a
    = Tree (Config a)


type alias Config a =
    { hashfs : List (String -> String)
    , encode : a -> JsonE.Value
    , bintree : BinaryTree a
    }


type Tuple a b
    = Tuple a b


type alias Hash =
    String


type BinaryTree a
    = Leaf (Maybe (Tuple a Hash))
    | Branch (Tuple Int Hash) (BinaryTree a) (BinaryTree a)


{-| Initialize an empty (`Merkle.Tree a`) that is bounded to the type assigned
in the hash functions (`hashfs`) and JSON encoder (`encode`). If no hash
functions are assigned or an empty list is given, then it will default to
`sha256sum`.

    initialize Nothing Json.Encode.int
    initialize (Just ([ sha256sum ])) Json.Encode.int
    initialize (Just ([ identity ])) Json.Encode.int
-}
initialize :
    Maybe (List (String -> String))
    -> (a -> Value)
    -> Tree a
initialize hashfs encode =
    Tree
        { hashfs = ensure hashfs
        , encode = encode
        , bintree = Leaf Nothing
        }


{-| Creates a (`Merkle.Tree a`) with the single element that is given which is
bounded to the type assigned in the hash functions (`hashfs`) and JSON encoder
(`encode`). If no hash functions are assigned or an empty list is given, then it
will default to `sha256sum`.

    singleton 42 Nothing Json.Encode.int
    singleton 42 (Just ([ sha256sum ])) Json.Encode.int
    singleton 42 (Just ([ identity ])) Json.Encode.int
-}
singleton :
    a
    -> Maybe (List (String -> String))
    -> (a -> Value)
    -> Tree a
singleton x hashfs encode =
    let
        bt =
            initbranch 0 x hashfs encode
    in
        init bt hashfs encode


{-| Creates a (`Merkle.Tree a`) from the elements in a list which are bounded to
the type assigned in the hash functions (`hashfs`) and JSON encoder
(`encode`). If no hash functions are assigned or an empty list is given, then it
will default to `sha256sum`. The order of the elements is significant, therefore
the first element in the list will be the first to be inserted in the
`Merkle.Tree`. This is an important property in order to maintain consistency
and to be able to recreate a tree.

    [42] |> fromList Nothing Json.Encode.int
    [42] |> fromList (Just ([ sha256sum ])) Json.Encode.int
    [42] |> fromList (Just ([ identity ])) Json.Encode.int
-}
fromList :
    Maybe (List (String -> String))
    -> (a -> Value)
    -> List a
    -> Tree a
fromList hashfs encode list =
    let
        init =
            initialize hashfs encode
    in
        list
            |> List.foldl (\x t -> t |> insert x) init


{-| Adds an element to the (`Merkle.Tree a`) in logarithmic time. It's important
to understand that this implementation of `Merkle.Tree` will ensure that all
elements have the same `depth` in the tree (in the same way as the
**Blockchain** for **BTC** does it). For more information, please see the
examples on GitHub.

    initialize Nothing Json.Encode.int
        |> insert 42
-}
insert : a -> Tree a -> Tree a
insert x (Tree ({ hashfs, encode, bintree } as config)) =
    let
        m =
            bintree |> max

        hlp bt m i hashfs encode =
            case bt of
                Leaf _ ->
                    Debug.crash "Don't call hlp with a Leaf"

                Branch (Tuple n h) lb rb ->
                    case n > 1 && n == m of
                        True ->
                            ( bt
                            , initbranch (log2 n) x (Just hashfs) encode
                            )

                        False ->
                            let
                                ln =
                                    case lb of
                                        Leaf _ ->
                                            1

                                        Branch (Tuple n_ _) _ _ ->
                                            n_
                            in
                                case ln < i of
                                    True ->
                                        ( lb |> rec (i // 2)
                                        , rb
                                        )

                                    False ->
                                        case n > 1 && (n == i) of
                                            True ->
                                                ( lb
                                                , initbranch (log2 n) x (Just hashfs) encode
                                                )

                                            False ->
                                                ( lb
                                                , rb |> rec (i // 2)
                                                )

        rec i bt =
            case bt of
                Leaf Nothing ->
                    let
                        h =
                            hashfs |> combine (stringify x encode)
                    in
                        (Tuple x h) |> Just |> Leaf

                (Leaf (Just (Tuple _ h))) as leaf ->
                    let
                        hx =
                            hashfs |> combine (stringify x encode)

                        h_ =
                            hashfs |> combine (h ++ hx)
                    in
                        Branch (Tuple 2 h_) leaf ((Tuple x hx) |> Just |> Leaf)

                (Branch (Tuple n h) lb rb) as branch ->
                    let
                        ( lb_, rb_ ) =
                            hlp branch m i hashfs encode

                        h_ =
                            hashHlp h lb_ rb_ hashfs
                    in
                        Branch (Tuple (n + 1) h_) lb_ rb_
    in
        Tree
            { config
                | bintree = rec ((bintree |> max) // 2) bintree
            }


{-| Adds elements from the list to the (`Merkle.Tree a`). The order of the
elements is significant, therefore the first element in the list will be the
first to be inserted in the `Merkle.Tree`. This is an important property in
order to maintain consistency and to be able to recreate a tree.

    [0 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)

-}
insertFromList : Tree a -> List a -> Tree a
insertFromList tree list =
    list
        |> List.foldl (\x a -> a |> insert x) tree


{-| Checks if the element is in the (`Merkle.Tree a`). As an element can appear
several times in a tree, therefore if we find early that the element is in the
tree we stop the recursion down this path.

    ([0 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> contains 42) == True

    ([0 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> contains 43) == False
-}
contains : a -> Tree a -> Bool
contains x (Tree ({ hashfs, encode, bintree } as config)) =
    let
        rec bt =
            case bt of
                Leaf Nothing ->
                    False

                Leaf (Just (Tuple data hash)) ->
                    x == data

                Branch _ lb rb ->
                    (rec lb) || (rec rb)
    in
        rec bintree


{-| Returns the element(s) together with its/their `hash`(es) if any. Otherwise
it's returning an empty list. As an element can appear several times in a tree
and order of the elements is significant, therefore as the first element
inserted in the `Merkle.Tree` will be the first element of the list.

    ([0 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> get 42) == [ (42, "Some hash") ]

    ([0 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> get 43) == []

    ([42, 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> get 42) == [ (42, "Some hash"), (42, "Same hash")  ]
-}
get : a -> Tree a -> List ( a, String )
get x (Tree ({ hashfs, encode, bintree } as config)) =
    let
        rec bt =
            case bt of
                Leaf Nothing ->
                    []

                Leaf (Just (Tuple data hash)) ->
                    case x == data of
                        True ->
                            [ ( data, hash ) ]

                        False ->
                            []

                Branch _ lb rb ->
                    (rec lb) ++ (rec rb)
    in
        rec bintree


{-| Returns all the elements together with their `hash`(es) if any. Otherwise it's
returning an empty list. As an element can appear several times in a tree and
order of the elements is significant, therefore as the first element inserted in
the `Merkle.Tree` will be the first element of the list.

    ([0 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> flatten) == [ (0, "Some hash"), ..., (42, "Some hash") ]
-}
flatten : Tree a -> List ( a, String )
flatten (Tree ({ hashfs, encode, bintree } as config)) =
    let
        rec bt =
            case bt of
                Leaf Nothing ->
                    []

                Leaf (Just (Tuple data hash)) ->
                    [ ( data, hash ) ]

                Branch _ lb rb ->
                    (rec lb) ++ (rec rb)
    in
        rec bintree


{-| Return the `depth` of the `Merkle.Tree`. We can calculate the max amount of
elements in the tree by power-of-two to the depth of the tree (`depth = 3` and
`2^3`) as we know that tree is logarithmic in `depth`.

    ([0 .. 7]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> depth) == 3

    ([0 .. 8]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> depth) == 4
-}
depth : Tree a -> Int
depth (Tree ({ hashfs, encode, bintree } as config)) =
    let
        rec bt =
            case bt of
                Leaf _ ->
                    0

                Branch _ lb rb ->
                    1 + Basics.max (rec lb) (rec rb)
    in
        rec bintree


{-| To ensure data consistency and correctness, we can at any given point check
if the root `hash` is actually based on all the underlying elements. This
function is important to use whenever trees are imported `fromJson` as `hash`es
can be modified in the text format.

    ([0 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> isValid Nothing) == True

    ([0 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> isValid (Just [ identity ])) == False

**Note**: If trees imported `fromJson` have both modified `data` and `hash`es
  consistently, it will not be possible for this function to ensure historic
  consistency but only correctness.
-}
isValid : Maybe (List (String -> String)) -> Tree a -> Bool
isValid hash (Tree ({ hashfs, encode, bintree } as config)) =
    let
        hashfs_ =
            case hash of
                Just _ ->
                    ensure hash

                Nothing ->
                    hashfs

        rec bt =
            case bt of
                Leaf Nothing ->
                    True

                Leaf (Just (Tuple x h)) ->
                    let
                        h_ =
                            hashfs_ |> combine (stringify x encode)
                    in
                        h == h_

                Branch (Tuple n h) lb rb ->
                    let
                        h_ =
                            hashHlp h lb rb hashfs_
                    in
                        (h == h_) && (rec lb) && (rec rb)
    in
        rec bintree


{-| Converts a `Merkle Tree a` into a prettified `JSON` string. The first
argument specifies the amount of indentation in the resulting string while the
second specifies a hash mask to ensure better readability. For compact, use `0`
as indentation (no line break) and for readability use `4` (with line
breaks). For `Git` alike `hash` mask use `7` and 0 for full `hash`.

    ([40 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> toJson 0 7)

Becomes:

    {
        "count": 3,
        "hash": "37536cf",
        "left": {
            "count": 2,
            "hash": "ae41226",
            "left": {
                "hash": "d59eced",
                "data": 40
            },
            "right": {
                "hash": "3d914f9",
                "data": 41
            }
        },
        "right": {
            "count": 1,
            "hash": "c6e4fe6",
            "left": {
                "hash": "73475cb",
                "data": 42
            },
            "right": null
        }
    }

**Note**: If data is needed to persist, don't add any mask and ensure that full
  length is saved. It's important in order to check for data consistency
  (`isValid`).
-}
toJson : Int -> Int -> Tree a -> String
toJson indentation mask (Tree { hashfs, encode, bintree }) =
    let
        shortHash h =
            case mask of
                0 ->
                    h

                _ ->
                    String.left mask h

        rec bt =
            case bt of
                Leaf Nothing ->
                    JsonE.null

                Leaf (Just (Tuple v h)) ->
                    JsonE.object
                        [ ( "hash", JsonE.string (shortHash h) )
                        , ( "data", encode v )
                        ]

                Branch (Tuple n_ h) lb rb ->
                    JsonE.object
                        [ ( "count", JsonE.int n_ )
                        , ( "hash", JsonE.string (shortHash h) )
                        , ( "left", rec lb )
                        , ( "right", rec rb )
                        ]
    in
        rec bintree |> JsonE.encode indentation


{-| Re-creates a `Merkle Tree a` from a prettified `JSON` string. It's important
to use the same initial parameters as the elements are bounded to the type
assigned in the hash functions (`hashfs`) and JSON encoder (`encode`). If any of
these parameters defer from the initial, the tree will neither be consistent nor
correct (`isValid`).

    (case ([40 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> toJson 0 0
        |> fromJson Nothing Json.Encode.int Json.Decode.int) of
            Ok t ->
                t |> isValid Nothing
            Err msg ->
                Debug.crash msg) == True

    (case ([40 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> toJson 0 7
        |> fromJson Nothing Json.Encode.int Json.Decode.int) of
            Ok t ->
                t |> isValid Nothing
            Err msg ->
                Debug.crash msg) == False

    (case ([40 .. 42]
        |> insertFromList
            (initialize Nothing Json.Encode.int)
        |> toJson 0 0
        |> fromJson (Just [ identity ]) Json.Encode.int Json.Decode.int) of
            Ok t ->
                t |> isValid Nothing
            Err msg ->
                Debug.crash msg) == False
-}
fromJson :
    Maybe (List (String -> String))
    -> (a -> Value)
    -> Decoder a
    -> String
    -> Result String (Tree a)
fromJson hash encode decode json =
    let
        result =
            json
                |> JsonD.decodeString (decodeTree decode)
    in
        case result of
            Ok bt ->
                Ok
                    (Tree
                        { hashfs = ensure hash
                        , encode = encode {- , decode = decode -}
                        , bintree = bt
                        }
                    )

            Err msg ->
                Err msg



-- HELPERS


initbranch :
    Int
    -> a
    -> Maybe (List (String -> String))
    -> (a -> Value)
    -> BinaryTree a
initbranch n x hash encode =
    let
        hl =
            (ensure hash) |> combine (stringify x encode)

        hb h =
            (ensure hash) |> combine (h ++ h)

        rec i h acc =
            case i of
                0 ->
                    acc

                _ ->
                    let
                        h_ =
                            hb h

                        t =
                            Branch (Tuple 1 h_) acc (Nothing |> Leaf)
                    in
                        rec (i - 1) h_ t
    in
        rec n hl ((Tuple x hl) |> Just |> Leaf)


init :
    BinaryTree a
    -> Maybe (List (String -> String))
    -> (a -> Value)
    -> Tree a
init bt hash encode =
    Tree
        { hashfs = ensure hash
        , encode = encode
        , bintree = bt
        }


ensure : Maybe (List (String -> String)) -> List (String -> String)
ensure hash =
    case hash of
        Nothing ->
            [ sha256sum ]

        Just xs ->
            case xs of
                [] ->
                    [ sha256sum ]

                _ ->
                    xs


stringify : a -> (a -> Value) -> String
stringify x encode =
    x |> encode |> JsonE.encode 0


combine : a -> List (a -> String) -> String
combine x =
    List.foldl (\f a -> a ++ (f x)) ""


max : BinaryTree a -> Int
max t =
    case t of
        Leaf Nothing ->
            1

        Leaf (Just _) ->
            2

        Branch (Tuple n _) _ _ ->
            n |> ceilPow


hashHlp :
    a
    -> BinaryTree b
    -> BinaryTree c
    -> List (String -> String)
    -> String
hashHlp h lh rh fs =
    case ( lh, rh ) of
        ( Leaf (Just (Tuple _ h1)), Leaf (Just (Tuple _ h2)) ) ->
            fs |> combine (h1 ++ h2)

        ( Leaf (Just (Tuple _ h1)), Leaf Nothing ) ->
            fs |> combine (h1 ++ h1)

        ( Branch (Tuple _ h1) _ _, Branch (Tuple _ h2) _ _ ) ->
            fs |> combine (h1 ++ h2)

        ( Branch (Tuple _ h1) _ _, Leaf (Just (Tuple _ h2)) ) ->
            fs |> combine (h1 ++ h2)

        ( Branch (Tuple _ h1) _ _, Leaf Nothing ) ->
            fs |> combine (h1 ++ h1)

        ( Leaf (Just _), Branch _ _ _ ) ->
            Debug.crash "It shouldn't be possible to reach this branch"

        ( Leaf Nothing, _ ) ->
            Debug.crash "It shouldn't be possible to reach this branch"


log2 : Int -> Int
log2 x =
    logBase 2.0 (x |> toFloat) |> truncate


ceilPow : Int -> Int
ceilPow n =
    let
        rec ( i, acc ) =
            case ( i, acc ) of
                ( 0, a ) ->
                    1 |> Bitwise.shiftLeftBy a

                ( m, a ) ->
                    let
                        m_ =
                            m |> Bitwise.and 1
                    in
                        ( (m - m_) |> Bitwise.shiftRightZfBy 1, a + 1 ) |> rec
    in
        rec ( n - 1, 0 )


decodeTree : Decoder a -> Decoder (BinaryTree a)
decodeTree decode =
    let
        decodeLeaf =
            JsonD.map Leaf
                (JsonD.maybe
                    (JsonD.map2 Tuple
                        (field "data" decode)
                        (field "hash" JsonD.string)
                    )
                )

        decodeBranch =
            JsonD.map3 Branch
                (JsonD.map2 Tuple
                    (field "count" JsonD.int)
                    (field "hash" JsonD.string)
                )
                (field "left" (decodeLazy (\_ -> decodeTree decode)))
                (field "right" (decodeLazy (\_ -> decodeTree decode)))
    in
        JsonD.oneOf [ decodeBranch, decodeLeaf ]


decodeLazy :
    (() -> JsonD.Decoder a)
    -> JsonD.Decoder a
decodeLazy thunk =
    customDecoder JsonD.value
        (\json -> JsonD.decodeValue (thunk ()) json)


customDecoder decoder toResult =
    JsonD.andThen
        (\a ->
            case toResult a of
                Ok b ->
                    JsonD.succeed b

                Err err ->
                    JsonD.fail err
        )
        decoder