(* Types *)

type tile = int * int * int

type terrain = {
  grid1: tile array array;
  grid2: int array array;
  dims: int * int;
}

type resource = {
  r_pos: int * int;
  r_kind: string;
}

type building = {
  b_pos: int * int;
  b_kind: string;
  b_owner: int;
}

type unite = {
  u_pos: int * int;
  u_kind: string;
  u_owner: int;
  u_tasks: (string * int * int * string) list;
  u_carry: string list;
}

let _tasks = ref []

(* Init-functions *)

let new_tile x y d =
  (x, y, d)

let new_terrain width height =
  let grid1 = Array.init height (fun y -> Array.init width (fun x -> new_tile x y 0)) in
  let grid2 = Array.init height (fun y -> Array.init width (fun x -> 0)) in
  { grid1; grid2; dims = (width, height); }

let new_resource pos kind = {
  r_pos = pos;
  r_kind = kind;
}

let new_building pos kind owner = {
  b_pos = pos;
  b_kind = kind;
  b_owner = owner;
}

let new_unite pos kind owner = {
  u_pos = pos;
  u_kind = kind;
  u_owner = owner;
  u_tasks = [];
  u_carry = [];
}


(* Game-state *)

type game_state = {
  terrain: terrain;
  resources: resource list;
  buildings: building list;
  unites: unite list;
}

(* Utils *)

let move_towards target (x, y) =
  let (tx, ty) = target in
  let dx = if x < tx then x + 1 else if x > tx then x - 1 else x in
  let dy = if y < ty then y + 1 else if y > ty then y - 1 else y in
  (dx, dy)

let sq_dist (x1, y1) (x2, y2) =
  let x = x2 - x1 in
  let y = y2 - y1 in
  ((x * x) + (y * y))

let rand_take lst =
  let n = List.length lst in
  let i = Random.int n in
  (List.nth lst i)

let ( >. ) v f = f v ;;


(* Build *)

let building_kind_2 = ["Mine"; "Refinery"; "Factory"]
let building_kind = ["MilitaryBase"; "Habitat"]

let add_building state b =
  { state with buildings = b :: state.buildings }


(* Tasks *)

let tasks = ["GoTo"; "Attack"; "Build"; "Harvest"]

let add_task task u =
  { u with u_tasks = task :: u.u_tasks }

let sort_resources pos resources =
  List.sort (fun r1 r2 ->
    let d1 = sq_dist r1.r_pos pos in
    let d2 = sq_dist r2.r_pos pos in
    compare d1 d2
  ) resources


let new_task state unite =
  if Random.int 100 > 8 then unite else
  match Random.int 3 with
  | 0 -> add_task ("GoTo", Random.int 68, Random.int 44, "") unite
  | 1 -> add_task ("Build", Random.int 68, Random.int 44, (rand_take building_kind)) unite
  | _ ->
      let resources = sort_resources unite.u_pos state.resources in
      match resources with
      | [] -> unite
      | r :: _ ->
          let x, y = r.r_pos in
          add_task ("Harvest", x, y, "") unite


let new_task state unite =
  match !_tasks with
  | [] -> new_task state unite
  | ("go-to", (x, y)) :: tl -> _tasks := tl; add_task ("GoTo", x, y, "") unite
  | ("build", (x, y)) :: tl -> _tasks := tl; add_task ("Build", x, y, (rand_take building_kind)) unite
  | ("harvest", (x, y)) :: tl ->
    begin
      _tasks := tl;
      let resources = sort_resources (x, y) state.resources in
      match resources with
      | [] -> unite
      | r :: _ ->
          let x, y = r.r_pos in
          add_task ("Harvest", x, y, "") unite
    end
  | _ -> assert false


(* Interactions *)

let resource_kind =
  [ "elementium"; "titanium"; "aetherium"; "nebulite"; "lunarite";
    "celestium"; "astralite"; "drakonite"; "neutronium"; "astronium"; ]



(* Proceeds *)

let proceed_goto state pos unite tl =
  if pos = unite.u_pos
  then state, { unite with u_tasks = tl }
  else begin
    let new_pos = move_towards pos unite.u_pos in
    state, { unite with u_pos = new_pos }
  end


