open! Import
include (val Ohlog.sublogs logger "Ren")


(* Buffer *)

type buffer = {
  bo : int;
} [@@unboxed]

let set_buffer_data ?(usage = Gl.stream_draw) {bo} data =
  let size = Array1.size_in_bytes data in
  Gl.bind_buffer Gl.array_buffer bo;
  Gl.buffer_data Gl.array_buffer size (Some data) usage

let make_buffer ?data () =
  let bo = Gl.gen_buffers 1 Gl._i32; Int32.to_int Gl._i32.{0} in
  let buf = {bo} in
  Option.iter (set_buffer_data buf ~usage:Gl.static_draw) data;
  buf

let buffer_size_in_bytes { bo } =
  Gl.get_named_buffer_parameteriv bo Gl.buffer_size Gl._i32;
  Int32.to_int Gl._i32.{0}


(* Vertex buffer attributes *)

type attr = {
  aidx : int;
  atype : [`float];
  asize : int;
}

let attr_type at = match at.atype with
  | `float -> Gl.float

let attr_size_in_bytes at = match at.atype with
  | `float -> at.asize * 4

type vertex_buffer = {
  vbuf : buffer;
  vats : attr list;
  vdiv : int;
}

let make_static_vertex_buffer data vats =
  let vbuf = make_buffer () ~data in
  { vbuf; vats; vdiv = 0 }

(* let make_instance_vertex_buffer vats = *)
(*   let vbuf = make_buffer () in *)
(*   {vbuf; vats; vdiv = 1} *)

let bind_vertex_buffer { vbuf; vdiv; vats } =
  Gl.bind_buffer Gl.array_buffer vbuf.bo;
  let stride =
    List.fold_left (fun n at -> n + attr_size_in_bytes at) 0 vats
  in
  let enable ofs at =
    Gl.vertex_attrib_pointer
      at.aidx at.asize (attr_type at)
      false stride (`Offset ofs);
    Gl.vertex_attrib_divisor at.aidx vdiv;
    Gl.enable_vertex_attrib_array at.aidx;
    ofs + attr_size_in_bytes at
  in
  ignore (List.fold_left enable 0 vats : int)


(* Shaders *)

type shader = {
  spo : int;
} [@@unboxed]

type 'a uniform = U : int -> 'a uniform [@@unboxed]

let compile_shader ~vert ~frag =
  let spo = Gl.create_program () in
  let compile type_ name src =
    let so = Gl.create_shader type_ in
    Gl.attach_shader spo so;
    (* "If a shader object to be deleted is attached to a program object, it will be
     * flagged for deletion, but it will not be deleted until it is no longer attached
     * to any program object" *)
    Gl.delete_shader so;
    Gl.shader_source so src;
    Gl.compile_shader so;
    match Gl.get_shader_compile_status so with
    | Ok () -> ()
    | Error info_log ->
      Gl.delete_program spo;
      Printf.ksprintf failwith "Error compiling %s shader: %s" name info_log
  in
  let link () =
    Gl.link_program spo;
    match Gl.get_program_link_status spo with
    | Ok () -> { spo }
    | Error info_log ->
      Gl.delete_program spo;
      Printf.ksprintf failwith "Error linking shader: %s" info_log
  in
  compile Gl.vertex_shader "vertex" vert;
  compile Gl.fragment_shader "fragment" frag;
  link ()

let load_shader ~name =
  let shd =
    compile_shader
      ~vert:(Printf.ksprintf Asset.load_file "shaders/%s.vert" name)
      ~frag:(Printf.ksprintf Asset.load_file "shaders/%s.frag" name)
  in
  debug (fun m -> m "loaded shader %S" name); shd

let use {spo} =
  Gl.use_program spo

let attr {spo} aname atype asize =
  match Gl.get_attrib_location spo aname with
  | -1 -> Printf.ksprintf failwith "No such attribute %S" aname
  | aidx -> { aidx; atype; asize }

let uniform {spo} name =
  match Gl.get_uniform_location spo name with
  | -1 -> Printf.ksprintf failwith "No such uniform: %S" name
  | loc -> U loc

type 'a set_fn = 'a uniform -> 'a -> unit

let set_int : int set_fn = fun (U l) x -> Gl.uniform1i l x
let set_ivec2 : (int*int) set_fn = fun (U l) (x, y) -> Gl.uniform2i l x y
let set_color : color set_fn = fun (U l) c -> Gl.uniform4f l c.r c.g c.b c.a
let set_aabb : aabb set_fn = fun (U l) b -> Gl.uniform4f l b.x0 b.y0 b.x1 b.y1

let set_mat2a : mat2a set_fn =
  let mat3 = Array1.create Float32 C_layout 9 in
  mat3.{6} <- 0.0;
  mat3.{7} <- 0.0;
  mat3.{8} <- 1.0;
  fun (U l) mat ->
    mat3.{0} <- mat.a0;
    mat3.{1} <- mat.a1;
    mat3.{2} <- mat.a2;
    mat3.{3} <- mat.a3;
    mat3.{4} <- mat.a4;
    mat3.{5} <- mat.a5;
    Gl.uniform_matrix3fv l 1 true mat3


(* Geometry *)

type geometry = {
  vao : int;
  draw_mode : Gl.enum;
  indices : [`count of int | `elems of int];
}

let make_geometry
      ~draw_mode
      ~index
      vertex_buffers
  =
  let vao = Gl.gen_vertex_arrays 1 Gl._i32; Int32.to_int Gl._i32.{0} in
  Gl.bind_vertex_array vao;
  let indices = match index with
    | `count n ->
      Gl.bind_buffer Gl.element_array_buffer 0;
      `count n
    | `elems ibuf ->
      Gl.bind_buffer Gl.element_array_buffer ibuf.bo;
      (* divide by 2 since the buffer should contain int16_unsigned *)
      `elems (buffer_size_in_bytes ibuf / 2)
  in
  List.iter bind_vertex_buffer vertex_buffers;
  {
    vao;
    draw_mode;
    indices
  }

let draw_geometry ?(instances = 1) { vao; draw_mode; indices } =
  Gl.bind_vertex_array vao;
  match indices with
  | `count n ->
    Gl.draw_arrays_instanced draw_mode 0 n instances
  | `elems n ->
    let offset = `Offset 0 in
    let type_ = Gl.unsigned_short in
    Gl.draw_elements_instanced draw_mode n type_ offset instances


(* Renderer *)

type t = {
  window : Sdl.window;
  gl_ctx : Sdl.gl_context;

  polygon : shader;
  rect : geometry;
}

let unit_square =
  (* 1--0
   * |  |
   * 3--2
  *)
  Array1.of_array Float32 C_layout @@
  Array.concat [
    [| 1.; 0.; |];
    [| 0.; 0.; |];
    [| 1.; 1.; |];
    [| 0.; 1.; |];
  ]

let unit_square_with_norm =
  (* 5    4
   *  1--0
   *  |  |
   *  3--2
   * 7    6
  *)
  Array1.of_array Float32 C_layout @@
  Array.concat [
    [| 1.; 0.;  0.;  0. |];
    [| 0.; 0.;  0.;  0. |];
    [| 1.; 1.;  0.;  0. |];
    [| 0.; 1.;  0.;  0. |];
    [| 1.; 0.;  1.; -1. |];
    [| 0.; 0.; -1.; -1. |];
    [| 1.; 1.;  1.;  1. |];
    [| 0.; 1.; -1.;  1. |];
  ]

let make ~(window : Sdl.window) : t =
  debug (fun m -> m "initializing");
  let gl_ctx = Sdl.gl_create_context_exn window in
  Sdl.gl_make_current_exn window gl_ctx;
  Sdl.gl_set_swap_interval_exn 1;

  Gl.enable Gl.blend;
  Gl.blend_func Gl.one Gl.one_minus_src_alpha;
  Gl.check_error "setup";

  let polygon =
    load_shader
      ~name:"polygon"
  in

  let rect =
    make_geometry [
      make_static_vertex_buffer unit_square_with_norm [
        attr polygon "Vert" `float 2;
        attr polygon "Norm" `float 2;
      ]
    ]
      ~draw_mode:Gl.triangle_strip
      ~index:(`count 4)
  in

  {
    window;
    gl_ctx;
    polygon;
    rect;
  }

let destroy t =
  Sdl.gl_delete_context t.gl_ctx

let pre_draw t =
  let viewport = Sdl.get_window_size t.window in
  begin
    Sdl.gl_make_current_exn t.window t.gl_ctx;
    set_ivec2 (uniform t.polygon "Viewport") viewport;
  end

let post_draw t =
  begin
    Sdl.gl_swap_window t.window;
  end

let clear _t (bg : color) =
  begin
    Gl.clear_color bg.r bg.g bg.b bg.a;
    Gl.clear Gl.color_buffer_bit;
  end

let draw_rect t ~(tf : mat2a) ~(rect : aabb) ~(fill : color) =
  let sh = t.polygon in
  begin
    (* TODO: cache/store uniform locations in some way *)
    use sh;
    set_mat2a (uniform sh "Transform") tf;
    set_aabb  (uniform sh "BoundingBox") rect;
    set_int   (uniform sh "Border") 0;
    set_color (uniform sh "Fill") fill;
    draw_geometry t.rect;
  end