open Ast

(* resulting value of eval *)
type value =
  | Unit
  | Int of int
  | Float of float
  | Bool of bool
  | String of string
  | Symbol of string
  (* (name), bound variables, expression, environment *)
  | Function of string option * string list * expr * env
  | External of string

and expr = Ast.t

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

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

  exception Invalid of t
  exception Expected of t

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

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

  let matches : t -> t -> bool = function
    | Any -> Fun.const true
    | t -> fun o -> o = t || o = Any
end

module Value = struct
  type t = value

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

  let typeof = function
    | Unit -> Type.Unit
    | Int _ -> Int
    | Float _ -> Float
    | Bool _ -> Bool
    | String _ -> String
    | Symbol _ -> Symbol
    | Function _ -> Function
    | External _ -> External

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

module Env = struct
  type t = env

  let empty = Env []

  let get_opt (Env e) name =
    List.assoc_opt name e

  let bind v (Env e) =
    Env (v::e)

  let bind_seq seq (Env e) =
    Env (List.of_seq seq @ e)
end

(* primitive methods *)
module Primitive = struct
  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)

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

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

  let floatfun f =
    [Type.Float],
    function
    | [Float n] -> Float (f n)
    | _ -> 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;
      Unit
    | _ -> 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;
      Unit
    | _ -> 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 print args =
    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 -> s
      | _ -> failwith "print"
    in
    List.map to_string args
    |> List.iter (Printf.printf "%s");
    Unit

  let println args =
    ignore @@ print args;
    Printf.printf "\n";
    Unit


  let methods =
    let open Type in
    let ip = [Int; Int] and fp = [Float; Float] in
    let number f = [ip, f; fp, f] in
    [
      "+", [ip, vi Int.add; fp, vf Float.add];
      "-", [ip, vi Int.sub; fp, vf Float.sub;
            [Int], negate; [Float], negate];
      "*", [ip, vi Int.mul; fp, vf Float.mul];
      "/", [ip, vi Int.div; fp, vf Float.div];
      "%", [ip, vi Int.rem; fp, vf Float.rem];
      "^", [fp, vf Float.pow];
      "=", number eq;
      "<>", number neq;
      ">=", number ge;
      "<=", number le;
      ">", number gt;
      "<", number lt;

      "sin", [floatfun Float.sin];
      "cos", [floatfun Float.cos];
      "tan", [floatfun Float.tan];
      "deg", [floatfun deg];
      "rad", [floatfun rad];
      "set_op_pre", [[Symbol; Int], set_op_pre];
      "get_op_pre", [[Symbol], get_op_pre];
      "set_op_assoc", [[Symbol; String], set_op_assoc];
      "get_op_assoc", [[Symbol], get_op_assoc];
      "print", [[Any], print];
      "println", [[Any], println];
    ]
    |> List.to_seq
    |> Hashtbl.of_seq

  let get op =
    Hashtbl.find methods op
end

(* find_method returns a method and (corresponding) type list that
 * satisfies ts. *)
let find_method m ts =
  let open List in
  let filter_type t i =
    filter (fun (ts, _) ->
        nth_opt ts i
        |> Option.map (Type.matches t)
        |> Option.value ~default: false)
  in
  let rec aux ms i = function
    | [] -> nth_opt ms 0
    | t::ts ->
      (match aux (filter_type t i ms) (i+1) ts with
       | None -> Option.bind (Type.supertype t)
                   (fun t -> aux ms i (t::ts))
       | Some _ as x -> x)
  in
  let ms =
    let len = length ts in
    Primitive.get m
    |> filter (fun (ts, _) -> length ts = len)
  in
  aux ms 0 ts

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

exception No_such_method of string * Type.t
let no_such m v = raise @@ No_such_method (m, Value.typeof v)

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

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

let extern f args =
  match find_method f (List.map Value.typeof args) with
  | None -> no_such f (List.hd args)
  | Some (ts, f) ->
    let vs = promote_values ts args in
    f vs

exception Unbound of string
exception Noop

let rec eval global env ast =
  let aux = eval global env in (* eval with current env *)
  match ast with
    | Nothing -> raise Noop
    | Nunit -> Unit
    | 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) ->
      let args = List.map (eval global env) args in
      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, vars, body, local_env) as f ->
    begin match args with
      | [] -> f
      | args ->
        (* bind arguments to variables *)
        let rec aux e = function
          | [], [] -> [], [], e
          | vars, [] -> vars, [], e
          | [], args -> [], args, e
          | v::vars, a::args ->
            let e = Env.bind (v, a) e in
            aux e (vars, args)
        in
        let vars, args, env = aux local_env (vars, args) in
        let env = (* binding itself into env for recursion *)
          itself |> Option.fold
            ~none: env
            ~some: (fun n -> Env.bind (n, f) env)
        in
        if vars <> [] then (* partially-applied function *)
          Function (None, vars, body, env)
        else if args <> [] then (* reapply *)
          apply global env body args
        else (* eval (vars = [], args = []) *)
          eval global env body
    end
  | External f -> extern 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