(* Types *)

type tile = int * int * int

type terrain = {
  grid: tile 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;
}

(* Init-functions *)

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

let new_terrain width height =
  let grid = Array.init height (fun y -> Array.init width (fun x -> new_tile x y 0)) in
  { grid; 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)


(* Tasks *)

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

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

let new_task state unite =
  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, "Habitat") unite
  | _ ->
      let resources =
        List.sort (fun r1 r2 ->
          let d1 = sq_dist r1.r_pos unite.u_pos in
          let d2 = sq_dist r2.r_pos unite.u_pos in
          compare d1 d2
        ) state.resources
      in
      match resources with
      | [] -> unite
      | r :: _ ->
          let x, y = r.r_pos in
          add_task ("Harvest", x, y, "") unite


(* Build *)

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

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


(* 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 unite tl =
  if pos = unite.u_pos
  then
    let b = new_building pos (rand_take building_kind) 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, bld) :: tl ->
            let state, u = proceed_build state (x, y) 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 180 = 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 24 = 0
  then r :: resources
  else resources


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
  (state)


(* Init *)

let init_state () =
  let pos1 = (Random.int 68, Random.int 44) in
  let pos2 = (Random.int 68, Random.int 44) in
  let unite1 = new_unite pos1 "worker" 0 in
  let resource1 = new_resource pos2 "elementium" in
  let unite1 = add_task ("GoTo", 18, 11, "") unite1 in
  {
    terrain = new_terrain 68 44;
    resources = [resource1];
    unites = unite1 :: [];
    buildings = [];
  }


(* 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;
  ()
;;

(* 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;
;;

(* Display *)

let display_game state =

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

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

  ()
;;

(* Events *)

let keychange_event ev =
  ()

let mousechange_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 true;

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