text: # Licensed to the Apache Software Foundation (ASF) under one
  # or more contributor license agreements.  See the NOTICE file
  # distributed with this work for additional information
  # regarding copyright ownership.  The ASF licenses this file
  # to you under the Apache License, Version 2.0 (the
  # "License"); you may not use this file except in compliance
  # with the License.  You may obtain a copy of the License at
  #
  #   http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing,
  # software distributed under the License is distributed on an
  # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  # KIND, either express or implied.  See the License for the
  # specific language governing permissions and limitations
  # under the License.
  # pylint: disable=redefined-outer-name, invalid-name
  """Start an RPC server intended for use as a microTVM debugger.
  
  microTVM aims to be runtime-agnostic, and to that end, frameworks often define command-line tools
  used to launch a debug flow. These tools often manage the process of connecting to an attached
  device using a hardware debugger, exposing a GDB server, and launching GDB connected to that
  server with a source file attached. It's also true that this debugger can typically not be executed
  concurrently with any flash tool, so this integration point is provided to allow TVM to launch and
  terminate any debuggers integrated with the larger microTVM compilation/autotuning flow.
  
  To use this tool, first launch this script in a separate terminal window. Then, provide the hostport
  to your compiler's Flasher instance.
  """
  
  import argparse
  import logging
  import socket
  import struct
  
  import tvm.micro.debugger as _  # NOTE: imported to expose global PackedFuncs over RPC.
  
  from .._ffi.base import py_str
  from ..rpc import base
  from ..rpc import _ffi_api
  
  
  _LOG = logging.getLogger(__name__)
  
  
  def parse_args():
      """Parse command line arguments to this script."""
      parser = argparse.ArgumentParser(description="microTVM debug-tool runner")
      parser.add_argument("--host", default="0.0.0.0", help="hostname to listen on")
      parser.add_argument("--port", type=int, default=9090, help="hostname to listen on")
      parser.add_argument(
          "--impl",
          help=(
              "If given, name of a module underneath tvm.micro.contrib "
              "which contains the Debugger implementation to use. For example, to enable a "
              "debugger named BarDebugger in python/tvm/micro/contrib/foo.py, specify either "
              "'tvm.micro.contrib.foo' or 'foo' here. To enable a debugger named BazDebugger in "
              "a third-party module ext_package.debugger, specify 'ext_package.debugger' here. "
              "NOTE: the module cannot be in a sub-package of tvm.micro.contrib."
          ),
      )
  
      return parser.parse_args()
  
  
  class ConnectionClosedError(Exception):
      """Raised when the connection is closed."""
  
  
  def handle_conn(conn, rpc_key):
      """Handle a single connection that has just been accept'd()."""
  
      def send(data):
          conn.sendall(data)
          return len(data)
  
      magic = struct.unpack("<i", base.recvall(conn, 4))[0]
      if magic != base.RPC_MAGIC:
          conn.close()
          return
  
      keylen = struct.unpack("<i", base.recvall(conn, 4))[0]
      key = py_str(base.recvall(conn, keylen))
      arr = key.split()
      expect_header = "client:"
      server_key = "server:" + rpc_key
      if arr[0] != expect_header:
          conn.sendall(struct.pack("<i", base.RPC_CODE_MISMATCH))
          _LOG.warning("mismatch key from %s", addr)
          return
  
      conn.sendall(struct.pack("<i", base.RPC_CODE_SUCCESS))
      conn.sendall(struct.pack("<i", len(server_key)))
      conn.sendall(server_key.encode("utf-8"))
      server = _ffi_api.CreateEventDrivenServer(send, "microtvm-rpc-debugger", key)
  
      def _readall(n):
          buf = bytearray()
          while len(buf) < n:
              x = conn.recv(n - len(buf))
              if not x:
                  raise ConnectionClosedError()
  
              buf = buf + x
  
          return buf
  
      while True:
          packet_length_bytes = _readall(8)
          packet_length = struct.unpack("<q", packet_length_bytes)[0]
          if not packet_length:
              break
  
          status = server(packet_length_bytes, 3)
          if status == 0:
              break
  
          packet_body = _readall(packet_length)
          status = server(packet_body, 3)
  
  
  def main():
      """Main entry point for microTVM debug shell."""
      args = parse_args()
      logging.basicConfig(level=logging.INFO)
      if args.impl:
          package = None
          if "." not in args.impl:
              package = f"tvm.micro.contrib.{args.impl}"
          importlib.import_module(args.impl, package)
  
      sock = socket.socket(base.get_addr_family([args.host, args.port]), socket.SOCK_STREAM)
      sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
      sock.bind((args.host, args.port))
      sock.listen(1)
      bind_addr, bind_port = sock.getsockname()
      _LOG.info("listening for connections on %s:%d", bind_addr, bind_port)
  
      while True:
          conn, peer = sock.accept()
          _LOG.info("accepted connection from %s", peer)
          try:
              handle_conn(conn, "")
          except ConnectionClosedError:
              pass
          finally:
              conn.close()
              _LOG.info("closed connection from %s", peer)
  
  
  if __name__ == "__main__":
      main()
  