open Ast

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

and expr = Ast.t

(* environment for eval *)
and env = {
  vars : (string, value) Hashtbl.t;
  parent : env option;
}

exception No_operation
exception Too_many_arguments

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 (vars, _) ->
      Printf.sprintf "function with %d arguments" @@ List.length vars
    | 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 init_global () = {
    vars = Hashtbl.create 100;
    parent = None;
  }

  let make parent = {
    vars = Hashtbl.create 100;
    parent = Some parent;
  }

  let rec get_opt e name =
    match Hashtbl.find_opt e.vars name with
    | None -> Option.bind e.parent (fun p -> get_opt p name)
    | Some _ as v -> v

  let set e name value =
    Hashtbl.replace e.vars name value

  let add_seq e seq =
    Hashtbl.add_seq e.vars seq
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 a b =
    match a, b with
    | Int a, Int b -> Int (f a b)
    | _ -> raise @@ Type.Invalid Int

  let vf f a b =
    match a, b with
    | Float a, Float b -> Float (f a b)
    | _ -> raise @@ Type.Expected Float

  let vb intf floatf a b =
    match a, b with
    | Int a, Int b -> Bool (intf a b)
    | Float a, Float b -> Bool (floatf a b)
    | _ -> raise @@ Type.Expected Bool

  let vnot = function
    | Bool b -> Bool (not b)
    | _ -> raise @@ Type.Expected Bool

  let map ?intf ?floatf ?boolf v =
    let app x f = f x in
    match v with
    | Int i -> Option.map (app i) intf
    | Float f -> Option.map (app f) floatf
    | Bool b -> Option.map (app b) boolf
    | _ -> invalid_arg "Operator.map"

  let eq = vb Int.equal Float.equal
  let neq a b = vnot @@ eq a b

  let compare a b =
    match a, b with
    | Int a, Int b -> Int.compare a b
    | Float a, Float b -> Float.compare a b
    | _ -> invalid_arg "Operator.compare"

  let ge a b = Bool (compare a b >= 0)
  let le a b = Bool (compare a b <= 0)
  let gt a b = Bool (compare a b > 0)
  let lt a b = Bool (compare a b < 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;
    ]
    |> List.to_seq
    |> Hashtbl.of_seq

  let get_types op =
    match Hashtbl.find_opt operators op with
    | None -> raise @@ Unavailable op
    | Some p -> List.map fst p

  let get_unary = function
    | Negate -> negate
    | op -> raise @@ Unavailable op

  let get_binary op typ =
    Hashtbl.find operators op
    |> List.assoc_opt typ
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)

  let apply f args =
    match f, args with
    | "sin", [n] -> floatfun Float.sin n
    | "cos", [n] -> floatfun Float.cos n
    | "tan", [n] -> floatfun Float.tan n
    | "deg", [n] -> floatfun deg n
    | "rad", [n] -> floatfun rad n
    | _ -> raise @@ Invalid f
end

let assert_same_length vars args =
  let vl = List.length vars
  and al = List.length args in
  if vl > al then
    failwith "assert_same_length"
  else if vl < al then
    raise Too_many_arguments

let resolve_type op tp =
  let optypes = Operator.get_types op in
  let q = Queue.create () in
  let rec aux (t1, t2) =
    if List.mem (t1, t2) optypes then
      t1, t2
    else begin
      [ Type.supertype t1 |> Option.map (fun t1 -> t1, t2);
        Type.supertype t2 |> Option.map (fun t2 -> t1, t2); ]
      |> List.filter_map Fun.id
      |> List.iter (Fun.flip Queue.push q);
      aux @@ Queue.pop q
    end
  in
  aux tp

let rec binop op l r =
  let open Value in
  let t1 = typeof l and t2 = typeof r in
  let t1, t2 = resolve_type op (t1, t2) in
  let rec promote_until t x =
    if typeof x = t
    then x
    else promote_until t (promote x)
  in
  let l = promote_until t1 l
  and r = promote_until t2 r in
  match Operator.get_binary op (t1, t2) with
  | None -> begin
      try binop op (promote l) (promote r)
      with _ -> raise No_operation
    end
  | Some f -> f l r

exception Unbound of string

let rec eval env ast : string * value =
  let rec aux = function
    | Nint n -> Int n
    | Nfloat n -> Float n
    | Nbool b -> Bool b
    | Nstring s -> String s
    | Nsymbol s -> Symbol s
    | Nfunction (args, e) -> Function (args, e)
    | Nexternal f -> External f
    | Var v -> begin match Env.get_opt env v with
        | None -> raise @@ Unbound v
        | Some v -> v
      end
    | Unary (op, t) ->
      let t = aux t in
      let op = Operator.get_unary op in
      op t
    | Binop (l, op, r) ->
      let l = aux l and r = aux r in
      binop op l 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) ->
      begin match aux v with
        | Function (vars, e) ->
          assert_same_length vars args;
          let args = List.map aux args in
          let nenv = Env.make env in
          List.combine vars args
          |> List.iter (fun (v, a) -> Env.set nenv v a);
          snd @@ eval nenv e
        | External f ->
          let args = List.map aux args in
          External.apply f args
        | v -> raise @@ Type.Invalid (Value.typeof v)
      end
    | Set_binop_pre (op, l) ->
      let l =
        match aux l with
        | Int n -> n
        | v -> raise @@ Type.Invalid (Value.typeof v)
      in
      Hashtbl.replace Parser.precedence op l;
      Nop
    | Get_binop_pre op ->
      Int (Hashtbl.find Parser.precedence op)
    | Set_binop_aso (op, a) ->
      Hashtbl.replace Parser.oper_assoc op @@ Parser.assoc_of_string a;
      Nop
    | Get_binop_aso op ->
      (match Hashtbl.find_opt Parser.oper_assoc op with
       | None -> String "left"
       | Some a -> String (Parser.assoc_to_string a))
    | _ -> failwith "Eval.eval"
  in

  match ast with
  | Let (var, e) ->
    let v = aux e in
    Env.set env var v;
    var, v
  | ast -> "-", aux ast