open Ast

(* resulting value of eval *)
type value =
  | Int of int
  | Float of float
  | Bool of bool
  | String of string
  | Symbol of string
  (* (name), arg, expression, name *)
  | Function of string option * string * expr * env
  | External of string
  | Nop (* return of system operations (will be deprecated) *)

and expr = Ast.t

(* environment for eval *)
and env = (string * value) list

exception No_operation
exception Too_many_arguments

(* TODO: add proper type system *)
module Type = struct
  type t =
    | Int
    | Float
    | Bool
    | String
    | Symbol
    | Function
    | External

  exception Invalid of t
  exception Expected of t

  let to_string = function
    | Int -> "int"
    | Float -> "float"
    | Bool -> "bool"
    | String -> "string"
    | Symbol -> "symbol"
    | Function -> "fun"
    | External -> "external"

  let supertype = function
    | Int -> Some Float
    | _ -> None
end

module Value = struct
  type t = value

  let to_string = function
    | Int n -> string_of_int n
    | Float n -> string_of_float n
    | Bool b -> string_of_bool b
    | String s -> "\"" ^ s ^ "\""
    | Symbol s -> "symbol " ^ s
    | Function _ -> "<fun>"
    | External f -> "external " ^ f
    | Nop -> "nop"

  let typeof = function
    | Int _ -> Type.Int
    | Float _ -> Type.Float
    | Bool _ -> Type.Bool
    | String _ -> Type.String
    | Symbol _ -> Type.Symbol
    | Function _ -> Type.Function
    | External _ -> Type.External
    | Nop -> failwith "Value.typeof"

  let promote = function
    | Int n -> Float (float n)
    | _ -> failwith "Value.promote"
end

module Env = struct
  type t = env

  let empty = []

  let get_opt e name =
    List.assoc_opt name e

  let bind v e =
    v::e

  let bind_seq seq e =
    List.of_seq seq @ e
end

(* operators *)
module Operator = struct
  type t = Ast.operator

  exception Unavailable of t

  let to_string = Ast.op_to_string

  let negate = function
    | [Int n] -> Int ~-n
    | [Float n] -> Float ~-.n
    | _ -> failwith "Operator.negate"

  let vi f = function
    | [Int a; Int b] -> Int (f a b)
    | _ -> raise @@ Type.Invalid Int

  let vf f = function
    | [Float a; Float b] -> Float (f a b)
    | _ -> raise @@ Type.Expected Float

  let compare = function
    | [Int a; Int b] -> Int.compare a b
    | [Float a; Float b] -> Float.compare a b
    | [Bool a; Bool b] -> Bool.compare a b
    | [String a; String b] -> String.compare a b
    | [Symbol a; Symbol b] -> String.compare a b
    | _ -> invalid_arg "Operator.compare"

  let eq vs = Bool (compare vs = 0)
  let neq vs = Bool (compare vs <> 0)
  let ge vs = Bool (compare vs >= 0)
  let le vs = Bool (compare vs <= 0)
  let gt vs = Bool (compare vs > 0)
  let lt vs = Bool (compare vs < 0)

  (* operator table *)
  let operators =
    let open Type in
    let ip = [Int; Int] and fp = [Float; Float] in
    let any f = [ip, f; fp, f] in
    [
      Add, [ip, vi Int.add; fp, vf Float.add];
      Sub, [ip, vi Int.sub; fp, vf Float.sub];
      Mul, [ip, vi Int.mul; fp, vf Float.mul];
      Div, [ip, vi Int.div; fp, vf Float.div];
      Mod, [ip, vi Int.rem; fp, vf Float.rem];
      Exp, [fp, vf Float.pow];
      Eq,  any eq;
      Neq, any neq;
      GE, any ge;
      LE, any le;
      GT, any gt;
      LT, any lt;
      Negate, [[Int], negate; [Float], negate];
    ]
    |> List.to_seq
    |> Hashtbl.of_seq

  let get op =
    Hashtbl.find operators op
end

