(*
    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 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)
            | x :: rest ->
                 match int_of_string_opt x with
                    | Some n ->
                        let quotation = Push (`JoyfulInt n) :: quotation in
                        parse_quotation rest ~quotation
                    | None ->
                        let quotation = Call x :: 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)
            | x :: rest ->
                match int_of_string_opt x with
                    | Some n ->
                        Push (`JoyfulInt n) :: prog |> helper rest
                    | None when String.starts_with ~prefix:"`" x ->
                        Push (`JoyfulSymbol (String.sub x 1 (String.length x - 1))) :: prog |> helper rest
                    | None ->
                        Call x :: 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
    end

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

        method new_primitive name f =
            Hashtbl.add words name f


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


let () =
    let new_math_primitives name op =
        dictionary#new_primitive 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#new_primitive "print-int" (fun runtime -> runtime#pop_int |> print_int) ;

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

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

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

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

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

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

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


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

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

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

let eval quote =
    let runtime = new runtime_class in
    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 mag_program = "
[ dup [ dip ] dip call ] `bi@ def
[ dup * ] `sq def
3 4 [ sq ] bi@ + sqrti print-int nl
"

let run_mag () =
    print_endline mag_program ;
    mag_program |> tokenize |> parse |> eval