open Adam
open S2

include (val Ohlog.logs "N2")

module Transform_graph = struct
  include (val Ohlog.sublogs logger "TG")

  type t = {
    mutable buffer : node array;
    mutable size : int;
  }

  and node =
    | Null
    | Node of {
        idx : int;
        parent_tf : mat2a option;
        model_tf : mat2a;
        world_tf : mat2a;
      }

  let make () = {
    buffer = Array.make 64 Null;
    size = 0;
  }

  let push t n =
    let buf = t.buffer in
    let cap = Array.length buf in
    if t.size >= cap then
      t.buffer <-
        Array.init (cap * 2)
          (fun i -> if i < cap then buf.(i) else Null);
    t.buffer.(t.size) <- n;
    t.size <- t.size + 1

  let add t parent =
    let model_tf = mat2a () in
    let parent_tf, world_tf =
      match parent with
      | Null -> None, model_tf
      | Node p -> Some p.world_tf, mat2a ()
    in
    let node =
      Node {
        idx = t.size;
        parent_tf;
        model_tf;
        world_tf;
      }
    in
    push t node;
    node

  let update t =
    let buf = t.buffer in
    for i = 0 to t.size - 1 do
      match buf.(i) with
      | Node { model_tf; world_tf; parent_tf = Some parent_tf; _ } ->
        Mat2A.multiply world_tf parent_tf model_tf
      | _ ->
        ()
    done

  let world = function
    | Null -> invalid_arg "null"
    | Node n -> n.world_tf

  let model = function
    | Null -> invalid_arg "null"
    | Node n -> n.model_tf
end

module Sprite_graph = struct
  include (val Ohlog.sublogs logger "SG")

  type t = {
    sprite_maps : (string, Sprite_map.t) Hashtbl.t;
    fonts : (string, Font.t) Hashtbl.t;
    mutable list_rev : node list;
    mutable list : node list;
  }

  and node =
    | Sprites of {
        tf : mat2a;
        map : Sprite_map.t;
        frames : Floatbuffer.t;
        (* tint : color *)
      }
    | Label of {
        tf : mat2a;
        font : Font.t;
        glyphs : Floatbuffer.t;
        fg : color;
      }

  let make () = {
    sprite_maps = Hashtbl.create 32;
    fonts = Hashtbl.create 32;
    list_rev = [];
    list = [];
  }

  let register_sprite_map t name map =
    Hashtbl.replace t.sprite_maps name map

  let register_font t name font =
    Hashtbl.replace t.fonts name font

  let get_sprite_map t name =
    try Hashtbl.find t.sprite_maps name with
      Not_found -> Format.kasprintf failwith "no sprite map %S" name

  let get_font t name =
    try Hashtbl.find t.fonts name with
      Not_found -> Format.kasprintf failwith "no font %S" name

  let push t node =
    begin
      t.list_rev <- node :: t.list_rev;
      t.list <- [];
    end

  let add_sprites t ~tf ~map ~sprites =
    let frames = Floatbuffer.make (List.length sprites * 8) in
    List.iter (fun (frame, pos) -> Sprite_map.emit_sprite map frames ~frame ~pos) sprites;
    let node = Sprites { tf; map; frames } in
    push t node;
    node

  let add_label t ~tf ~font ~text ~fg =
    let glyphs = Floatbuffer.make (String.length text * 10) in
    Font.emit_glyphs font glyphs ~text;
    let node = Label { tf; font; glyphs; fg } in
    push t node;
    node

  let _white = Color.white ()

  let rec render_rec ren = function
    | [] -> ()
    | Sprites { tf; map; frames } :: nodes ->
      Renderer.draw_sprites ren map frames ~tf ~tint:_white;
      render_rec ren nodes
    | Label { tf; font; glyphs; fg } :: nodes ->
      Renderer.draw_text ren font glyphs ~tf ~fg;
      render_rec ren nodes

  let render t ~ren =
    if t.list = [] then
      t.list <- List.rev t.list_rev;
    render_rec ren t.list
end