module External = struct
  exception Invalid of string

  let rad r =
    r *. 180. /. Float.pi

  let deg d =
    d /. 180. *. Float.pi

  let floatfun f = function
    | [Float n] -> Float (f n)
    | [v] -> raise @@ Type.Invalid (Value.typeof v)
    | _ -> invalid_arg "External.floatfun"

  let symbol_to_op op =
    op
    |> String.to_seqi
    |> Lex.find_token
    |> Option.get
    |> fst
    |> Parser.token_to_op

  let set_op_pre = function
    | [Symbol op; Int l] ->
      let op = symbol_to_op op in
      Hashtbl.replace Parser.precedence op l;
      Nop
    | _ -> failwith "set_op_pre"

  let get_op_pre = function
    | [Symbol op] ->
      let op = symbol_to_op op in
      Int (Hashtbl.find Parser.precedence op)
    | _ -> failwith "get_op_pre"

  let set_op_assoc = function
    | [Symbol op; String a] ->
      let op = symbol_to_op op in
      Hashtbl.replace Parser.oper_assoc op @@ Parser.assoc_of_string a;
      Nop
    | _ -> failwith "set_op_assoc"

  let get_op_assoc = function
    | [Symbol op] ->
      let op = symbol_to_op op in
      Hashtbl.find_opt Parser.oper_assoc op
      |> Option.value ~default: Parser.Left_to_right
      |> (fun a -> String (Parser.assoc_to_string a))
    | _ -> failwith "get_op_assoc"

  let apply f args =
    let f = match f with
      | "sin" -> floatfun Float.sin
      | "cos" -> floatfun Float.cos
      | "tan" -> floatfun Float.tan
      | "deg" -> floatfun deg
      | "rad" -> floatfun rad
      | "set_op_pre" -> set_op_pre
      | "get_op_pre" -> get_op_pre
      | "set_op_assoc" -> set_op_assoc
      | "get_op_assoc" -> get_op_assoc
      | _ -> raise @@ Invalid f
    in
    f args
end

let find_operator op ts =
  let filter t =
    List.filter (fun (ts, _) ->
        match ts with [] -> false | x::_ -> t=x)
  in
  let rec aux ops = function
    | [] -> List.nth_opt ops 0
    | t::ts ->
      (match aux (filter t ops) ts with
      | None -> Option.bind (Type.supertype t) (fun t -> aux ops (t::ts))
      | Some _ as x -> x)
  in
  aux (Operator.get op) ts

let promote_values =
  let rec promote_until t v =
    if Value.typeof v = t
    then v
    else promote_until t @@ Value.promote v
  in
  List.map2 promote_until

let unary op v =
  match find_operator op [Value.typeof v] with
  | None -> raise No_operation
  | Some (ts, f) ->
    let vs = promote_values ts [v] in
    f vs

let binop op l r =
  let open Value in
  match find_operator op [typeof l; typeof r] with
  | None -> raise No_operation
  | Some (ts, f) ->
    let vs = promote_values ts [l; r] in
    f vs

exception Unbound of string

let rec eval global env ast =
  let aux = eval global env in (* eval with current env *)
  match ast with
    | Nothing -> Nop
    | Nint n -> Int n
    | Nfloat n -> Float n
    | Nbool b -> Bool b
    | Nstring s -> String s
    | Nsymbol s -> Symbol s
    | Nfunction (arg, e) -> Function (None, arg, e, env)
    | Nexternal f -> External f

    | Var v -> begin match Env.get_opt env v with
        | None -> (try Hashtbl.find global v
                   with Not_found -> raise @@ Unbound v)
        | Some v -> v
      end
    | Letin (v, e, f) ->
      let env = Env.bind (v, aux e) env in
      eval global env f

    | Unary (op, v) -> unary op (aux v)
    | Binop (l, op, r) -> binop op (aux l) (aux r)

    | If (co, th, el) ->
      begin match aux co with
        | Bool true -> aux th
        | Bool false -> aux el
        | v -> raise @@ Type.Invalid (Value.typeof v)
      end
    | Apply (v, args) -> apply global env v args

    | _ -> failwith "Eval.eval"

(* apply args to result of expr *)
and apply global env expr args =
  match eval global env expr with
  | Function (itself, var, body, local_env) as f ->
    begin match args with
      | [] -> f
      | a::args ->
        let value = eval global env a in
        let env = (* binding itself into env for recursion *)
          itself
          |> Option.fold
            ~none: local_env
            ~some: (fun n -> Env.bind (n, f) local_env)
          |> Env.bind (var, value)
        in
        apply global env body args
    end
  | External f ->
    let args = List.map (eval global env) args in
    External.apply f args
  | v ->
    if args = [] then v
    else raise @@ Type.Invalid (Value.typeof v)

(* toplevel for global let *)
let eval_top global ast =
  let var, v = match ast with
    | Let (var, Nfunction (arg, e)) -> (* named function *)
      var, Function (Some var, arg, e, Env.empty)
    | Let (var, e) -> var, eval global Env.empty e
    | ast -> "-", eval global Env.empty ast
  in
  if var <> "-" then Hashtbl.replace global var v;
  var, v