(*
    1 2 + print-int
*)

let rec skip_whitespace chars =
    let open Seq in
    match chars () with
        | Cons (' ', rest) | Cons ('\t', rest) | Cons ('\n', rest) ->
            skip_whitespace rest
        | _ ->
            chars

let collect_symbol chars =
    let rec helper chars acc =
        let open Seq in
        match chars () with
            | Cons (' ', rest) | Cons ('\t', rest) | Cons ('\n', rest) ->
                acc |> Buffer.contents, rest
            | Cons (ch, rest) ->
                Buffer.add_char acc ch ;
                helper rest acc
            | Nil ->
                (acc |> Buffer.contents, Seq.empty) in
    helper chars @@ Buffer.create 16

let scan_symbol chars =
    let chars = skip_whitespace chars in
    collect_symbol chars

let tokenize str =
    let stream = String.to_seq str in
    let rec helper stream acc =
        let open Seq in
        match scan_symbol stream with
            | "", seq ->
                acc |> List.rev
            | str, more ->
                helper more (str :: acc) in
    helper stream []

type joyful_quote = [ `JoyfulQuotation of joyful_action list ]
and joyful_data =
    [
    | `JoyfulSymbol of string
    | `JoyfulInt of int
    | `JoyfulQuotation of joyful_action list
    ]
and joyful_action =
    | Push of joyful_data
    | Call of string

(* todo, implmenent stackless qoutation eval *)
exception Unclosed_quotation

let parse tokens =
    let to_action token =
        match int_of_string_opt token with
            | Some n ->
                Push (`JoyfulInt n)
            | None when String.starts_with ~prefix:"`" token ->
                Push (`JoyfulSymbol (String.sub token 1 (String.length token - 1)))
            | None ->
                Call token in

    let rec parse_quotation ?(quotation=[]) tokens =
        match tokens with
            | [] ->
                raise Unclosed_quotation
            | "[" :: rest ->
                let rest, inner_quotation = parse_quotation rest in
                parse_quotation rest ~quotation:(Push inner_quotation :: quotation)
            | "]" :: rest ->
                rest, `JoyfulQuotation (List.rev quotation)
            | token :: rest ->
                let quotation = to_action token  :: quotation in
                parse_quotation rest ~quotation in

    let rec helper tokens prog =
        match tokens with
            | [] ->
                `JoyfulQuotation (List.rev prog)
            | "[" :: rest ->
               let rest, quotation = parse_quotation rest in
               helper rest (Push (quotation) :: prog)
            | token :: rest ->
                to_action token :: prog |> helper rest in

    helper tokens []

class runtime_class =
    let transfer sa sb = Stack.push (Stack.pop sa) sb in
    object(self)
        val data_stack : joyful_data Stack.t = Stack.create ()
        val retain_stack : joyful_data Stack.t = Stack.create ()
        val call_stack : joyful_action Stack.t = Stack.create ()

        method stash = transfer data_stack retain_stack
        method unstash = transfer retain_stack data_stack
        method push = Fun.flip Stack.push data_stack
        method pop = Stack.pop data_stack

        method pop_int =
            match self#pop with
                | `JoyfulInt i -> i
                | _ -> raise @@ Invalid_argument "Expected int"

        method pop_quote =
            match self#pop with
                | `JoyfulQuotation q -> q
                | _ -> raise @@ Invalid_argument "Expected quote"

        method pop_symbol =
            match self#pop with
                | `JoyfulSymbol s -> s
                | _ -> raise @@ Invalid_argument "Expected symbol"

        method dup = Stack.push (Stack.top data_stack) data_stack

        method swap =
            let tos = self#pop in
            let nos = self#pop in
            self#push tos ;
            self#push nos

        method drop = Stack.drop data_stack

        method queue_quotation (quote : joyful_quote) =
            match quote with
                | `JoyfulQuotation actions -> self#queue_actions actions

        method queue_actions actions =
            actions |> List.rev |> List.to_seq |> Stack.add_seq call_stack

        method poll =
            Stack.pop_opt call_stack

        method data_stack = data_stack
        method retain_stack = retain_stack
        method call_stack = call_stack
    end

let dictionary =
    object(self)
        val words : (string, runtime_class -> unit) Hashtbl.t = Hashtbl.create 16

        method add_word name f =
            Hashtbl.add words name f

        method remove_word name =
            Hashtbl.remove words name

        method call name runtime =
            (Hashtbl.find words name) runtime
    end


let () =
    let new_math_primitives name op =
        dictionary#add_word name (fun runtime ->
            let y = runtime#pop_int in
            let x = runtime#pop_int in
            let z = `JoyfulInt (op x y) in
            runtime#push z
        ) in

    List.iter2 new_math_primitives
        ["+" ; "-" ; "/" ; "*" ]
        [ (+) ; (-) ; (/) ; ( * ) ] ;

    dictionary#add_word "print-int" (fun runtime -> runtime#pop_int |> print_int) ;

    dictionary#add_word "dup" (fun runtime -> runtime#dup);

    dictionary#add_word "swap" (fun runtime -> runtime#swap);

    dictionary#add_word "drop" (fun runtime -> runtime#drop);

    dictionary#add_word "call" (fun runtime -> runtime#pop_quote |> runtime#queue_actions) ;

    dictionary#add_word "compose" (fun runtime ->
        let p = runtime#pop_quote in
        let q = runtime#pop_quote in
        `JoyfulQuotation (q @ p) |> runtime#push
    ) ;

    dictionary#add_word "quote" (fun runtime ->
        `JoyfulQuotation [ Push runtime#pop ] |> runtime#push
    ) ;

    dictionary#add_word "stash" (fun runtime -> runtime#stash) ;

    dictionary#add_word "unstash" (fun runtime -> runtime#unstash) ;

    dictionary#add_word "def" (fun runtime ->
        let symbol = runtime#pop_symbol in
        let actions = runtime#pop_quote in
        dictionary#add_word symbol (fun runtime ->
            actions |> runtime#queue_actions
        )
    ) ;

    dictionary#add_word "undef" (fun runtime ->
        let symbol = runtime#pop_symbol in
        dictionary#remove_word symbol) ;

    dictionary#add_word "dip" (fun runtime ->
        [ Call "unstash" ] |> runtime#queue_actions ;
        runtime#pop_quote |> runtime#queue_actions ;
        [ Call "stash" ] |> runtime#queue_actions
    ) ;

    dictionary#add_word "sqrti" (fun runtime ->
        `JoyfulInt (runtime#pop_int |> Float.of_int |> sqrt |> Int.of_float) |> runtime#push) ;

    dictionary#add_word "nl" (fun _ -> print_newline ())


let eval ?(runtime = new runtime_class) quote =
    runtime#queue_quotation quote ;
    let rec helper () =
        let action = runtime#poll in
        match action with
            | None ->
                ()

            | Some (Push value) ->
                runtime#push value ;
                helper ()

            | Some (Call name) ->
                dictionary#call name runtime;
                helper () in
    helper ()

let program = "
    [ quote [ call ] compose swap quote compose call ] `dip def
    [ dup [ dip ] dip call ] `bi@ def
    [ dup * ] `sq def
    [ [ sq ] bi@ + sqrti ] `mag def
    3 4 mag print-int nl

"

let spy_eval program =
    let r = new runtime_class in
    program |> tokenize |> parse |> eval ~runtime:r ;
    ( r#data_stack |> Stack.to_seq |> List.of_seq
    , r#retain_stack |> Stack.to_seq |> List.of_seq
    , r#call_stack |> Stack.to_seq |> List.of_seq
    )