let proceed_harvest state pos unite tl =
  if pos = unite.u_pos
  then
    let resources = List.filter (fun r -> r.r_pos <> pos) state.resources in
    let state = { state with resources } in
    state, { unite with u_tasks = tl }
  else begin
    let new_pos = move_towards pos unite.u_pos in
    state, { unite with u_pos = new_pos }
  end


let proceed_build state pos k unite tl =
  if pos = unite.u_pos
  then
    let b = new_building pos k unite.u_owner in
    let state = add_building state b in
    state, { unite with u_tasks = tl }
  else begin
    let new_pos = move_towards pos unite.u_pos in
    state, { unite with u_pos = new_pos }
  end


(* Updates *)

let update_unites state =
  let rec aux acc state unites =
    match unites with
    | [] -> { state with unites = acc }
    | unite :: unites ->
        match unite.u_tasks with
        | ("GoTo", x, y, _) :: tl ->
            let state, u = proceed_goto state (x, y) unite tl in
            aux (u::acc) state unites
        | ("Harvest", x, y, _) :: tl ->
            let state, u = proceed_harvest state (x, y) unite tl in
            aux (u::acc) state unites
        | ("Build", x, y, k) :: tl ->
            let state, u = proceed_build state (x, y) k unite tl in
            aux (u::acc) state unites
        | [] ->
            let u = new_task state unite in
            aux (u::acc) state unites
        | _ ->
            aux (unite::acc) state unites
  in
  aux [] state state.unites


let update_building state building =
  let pos = building.b_pos in
  let own = building.b_owner in
  if Random.int 280 = 0
  then begin
    match building.b_kind with
    | "Habitat" ->
        let u = new_unite pos "worker" own in
        { state with unites = u :: state.unites }
    | "MilitaryBase" ->
        let u = new_unite pos "soldier" own in
        { state with unites = u :: state.unites }
    | _ -> (state)
  end
  else (state)


let update_resources resources =
  let resources = List.filter (fun _ -> Random.int 160 <> 0) resources in
  let k = rand_take resource_kind in
  let r = new_resource (Random.int 68, Random.int 44) k in
  if Random.int 8 = 0
  then r :: resources
  else resources


let acid_rain state =
  if List.length state.buildings > 20
  && List.length state.unites > 20 then
    let buildings = List.filter (fun _ -> Random.int 120 <> 0) state.buildings in
    let unites = List.filter (fun _ -> Random.int 120 <> 0) state.unites in
    { state with buildings; unites; }
  else state


let update_state state =
  let resources = update_resources state.resources in
  let state = { state with resources } in
  let state = List.fold_left update_building state state.buildings in
  let state = update_unites state in
  let state = acid_rain state in
  (state)


(* Init *)

let init_state () =
  let pos1 = (Random.int 68, Random.int 44) in
  let pos2 = (Random.int 68, Random.int 44) in
  let unite = new_unite pos1 "worker" 0 in
  let resource1 = new_resource pos2 "elementium" in
  let unite = add_task ("GoTo", 18, 11, "") unite in
  let terrain = new_terrain 68 44 in
  for _ = 0 to pred 16 do
    let x, y = (Random.int 68, Random.int 44) in
    terrain.grid2.(y).(x) <- 2;
  done;
  {
    resources = [resource1];
    unites = unite :: [];
    buildings = [];
    terrain;
  }


(* Canvas *)

let width, height = (680, 440)

let canvas = Canvas.getElementById Canvas.document "my_canvas"
let ctx = Canvas.getContext canvas "2d"

let bg_color = "#111"

let fill_rect color (x, y) =
  Canvas.fillStyle ctx color;
  Canvas.fillRect ctx (x * 10) (y * 10) 10 10;
  ()
;;

let fill_rect2 color (x, y, w, h) =
  Canvas.fillStyle ctx color;
  Canvas.fillRect ctx (x * 10) (y * 10) (w * 10) (h * 10);
  ()
;;

let stroke_rect color (x, y) =
  Canvas.lineWidth ctx 1;
  Canvas.strokeStyle ctx color;
  Canvas.strokeRect ctx (x * 10) (y * 10) 10 10;
  ()
;;

let fill_color color =
  let r, g, b = color in
  let color = Printf.sprintf "rgba(%d,%d,%d,%d)" r g b 1 in
  Canvas.fillStyle ctx color;
  ()
;;

