open Uxn
open Effect.Deep

let debug = Option.is_some (Sys.getenv_opt "DBG")
let banks = Array.init 15 (fun _ -> Bytes.create 65536)

let get_bank_memory mach bank =
  if bank = 0 then Machine.ram mach
  else if bank > 0 && bank < 16 then banks.(bank - 1)
  else Bytes.create 0

let system_expansion mach cmd_addr =
  let ram = Machine.ram mach in
  let cmd = Bytes.get_uint8 ram cmd_addr in
  match cmd with
  | 0x00 ->
      let length = Util.get_uint16_wrap ~wrap:0xffff ram (cmd_addr + 1) in
      let bank = Bytes.get_uint16_be ram (cmd_addr + 3) in
      let addr = Util.get_uint16_wrap ~wrap:0xffff ram (cmd_addr + 5) in
      let value = Bytes.get_uint8 ram (cmd_addr + 7) in
      if bank < 16 then begin
        let mem = get_bank_memory mach bank in
        for i = 0 to length - 1 do
          let pos = (addr + i) land 0xffff in
          Bytes.set_uint8 mem pos value
        done
      end
  | 0x01 ->
      let length = Util.get_uint16_wrap ~wrap:0xffff ram (cmd_addr + 1) in
      let src_bank = Bytes.get_uint16_be ram (cmd_addr + 3) in
      let src_addr = Util.get_uint16_wrap ~wrap:0xffff ram (cmd_addr + 5) in
      let dst_bank = Bytes.get_uint16_be ram (cmd_addr + 7) in
      let dst_addr = Util.get_uint16_wrap ~wrap:0xffff ram (cmd_addr + 9) in
      if src_bank < 16 && dst_bank < 16 then begin
        let src_mem = get_bank_memory mach src_bank in
        let dst_mem = get_bank_memory mach dst_bank in
        for i = 0 to length - 1 do
          let src_pos = (src_addr + i) land 0xffff in
          let dst_pos = (dst_addr + i) land 0xffff in
          let v = Bytes.get_uint8 src_mem src_pos in
          Bytes.set_uint8 dst_mem dst_pos v
        done
      end
  | 0x02 ->
      let length = Util.get_uint16_wrap ~wrap:0xffff ram (cmd_addr + 1) in
      let src_bank = Bytes.get_uint16_be ram (cmd_addr + 3) in
      let src_addr = Util.get_uint16_wrap ~wrap:0xffff ram (cmd_addr + 5) in
      let dst_bank = Bytes.get_uint16_be ram (cmd_addr + 7) in
      let dst_addr = Util.get_uint16_wrap ~wrap:0xffff ram (cmd_addr + 9) in
      if src_bank < 16 && dst_bank < 16 then begin
        let src_mem = get_bank_memory mach src_bank in
        let dst_mem = get_bank_memory mach dst_bank in
        for i = length - 1 downto 0 do
          let src_pos = (src_addr + i) land 0xffff in
          let dst_pos = (dst_addr + i) land 0xffff in
          let v = Bytes.get_uint8 src_mem src_pos in
          Bytes.set_uint8 dst_mem dst_pos v
        done
      end
  | _ -> Fmt.epr "System/expansion: unknown command #%02x" cmd

let print_stack ~name (Machine.Stack { data; sp }) =
  Fmt.epr "%s: @[%a@]@." name
    (Fmt.on_bytes (Fmt.octets ()))
    (Bytes.sub data 0 sp)

let rec run m pc =
  let dev = Machine.dev m in
  let console_vector = ref 0 in
  let console_input ch ty =
    Bytes.set_uint8 dev 0x12 ch;
    Bytes.set_uint8 dev 0x17 ty;
    if Bytes.get_uint8 dev 0x0f = 0 then run m !console_vector
  in
  try Machine.dispatch ~trace:debug m pc with
  | effect Machine.Trace (pc, instr, args), k when debug ->
      Fmt.epr "PC = %04x | %6s %a@." pc (Instr.to_string instr)
        (Fmt.list ~sep:(Fmt.any " ") (Fmt.fmt "%02x"))
        args;
      Out_channel.flush stderr;
      continue k ()
  | effect Machine.Trace _, k -> continue k ()
  | effect Machine.BRK, _ when !console_vector != 0 -> (
      try
        while Bytes.get_uint8 dev 0x0f = 0 do
          match In_channel.input_byte stdin with
          | None -> raise Exit
          | Some c -> console_input c 1
        done
      with Exit -> console_input 0 4)
  | effect Machine.BRK, _ -> ()
  | effect Machine.DEI (`Byte, port), k ->
      let value =
        match port with
        | 0x04 ->
            let (Machine.Stack { sp; _ }) = Machine.wst m in
            sp
        | 0x05 ->
            let (Machine.Stack { sp; _ }) = Machine.rst m in
            sp
        | _ -> Bytes.get_uint8 dev port
      in
      continue k value
  | effect Machine.DEI (`Short, port), k ->
      continue k (Util.get_uint16_wrap ~wrap:0xffff dev port)
  | effect Machine.DEO (port, value), k ->
      (match port with
      | 0x02 -> system_expansion m value
      | 0x04 ->
          let (Machine.Stack s) = Machine.wst m in
          s.sp <- value land 0xff
      | 0x05 ->
          let (Machine.Stack s) = Machine.rst m in
          s.sp <- value land 0xff
      | 0x0e ->
          if value <> 0 then begin
            print_stack ~name:"WST" (Machine.wst m);
            print_stack ~name:"RST" (Machine.rst m);
            Out_channel.flush stderr
          end
      | 0x0f -> Bytes.set_uint8 dev 0x0f value
      | 0x10 -> console_vector := value
      | 0x18 -> print_char (Char.chr value)
      | 0x19 -> prerr_char (Char.chr value)
      | _ -> ());
      continue k ()

let main () =
  if Array.length Sys.argv < 2 then (
    Fmt.epr "usage: uxnemu file.rom ...\n";
    exit 1);

  let code =
    In_channel.with_open_bin Sys.argv.(1) (fun i -> In_channel.input_all i)
  in

  In_channel.set_binary_mode stdin true;
  Out_channel.set_binary_mode stdout true;

  let mach = Machine.create code in
  Bytes.set (Machine.dev mach) 0 '\x00';
  run mach 0x100;
  if debug then begin
    print_stack ~name:"wst" (Machine.wst mach);
    print_stack ~name:"rst" (Machine.rst mach)
  end;
  exit (Bytes.get_uint8 (Machine.dev mach) 0x0f land 0x7f)

let _ = main ()