(* This Source Code Form is subject to the terms of the Mozilla Public
   License, v. 2.0. If a copy of the MPL was not distributed with this
   file, You can obtain one at https://mozilla.org/MPL/2.0/. *)

open MParser
open Common
open Utils

let ws = skip_many1 blank

let string_return str result = string str >>$ result

let p_registers =
  choice
    [
      string_return "X" X; string_return "T" T; string_return "F" F; string_return "M" M;
    ]

let p_R = ws >> p_registers

let p_RN_ =
  p_registers |>> (fun r -> R r) <|> (many1_chars digit |>> fun n -> N (int_of_string n))

let p_RN = ws >> p_RN_

let p_label = ws >> many1_chars uppercase

let p_test_op =
  ws >> choice [ string_return "<" `LT; string_return "=" `EQ; string_return ">" `GT ]

let comment = string ";" >> skip_many_chars_until any_char newline

let instructions =
  [
    string_return "NOOP" (Some NOOP);
    ( string "COPY" >> p_RN >>= fun rn ->
      p_R >>= fun r ->
      match (rn, r) with
      | R src, dest when Register.equal src dest -> fail "CANNOT REPEAT REGISTERS"
      | _ -> return (Some (COPY (rn, r))) );
    (* Math *)
    string "ADDI" >> pipe3 p_RN p_RN p_R (fun src i dest -> Some (ADDI (src, i, dest)));
    string "SUBI" >> pipe3 p_RN p_RN p_R (fun src i dest -> Some (SUBI (src, i, dest)));
    string "MULI" >> pipe3 p_RN p_RN p_R (fun src i dest -> Some (MULI (src, i, dest)));
    string "DIVI" >> pipe3 p_RN p_RN p_R (fun src i dest -> Some (DIVI (src, i, dest)));
    string "MODI" >> pipe3 p_RN p_RN p_R (fun src i dest -> Some (MODI (src, i, dest)));
    string "SWIZ" >> pipe3 p_RN p_RN p_R (fun src i dest -> Some (SWIZ (src, i, dest)));
    (* Branching *)
    (string "MARK" >> p_label |>> fun label -> Some (MARK label));
    (string "JUMP" >> p_label |>> fun label -> Some (JUMP label));
    (string "TJMP" >> p_label |>> fun label -> Some (TJMP label));
    (string "FJMP" >> p_label |>> fun label -> Some (FJMP label));
    (* Testing *)
    string "TEST" >> ws
    >> choice
         [
           attempt (string "MRD" >>$ Some TEST_MRD);
           attempt (string "EOF" >>$ Some TEST_EOF);
           pipe3 p_RN_ p_test_op p_RN (fun left op right ->
               Some
                 (match op with
                 | `LT -> TEST_LT (left, right)
                 | `EQ -> TEST_EQ (left, right)
                 | `GT -> TEST_GT (left, right)));
         ];
    (* Lifecycle *)
    (string "REPL" >> p_label |>> fun label -> Some (REPL label));
    string_return "HALT" (Some HALT);
    (* Movement *)
    (string "LINK" >> p_RN |>> fun rn -> Some (LINK rn));
    (string "HOST" >> p_R |>> fun r -> Some (HOST r));
    (* Communication *)
    string_return "MODE" (Some MODE);
    ( string "VOID" >> ws >> choice [ string_return "M" `M; string_return "F" `F ]
    |>> fun mf ->
      Some
        (match mf with
        | `M -> VOID_M
        | `F -> VOID_F) );
    (* File Manipulation *)
    string_return "MAKE" (Some MAKE);
    (string "GRAB" >> p_RN |>> fun rn -> Some (GRAB rn));
    (string "FILE" >> p_R |>> fun r -> Some (FILE r));
    (string "SEEK" >> p_RN |>> fun rn -> Some (SEEK rn));
    string_return "DROP" (Some DROP);
    string_return "WIPE" (Some WIPE);
    string "NOTE" >> skip_many_chars_until any_char newline >>$ None;
    many1 blank >>$ None;
  ]
  |> List.map (fun inst -> inst << optional comment << optional newline)

let exa_grammar = many (spaces >> ((comment >>$ None) <|> choice instructions))

let parse_code str =
  match parse_string exa_grammar (String.trim str) () with
  | Success result ->
      Ok
        (result
        |> List.mapi (fun idx line -> (idx, line))
        |> List.filter_map (fun (idx, line) ->
               match line with
               | None -> None
               | Some line -> Some (idx, line)))
  | Failed (error, _) ->
      Error error

type op_loc = {
  inst : op_code;
  line_num : int;
}

let parse_code (input : string) : code =
  let instructions = parse_code input in
  let labels = StringMap.create 0 in
  let marks = ref StringSet.empty in
  let targets = ref StringSet.empty in
  let ops = ref [] in
  let errors = ref [] in
  (match instructions with
  | Error err -> errors := { msg = err; line_num = 0; exa = None } :: !errors
  | Ok instructions ->
      List.iter
        (fun (line_num, line) ->
          match line with
          | MARK label as inst ->
              marks := StringSet.add label !marks;
              StringMap.replace labels label line_num;
              ops := { inst; line_num } :: !ops
          | (JUMP label as inst)
          | (TJMP label as inst)
          | (FJMP label as inst) ->
              targets := StringSet.add label !targets;
              ops := { inst; line_num } :: !ops
          | inst -> ops := { inst; line_num } :: !ops)
        instructions);
  let extra_targets = StringSet.diff !targets !marks in
  if not (StringSet.is_empty extra_targets) then
    !ops
    |> List.filter (fun { inst; _ } ->
           match inst with
           | JUMP label
           | TJMP label
           | FJMP label ->
               StringSet.find_opt label extra_targets |> Option.is_some
           | _ -> false)
    |> List.iter (fun { inst; line_num } ->
           match inst with
           | JUMP label
           | TJMP label
           | FJMP label ->
               let msg = Printf.sprintf "MISMATCHED TARGET: %s" label in
               errors := { msg; line_num; exa = None } :: !errors
           | _ -> ());
  let ops = !ops |> List.rev |> List.map (fun { inst; _ } -> inst) in
  { ops; labels; errors = List.rev !errors }