open Space_wire

let log_src = Logs.Src.create "space-net.switch"

module Log = (val Logs.src_log log_src)

type t = {
  config : Config.t;
  router : Router.t;
  on_system : Msg.t -> Msg.t option;
  on_uplink : Msg.t -> unit;
  connections : (string, Connection.t) Hashtbl.t;
  mutable running : bool;
}

let pp fmt t =
  Fmt.pf fmt "Switch(%d tenants, %s)"
    (List.length t.config.tenants)
    t.config.socket_dir

let default_on_system frame =
  Log.info (fun m -> m "system frame: %a" Msg.pp frame);
  None

let default_on_uplink frame =
  Log.info (fun m -> m "uplink frame: %a" Msg.pp frame)

let v ~config ?(on_system = default_on_system) ?(on_uplink = default_on_uplink)
    () =
  let router = Router.of_config config in
  {
    config;
    router;
    on_system;
    on_uplink;
    connections = Hashtbl.create 16;
    running = true;
  }

let connection t (tenant : Config.tenant) =
  Hashtbl.find_opt t.connections tenant.name

(** In v0, the frame carries a source APID in [frame.apid] and a destination
    APID in [frame.reserved]. The source APID identifies the sending
    application; the destination APID determines routing. *)

let dest_apid (frame : Msg.t) = frame.reserved

let handle_frame t conn frame =
  let (frame : Msg.t) = frame in
  let source = Connection.tenant conn in
  (* 1. Validate version *)
  if frame.version <> 0x01 then
    Connection.write_error conn Unknown_version ~offending:frame
  else
    (* 2. Validate msg_type *)
    match Msg.kind_of_int frame.msg_type with
    | None -> Connection.write_error conn Unknown_type ~offending:frame
    | Some _ -> (
        (* 3. Validate source APID ownership *)
        match Router.validate_source t.router ~source ~apid:frame.apid with
        | Some (Apid_not_allocated _) ->
            Connection.write_error conn Apid_not_allocated ~offending:frame
        | Some (Policy_denied _) ->
            Connection.write_error conn Apid_not_allocated ~offending:frame
        | None -> (
            (* 4. Route by destination APID (carried in reserved field) *)
            let dest = dest_apid frame in
            match Router.route t.router ~source ~dest_apid:dest with
            | Error (Policy_denied { src_tenant; dest_apid }) ->
                Log.warn (fun m ->
                    m "policy denied: %s -> APID %d" src_tenant dest_apid);
                Connection.write_error conn Apid_not_allocated ~offending:frame
            | Error (Apid_not_allocated { apid; tenant }) ->
                Log.warn (fun m -> m "APID %d not allocated to %s" apid tenant);
                Connection.write_error conn Apid_not_allocated ~offending:frame
            | Ok (Router.Local dest_tenant) -> (
                match connection t dest_tenant with
                | Some dest_conn -> Connection.write_frame dest_conn frame
                | None ->
                    Log.warn (fun m ->
                        m "dest %s not connected" dest_tenant.Config.name))
            | Ok Router.System -> (
                match t.on_system frame with
                | Some response -> Connection.write_frame conn response
                | None -> ())
            | Ok Router.Uplink -> t.on_uplink frame
            | Ok (Router.Drop reason) ->
                Log.debug (fun m -> m "drop: %s (apid=%d)" reason dest)))

let inject t frame =
  let dest = dest_apid frame in
  match Router.lookup t.router dest with
  | Router.Local tenant -> (
      match connection t tenant with
      | Some conn -> Connection.write_frame conn frame
      | None ->
          Log.warn (fun m ->
              m "inject: dest %s not connected" tenant.Config.name))
  | Router.System -> ( match t.on_system frame with Some _ -> () | None -> ())
  | Router.Uplink -> t.on_uplink frame
  | Router.Drop reason ->
      Log.debug (fun m -> m "inject drop: %s (apid=%d)" reason dest)

let stop t = t.running <- false

let run t ~sw ~net =
  (try Unix.mkdir t.config.socket_dir 0o755
   with Unix.Unix_error (Unix.EEXIST, _, _) -> ());
  List.iter
    (fun (tenant : Config.tenant) ->
      let path = Config.socket_path t.config tenant in
      Log.info (fun m -> m "listening for %s on %s" tenant.name path);
      Connection.listen ~sw ~net ~path ~tenant
        ~on_connect:(fun conn -> Hashtbl.replace t.connections tenant.name conn)
        ~on_frame:(fun conn frame -> handle_frame t conn frame)
        ~on_disconnect:(fun _conn ->
          Log.info (fun m -> m "guest %s disconnected" tenant.name);
          Hashtbl.remove t.connections tenant.name))
    t.config.tenants

let space_packet_uplink ~seq_count frame =
  let apid = frame.Msg.apid in
  if apid < 0 || apid > Space_packet.apid_max then None
  else
    let data =
      let buf = Bytes.make Msg.frame_size '\x00' in
      Wire.Codec.encode Msg.codec frame buf 0;
      Bytes.unsafe_to_string buf
    in
    match
      Space_packet.v ~packet_type:Telemetry ~apid ~sequence_flags:Unsegmented
        ~sequence_count:(seq_count ()) data
    with
    | Ok pkt -> Some pkt
    | Error _ -> None