open Uxn

type state = { banks : bytes array }

module Make () : Uxn.Device.DEVICE with type state = state = struct
  type nonrec state = state

  let state = { banks = Array.init 15 (fun _ -> Bytes.create 65536) }
  let can_handle port = port >= 0x00 && port <= 0x0f

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

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

  let 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 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 mach src_bank in
          let dst_mem = get_bank 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 mach src_bank in
          let dst_mem = get_bank 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
    | _ -> Format.eprintf "System/expansion: unknown command #%02x" cmd

  let dei m port =
    match port with
    | 0x04 ->
        let (Machine.Stack { sp; _ }) = Machine.wst m in
        Some sp
    | 0x05 ->
        let (Machine.Stack { sp; _ }) = Machine.rst m in
        Some sp
    | _ -> None

  let dei2 _ _ = None

  let deo mach port value =
    match port with
    | 0x02 -> expansion mach value
    | 0x04 ->
        let (Machine.Stack s) = Machine.wst mach in
        s.sp <- value land 0xff
    | 0x05 ->
        let (Machine.Stack s) = Machine.rst mach in
        s.sp <- value land 0xff
    | 0x0e ->
        if value <> 0 then begin
          print_stack ~name:"wst" (Machine.wst mach);
          print_stack ~name:"rst" (Machine.rst mach);
          Out_channel.flush stderr
        end
    | 0x0f -> Bytes.set_uint8 (Machine.dev mach) 0x0f value
    | _ -> ()
end