module Bergamot.Rules exposing (..)

import Bergamot.Syntax exposing (Term(..), Metavariable, UnificationVar, unify, emptyUnificationState, instantiate, instantiateList, emptyInstantiationState, resetVars, InstantiationState, UnificationState, reify)
import Bergamot.Search as Search exposing (Search)
import Bergamot.Utils exposing (encodeStr, encodeLatex)

import Debug

type alias Rule =
    { name : String
    , conclusion : Term Metavariable
    , premises : List (Term Metavariable)
    }

type ProofTree = MkProofTree
    { name : String
    , conclusion : Term UnificationVar
    , premises : List ProofTree
    }

type alias RuleEnv =
    { rules : List Rule
    }

type alias ProveState =
    { unificationState : UnificationState
    , instantiationState : InstantiationState
    , gas: Int
    }

type alias Prover a = RuleEnv -> ProveState -> Search (a, ProveState)

andThen : (a -> Prover b) -> Prover a -> Prover b
andThen f p env ps =
    p env ps
    |> Search.andThen (\(a, psp) -> if psp.gas > 0 then (f a) env psp else Search.fail)

join : Prover (Prover a) -> Prover a
join p = andThen (\x -> x) p

apply : Prover a -> Prover (a -> b) -> Prover b
apply pa pf = pf |> andThen (\f -> map f pa)

map : (a -> b) -> Prover a -> Prover b
map f p env ps =
    p env ps
    |> Search.map (Tuple.mapFirst f)

mapM : (a -> Prover b) -> List a -> Prover (List b)
mapM f l =
    case l of
        x :: xs ->
            pure (::)
            |> apply (f x)
            |> apply (mapM f xs)
        [] -> pure []

interleave : Prover a -> Prover a -> Prover a
interleave p1 p2 env ps =
    Search.interleave (p1 env ps) (p2 env ps)

pure : a -> Prover a
pure a env ps = Search.pure (a, ps)

fail : Prover a
fail env ps = Search.fail

yield : Prover a -> Prover a
yield p env ps = Search.yield (p env ps)

getEnv : Prover RuleEnv
getEnv env ps = Search.pure (env, ps)

getUnificationState : Prover UnificationState
getUnificationState env ps = Search.pure (ps.unificationState, ps)

burn : Prover ()
burn env ps = Search.pure ((), { ps | gas = ps.gas - 1})

liftInstantiation : (a -> InstantiationState -> (b, InstantiationState)) -> a -> Prover b
liftInstantiation f a env ps =
    let
        (b, is) = f a ps.instantiationState
    in
        Search.pure (b, { ps | instantiationState = is })

liftUnification : (a -> a -> UnificationState -> Maybe (a, UnificationState)) -> a -> a -> Prover a
liftUnification f a1 a2 env ps =
    case f a1 a2 ps.unificationState of
        Just (a, us) -> Search.pure (a, { ps | unificationState = us })
        Nothing -> Search.fail

withVars : Prover a -> Prover a
withVars p env ps =
    p env ps
    |> Search.map (Tuple.mapSecond (\psp -> { psp | instantiationState = resetVars psp.instantiationState }))


reifyProofTree : ProofTree -> Prover ProofTree
reifyProofTree (MkProofTree node) =
    pure (\t trees -> MkProofTree { node | conclusion = t, premises = trees })
    |> apply (map (reify node.conclusion) getUnificationState)
    |> apply (mapM reifyProofTree node.premises)

rule : Term UnificationVar -> Rule -> Prover ProofTree
rule t r =
        withVars (pure Tuple.pair
            |> apply (liftInstantiation instantiate r.conclusion
                |> andThen (\conc -> liftUnification unify t conc))
            |> apply (liftInstantiation instantiateList r.premises))
        |> andThen (\(conc, prems) ->
            provePremises prems)
            |> map (\trees -> MkProofTree { name = r.name, conclusion = t, premises = trees })

collectStrings : Term UnificationVar -> Prover (List String)
collectStrings t =
    case t of
        Call "cons" [StringLit s, rest] -> map (\ss -> s :: ss) <| collectStrings rest
        Call "nil" [] -> pure []
        _ -> fail

builtinRules : Term UnificationVar -> Prover ProofTree
builtinRules t =
    case t of
        Call "concat" [StringLit s1, StringLit s2, Var output] ->
            liftUnification unify (Var output) (StringLit (s1 ++ s2))
            |> map (\_ -> MkProofTree { name = "BuiltinConcat", conclusion = t, premises = []})
        Call "join" [tp, Var output] ->
            collectStrings tp
            |> andThen (\ss -> liftUnification unify (Var output) (StringLit (String.concat ss)))
            |> map (\_ -> MkProofTree { name = "BuiltinJoin", conclusion = t, premises = []})
        Call "int" [IntLit i] ->
            MkProofTree { name = "BuiltinInt", conclusion = t, premises = [] }
            |> pure
        Call "str" [StringLit s] ->
            MkProofTree { name = "BuiltinStr", conclusion = t, premises = [] }
            |> pure
        Call "sym" [Call s []] ->
            MkProofTree { name = "BuiltinSym", conclusion = t, premises = [] }
            |> pure
        Call "call" [Call s ts, Var name, Var args] ->
            pure (\_ _ -> MkProofTree { name = "BuiltinCall", conclusion = t, premises = [] })
            |> apply (liftUnification unify (Var name) (StringLit <| encodeStr s))
            |> apply (liftUnification unify (Var args) (List.foldr (\x xs -> Call "cons" [x, xs]) (Call "nil" []) ts))
        Call "tostring" [IntLit i, Var output] ->
            liftUnification unify (Var output) (StringLit (String.fromInt i))
            |> map (\_ -> MkProofTree { name = "BuiltinToString", conclusion = t, premises = []})
        Call "tostring" [Call s [], Var output] ->
            liftUnification unify (Var output) (StringLit <| encodeStr s)
            |> map (\_ -> MkProofTree { name = "BuiltinToString", conclusion = t, premises = []})
        Call "escapestring" [StringLit s, Var output] ->
            liftUnification unify (Var output) (StringLit (encodeStr s))
            |> map (\_ -> MkProofTree { name = "BuiltinEscapeString", conclusion = t, premises = []})
        Call "latexifystring" [StringLit s, Var output] ->
            liftUnification unify (Var output) (StringLit (encodeLatex s))
            |> map (\_ -> MkProofTree { name = "BuiltinLatexifyeString", conclusion = t, premises = []})
        _ -> fail

provePremises : List (Term UnificationVar) -> Prover (List ProofTree)
provePremises = mapM proveTerm

proveTerm : Term UnificationVar -> Prover ProofTree
proveTerm t =
    map (reify t) getUnificationState
    |> andThen (\tp ->
        burn
        |> andThen (\() -> getEnv)
        |> andThen (\env -> List.foldl (\r -> interleave (yield (rule tp r))) (builtinRules tp) env.rules))

prove : Term Metavariable -> Prover ProofTree
prove mt =
    withVars (liftInstantiation instantiate mt)
    |> andThen proveTerm

single : RuleEnv -> Prover a -> Maybe a
single env p =
    p env { instantiationState = emptyInstantiationState, unificationState = emptyUnificationState, gas = 1000 }
    |> Search.one
    |> Maybe.map (Tuple.first << Tuple.first)