module Scene = struct
  module Sexp = Sexplib0.Sexp
  open Sexplib0.Sexp_conv

  include (val Ohlog.sublogs logger "Scene")

  type t = {
    objs : (string, obj) Hashtbl.t;
  }

  and obj = {
    name : string option;
    transform : Transform_graph.node;
  }

  let get t name = try Hashtbl.find t.objs name
    with Not_found ->
      Format.ksprintf failwith "no node %S in scene" name

  let transform obj = obj.transform

  let parse_transform_arg ~(tf : mat2a) = function
    | Sexp.List [Atom "tx"; arg] ->
      tf.a2 <- float_of_sexp arg
    | Sexp.List [Atom "ty"; arg] ->
      tf.a5 <- float_of_sexp arg
    | Sexp.List [Atom "sx"; arg] ->
      tf.a0 <- float_of_sexp arg
    | Sexp.List [Atom "sy"; arg] ->
      tf.a4 <- float_of_sexp arg
    | sexp ->
      of_sexp_error "bad argument to transform" sexp

  let parse_sprite_arg ~map = function
    | Sexp.List [Atom frame; x; y] ->
      let pos = vec2 (float_of_sexp x) (float_of_sexp y) in
      frame, pos
    | Atom frame ->
      frame, vec2 0.0 0.0
    | sexp ->
      of_sexp_error "bad sprite argument" sexp

  let parse_sprites ~sg ~tf = function
    | Sexp.List (Atom "sprites" :: Atom map :: args) ->
      let map = Sprite_graph.get_sprite_map sg map in
      let sprites = List.map (parse_sprite_arg ~map) args in
      Sprite_graph.add_sprites sg ~tf ~map ~sprites |> ignore
    | sexp ->
      of_sexp_error "invalid sprites" sexp

  let parse_label_arg ~(fg : color) = function
    | Sexp.List [Atom "fg"; rgb] ->
      Color.set_rgb24 fg (int_of_sexp rgb)
    | sexp ->
      of_sexp_error "bad argument to label" sexp

  let parse_label ~sg ~tf = function
    | Sexp.List (Atom "label" :: Atom font :: Atom text :: args) ->
      let font = Sprite_graph.get_font sg font in
      let fg = rgb24 0xffffff in
      List.iter (parse_label_arg ~fg) args;
      Sprite_graph.add_label sg ~tf ~font ~text ~fg |> ignore
    | sexp ->
      of_sexp_error "invalid label" sexp

  let rec parse_obj t ~tg ~sg ~parent = function
    | Sexp.List (Atom "object" :: args) ->
      trace (fun m -> m "parse_obj");
      let ptgn = match parent with
        | None -> Transform_graph.Null
        | Some parent -> parent.transform
      in
      let transform = Transform_graph.add tg ptgn in
      let obj = { name = None; transform } in
      let obj = parse_obj_args t obj args ~tg ~sg in
      Option.iter
        (fun name -> Hashtbl.replace t.objs name obj)
        obj.name;
      obj
    | sexp ->
      of_sexp_error "invalid object" sexp

  and parse_obj_args t obj ~tg ~sg = function
    | [] -> obj
    | arg :: args ->
      match arg with
      | Atom name ->
        trace (fun m -> m "parse_obj_args: name=%S" name);
        let obj = { obj with name = Some name } in
        parse_obj_args t obj args ~tg ~sg

      | Sexp.List (Atom "transform" :: tf_args) ->
        let tf = Transform_graph.model obj.transform in
        List.iter (parse_transform_arg ~tf) tf_args;
        trace (fun m -> m "parse_obj_args: tf=%a" Mat2A.pp tf);
        parse_obj_args t obj args ~tg ~sg

      | Sexp.List (Atom "sprites" :: _) ->
        let tf = Transform_graph.world obj.transform in
        parse_sprites arg ~sg ~tf;
        parse_obj_args t obj args ~tg ~sg

      | Sexp.List (Atom "label" :: _) ->
        let tf = Transform_graph.world obj.transform in
        parse_label arg ~sg ~tf;
        parse_obj_args t obj args ~tg ~sg

      | Sexp.List (Atom "object" :: _) ->
        trace (fun m -> m "parse_obj_args: begin object");
        let _child = parse_obj t arg ~tg ~sg ~parent:(Some obj) in
        trace (fun m -> m "parse_obj_args: end object");
        parse_obj_args t obj args ~tg ~sg

      | sexp ->
        of_sexp_error "bad argument to object" sexp

  let of_sexp ~tg ~sg = function
    | Sexp.List [Atom "scene"; root] ->
      let t = {
        objs = Hashtbl.create 128;
      } in
      let _root = parse_obj t root ~tg ~sg ~parent:None in
      t

    | sexp ->
      of_sexp_error "invalid scene" sexp

  let load name ~tg ~sg =
    S2.Asset.load_sexp_conv
      (name ^ ".scene")
      (of_sexp ~tg ~sg)
end