let draw_text color (x, y) m =
  fill_color color;
  Canvas.font ctx "bold 9px Arial";
  Canvas.fillText ctx m (2 + x * 10) (8 + y * 10);
  ()
;;

(* Draw *)

let draw_resource r =
  fill_rect "#ec4" r.r_pos;
;;

let draw_habitat b =
  fill_rect "#59e" b.b_pos;
;;

let draw_militarybase b =
  fill_rect "#a2d" b.b_pos;
;;

let draw_worker u =
  fill_rect "#4d8" u.u_pos;
;;

let draw_soldier u =
  fill_rect "#f62" u.u_pos;
;;

let draw_building b =
  match b.b_kind with
  | "Habitat" -> draw_habitat b
  | "MilitaryBase" -> draw_militarybase b
  | _ -> fill_rect "#888" b.b_pos;
;;

let draw_unite u =
  match u.u_kind with
  | "worker" -> draw_worker u
  | "soldier" -> draw_soldier u
  | _ -> fill_rect "#eee" u.u_pos;
;;

let menu = ref None

(* Display *)

let display_game state =

  (* Background *)
  Canvas.fillStyle ctx bg_color;
  Canvas.fillRect ctx 0 0 width height;

  let w, h = state.terrain.dims in
  for x = 0 to pred w do
    for y = 0 to pred h do
      match state.terrain.grid2.(y).(x) with
      | 2 -> fill_rect "#151515" (x, y)
      | _ -> ()
    done;
  done;

  (* Draw Elements *)
  List.iter draw_resource state.resources;
  List.iter draw_unite state.unites;
  List.iter draw_building state.buildings;

  (* Draw Tasks *)
  begin
    List.iter (fun (m, p) ->
      let c = String.make 1 m.[0] in
      stroke_rect "#444" p;
      draw_text (60, 60, 60) p c;
    ) !_tasks
  end;
  begin
    List.iter (fun u ->
      match u.u_tasks with [] -> ()
      | (m, x, y, _) :: _ ->
          let m = String.lowercase_ascii m in
          let c = String.make 1 m.[0] in
          stroke_rect "#666" (x, y);
          draw_text (90, 90, 90) (x, y) c;
    ) state.unites
  end;

  (* Draw Menu *)
  begin
    match !menu with
    | None -> ()
    | Some (p, ms) ->
        stroke_rect "#444" p;
        List.iter (fun (r, m) ->
          let x, y, _, _ = r in
          fill_rect2 "#338" r;
          draw_text (140, 190, 255) (x, y) m;
        ) ms
  end;

  ()
;;

(* Contextual-Menu *)

let in_rect (x1, y1) (x, y, w, h) =
  if x1 < x then false else
  if y1 < y then false else
  if x1 > x + w then false else
  if y1 > y + h then false else
  true

let r1, m1 = (2, 0, 8, 1), "build"
let r2, m2 = (2, 2, 8, 1), "go-to"
let r3, m3 = (2, 4, 8, 1), "harvest"

let translate x1 y1 (x, y, w, h) =
  (x + x1, y + y1, w, h)

let cont_menu x y =
  let r1 = translate x y r1 in
  let r2 = translate x y r2 in
  let r3 = translate x y r3 in
  match !menu with
  | Some (p, ms) ->
      List.iter (fun (r, m) ->
        if in_rect (x, y) r then
          _tasks := (m, p) :: !_tasks
      ) ms;
      menu := None
  | None ->
      menu :=
        Some ((x, y), [
          (r1, m1);
          (r2, m2);
          (r3, m3);
        ])
;;

(* Events *)

let mousechange_event state ev =
  let rect = Canvas.getBoundingClientRect canvas in
  let x = ev.Canvas.clientX - rect.Canvas.left in
  let y = ev.Canvas.clientY - rect.Canvas.top in
  let x, y = (x / 10, y / 10) in
  cont_menu x y;
  ()

let keychange_event ev =
  ()


(* Main *)

let () =
  Random.self_init ();

  let _state = init_state () in
  let state = ref _state in

  let animate () =
    state := update_state !state;
    display_game !state;
    ()
  in

  Canvas.addKeyEventListener Canvas.window "keydown" keychange_event true;
  Canvas.addMouseEventListener Canvas.window "mousedown" (mousechange_event !state) true;

  let _ = Canvas.setInterval animate (1000/3) in
  ()
;;