open Space_net

let camera : Config.tenant =
  {
    name = "camera";
    apids = Config.apid_range 0x010 0x01F;
    can_send_to = [ Config.apid_range 0x020 0x02F ];
  }

let processor : Config.tenant =
  {
    name = "processor";
    apids = Config.apid_range 0x020 0x02F;
    can_send_to = [ Config.apid_range 0x010 0x01F ];
  }

let test_config socket_dir : Config.t =
  { tenants = [ camera; processor ]; socket_dir }

let routable_frame ~src_apid ~dest_apid typ payload =
  let base = Space_wire.Msg.v typ ~apid:src_apid payload in
  { base with reserved = dest_apid }

let read_frame reader =
  let s = Eio.Buf_read.take Space_wire.Msg.frame_size reader in
  Wire.Codec.decode Space_wire.Msg.codec (Bytes.of_string s) 0

let write_frame dst frame =
  let buf = Bytes.make Space_wire.Msg.frame_size '\x00' in
  Wire.Codec.encode Space_wire.Msg.codec frame buf 0;
  Eio.Flow.copy_string (Bytes.unsafe_to_string buf) dst

let with_timeout clock f =
  match Eio.Time.with_timeout clock 2.0 (fun () -> Ok (f ())) with
  | Ok v -> v
  | Error `Timeout -> Alcotest.fail "timeout"

let socket_dir prefix =
  let dir =
    Filename.concat
      (Filename.get_temp_dir_name ())
      (Fmt.str "%s-%d" prefix (Unix.getpid ()))
  in
  (try Unix.mkdir dir 0o755 with Unix.Unix_error (Unix.EEXIST, _, _) -> ());
  dir

let setup_switch ~sw ~net ~clock config =
  let switch = Switch.v ~config () in
  Switch.run switch ~sw ~net;
  let cam_path = Config.socket_path config camera in
  let proc_path = Config.socket_path config processor in
  let cam_sock = Eio.Net.connect ~sw net (`Unix cam_path) in
  let proc_sock = Eio.Net.connect ~sw net (`Unix proc_path) in
  Eio.Time.sleep clock 0.1;
  (switch, cam_sock, proc_sock)

let test_inter_guest () =
  Eio_main.run @@ fun env ->
  Eio.Switch.run @@ fun sw ->
  let clock = Eio.Stdenv.clock env in
  let net = Eio.Stdenv.net env in
  let sd = socket_dir "space-net-inter" in
  let config = test_config sd in
  let _switch, cam_sock, proc_sock = setup_switch ~sw ~net ~clock config in
  let proc_reader =
    Eio.Buf_read.of_flow
      ~max_size:(Space_wire.Msg.frame_size * 10)
      (proc_sock :> _ Eio.Flow.source)
  in
  let frame =
    routable_frame ~src_apid:0x010 ~dest_apid:0x020 TM "hello processor"
  in
  write_frame (cam_sock :> _ Eio.Flow.sink) frame;
  with_timeout clock (fun () ->
      let received = read_frame proc_reader in
      Alcotest.(check string)
        "payload" "hello processor"
        (Space_wire.Msg.payload_bytes received))

let test_source_enforcement () =
  Eio_main.run @@ fun env ->
  Eio.Switch.run @@ fun sw ->
  let clock = Eio.Stdenv.clock env in
  let net = Eio.Stdenv.net env in
  let sd = socket_dir "space-net-src" in
  let config = test_config sd in
  let _switch, cam_sock, _proc_sock = setup_switch ~sw ~net ~clock config in
  let cam_reader =
    Eio.Buf_read.of_flow
      ~max_size:(Space_wire.Msg.frame_size * 10)
      (cam_sock :> _ Eio.Flow.source)
  in
  let bad_frame =
    routable_frame ~src_apid:0x020 ~dest_apid:0x020 TM "spoofed"
  in
  write_frame (cam_sock :> _ Eio.Flow.sink) bad_frame;
  with_timeout clock (fun () ->
      let response = read_frame cam_reader in
      Alcotest.(check int)
        "error type"
        (Space_wire.Msg.kind_to_int ERROR)
        response.Space_wire.Msg.msg_type)

let test_inject () =
  Eio_main.run @@ fun env ->
  Eio.Switch.run @@ fun sw ->
  let clock = Eio.Stdenv.clock env in
  let net = Eio.Stdenv.net env in
  let sd = socket_dir "space-net-inj" in
  let config = test_config sd in
  let switch, cam_sock, _proc_sock = setup_switch ~sw ~net ~clock config in
  let cam_reader =
    Eio.Buf_read.of_flow
      ~max_size:(Space_wire.Msg.frame_size * 10)
      (cam_sock :> _ Eio.Flow.source)
  in
  let injected =
    { (Space_wire.Msg.v TC ~apid:0 "from DTN") with reserved = 0x010 }
  in
  Switch.inject switch injected;
  with_timeout clock (fun () ->
      let received = read_frame cam_reader in
      Alcotest.(check string)
        "injected payload" "from DTN"
        (Space_wire.Msg.payload_bytes received))

let test_system_handler () =
  Eio_main.run @@ fun env ->
  Eio.Switch.run @@ fun sw ->
  let clock = Eio.Stdenv.clock env in
  let net = Eio.Stdenv.net env in
  let sd = socket_dir "space-net-sys" in
  let config = test_config sd in
  let system_called = ref false in
  let on_system _frame =
    system_called := true;
    Some (Space_wire.Msg.v Space_wire.Msg.EVR ~apid:0x001 "system ack")
  in
  let switch = Switch.v ~config ~on_system () in
  Switch.run switch ~sw ~net;
  let cam_path = Config.socket_path config camera in
  let proc_path = Config.socket_path config processor in
  let cam_sock = Eio.Net.connect ~sw net (`Unix cam_path) in
  let _proc_sock = Eio.Net.connect ~sw net (`Unix proc_path) in
  Eio.Time.sleep clock 0.1;
  let cam_reader =
    Eio.Buf_read.of_flow
      ~max_size:(Space_wire.Msg.frame_size * 10)
      (cam_sock :> _ Eio.Flow.source)
  in
  let frame =
    routable_frame ~src_apid:0x010 ~dest_apid:0x001 TM "ping system"
  in
  write_frame (cam_sock :> _ Eio.Flow.sink) frame;
  with_timeout clock (fun () ->
      let response = read_frame cam_reader in
      Alcotest.(check string)
        "system response" "system ack"
        (Space_wire.Msg.payload_bytes response));
  Alcotest.(check bool) "system handler called" true !system_called

let suite =
  ( "switch",
    [
      Alcotest.test_case "inter-guest" `Quick test_inter_guest;
      Alcotest.test_case "source enforcement" `Quick test_source_enforcement;
      Alcotest.test_case "inject" `Quick test_inject;
      Alcotest.test_case "system handler" `Quick test_system_handler;
    ] )