let prelude = import! std.prelude
let { Functor, Applicative, Monad } = prelude
let { map } = import! std.functor
let { Foldable } = import! std.foldable
let { traverse } = import! std.traversable
let { id, (<|) } = import! std.function
let { ? } = import! std.array

let string = import! std.string
let { (<>) } = import! std.prelude

let { Bool } = import! std.bool
let float = import! std.float
let int = import! std.int
let result @ { Result, ? } = import! std.result
let { Option } = import! std.option
let { Eff, run_pure, ? } = import! std.effect
let { get, put, modify } = import! std.effect.state
let { throw, catch, run_error } = import! std.effect.error
let { eval_state, run_state } = import! std.effect.state

let list @ { List, ? } = import! std.list

let std_map @ { Map, ? } = import! std.map

let { Expr, Function, LispEffect, LispState } = import! "examples/lisp/types.glu"
let lisp_parser = import! "examples/lisp/parser.glu"
let parser = import! std.parser

let eq : Eq Expr =
    let eq_expr l r : Expr -> Expr -> Bool =
        match (l, r) with
        | (Atom ls, Atom rs) -> ls == rs
        | (Int ls, Int rs) -> ls == rs
        | (Float ls, Float rs) -> ls == rs
        | (List ls, List rs) ->
            let list_eq : Eq (List Expr) = list.eq ?{ (==) = eq_expr }
            list_eq.(==) ls rs
        | _ -> False
    { (==) = eq_expr }

let show_expr : Show Expr =
    rec
    let spaced show_ xs =
        match xs with
        | Cons y ys -> show_ y <> spaced1 show_ ys
        | Nil -> ""
    let spaced1 show_ xs =
        match xs with
        | Cons y ys -> " " <> show_ y <> spaced1 show_ ys
        | Nil -> ""
    in

    let show expr =
        match expr with
        | Atom s -> s
        | Int i -> int.show.show i
        | Float f -> float.show.show f
        | List ls -> "(" <> spaced show ls <> ")"
        | Function f ->
            let vararg =
                match f.vararg with
                | Some arg -> " . " <> arg
                | None -> ""
            "(lambda (" <> spaced id f.params <> ")" <> vararg
                <> ") ...)"
        | Primitive _ -> "<primitive>"
    { show }

let { wrap } = import! std.applicative
let { flat_map, (>>=) } = import! std.monad
let { fold_m } = import! std.foldable

let scope_state run : Eff (LispEffect r) a -> Eff (LispEffect r) a =
    do original = get
    do x = run
    seq put original
    wrap x

let primitive name f : String -> _ -> Map String Expr = std_map.singleton name (Primitive f)

type Binop a = a -> a -> a

let primitive_binop name int_op float_op : _ -> Binop Int -> Binop Float -> Map String Expr =
    let unpack_int x : Expr -> Eff (LispEffect r) Int =
        match x with
        | Int i -> wrap i
        | _ -> throw "Expected integer"
    let unpack_float x : Expr -> Eff (LispEffect r) Float =
        match x with
        | Float f -> wrap f
        | _ -> throw "Expected float"

    let fold unpack op : (Expr -> Eff (LispEffect r) a)
            -> Binop a
            -> a
            -> List Expr
            -> _
        = fold_m (\acc x -> map (\y -> op acc y) (unpack x))

    let f xs : List Expr -> Eff (LispEffect r) Expr =
        match xs with
        | Cons l ys ->
            match l with
            | Int li -> map Int (fold unpack_int int_op li ys)
            | Float lf -> map Float (fold unpack_float float_op lf ys)
            | _ -> throw ("Cant add " <> show l)
        | _ -> throw ("Expected two arguments to binop, got " <> show (List xs))
    primitive name f

let define xs =
    match xs with
    | Cons (Atom name) (Cons value Nil) ->
        do state = get
        let new_state = std_map.insert name value state
        seq put new_state
        wrap value
    | Cons (List (Cons (Atom name) params)) body ->
        do closure = get

        let function = Function {
                    params = map show params,
                    vararg = None,
                    body,
                    closure,
                }
        let new_state = std_map.insert name function closure

        seq put new_state

        wrap function
    | _ -> throw "Unexpected parameters to define `define`"

let primitives : LispState =
    let { concat } = import! std.foldable
    concat [
        primitive_binop "+" (+) (+),
        primitive_binop "-" (-) (-),
        primitive_binop "*" (*) (*),
        primitive_binop "/" (/) (/),
        primitive "define" define,
    ]

rec
let apply f xs : forall r . Expr -> List Expr -> Eff (LispEffect r) Expr =
    let add_args names values : List String -> _ =
        match (names, values) with
        | (Cons name names, Cons value values) ->
            seq modify (\state -> std_map.insert name value state)
            add_args names values
        | (Nil, _) -> wrap ()
        | _ -> throw "Not enough arguments to function"

    match f with
    | Primitive primitive -> primitive xs
    | Function function ->
        scope_state (
            seq add_args function.params xs
            eval_exprs function.body)
    | _ -> throw ("Can\'t call value: " <> show f)

let eval_lisp expr : Expr -> Eff (LispEffect r) Expr =
    match expr with
    | Atom name ->
        do state = get
        match std_map.find name state with
        | Some value -> wrap value
        | None -> throw ("Binding `" <> name <> "` is not defined")
    | Int _ -> wrap expr
    | Float _ -> wrap expr
    | Function _ -> wrap expr
    | List list ->
        match list with
        | Cons x xs ->
            match x with
            | Atom name ->
                do state = get
                if name == "define" then
                    define xs
                else
                    match std_map.find name state with
                    | Some prim ->
                        do evaluated_args = traverse eval_lisp xs
                        apply prim evaluated_args
                    | None -> throw ("Variable `" <> name <> "` does not exist")
            | _ -> throw ("Cant apply " <> show x)
        | Nil -> wrap expr

let eval_exprs exprs = fold_m (\_result expr -> eval_lisp expr) (List Nil) exprs
in

let eval_env expr env : Eff (LispEffect r) a -> Map String Expr -> Result String a =
    run_pure <| run_error <| eval_state env expr

let eval expr : Expr -> Result String Expr = eval_env (eval_lisp expr) primitives
let eval_seq exprs =
    match exprs with
    | Cons _ _ -> eval_env (eval_exprs exprs) primitives
    | Nil -> error "Expected at least one lisp expression"

let eval_string s = parser.parse lisp_parser.expr s >>= eval
let eval_env_string s env =
    do e = parser.parse lisp_parser.expr s
    do l = run_pure <| run_error <| run_state env <| eval_lisp e
    wrap (l.value, l.state)

{
    LispState,
    Expr,
    eq,
    show = show_expr,
    expr = lisp_parser.expr,

    default_env = primitives,

    eval,
    eval_seq,
    eval_string,
    eval_env_string,
}