import
  ws, asyncdispatch, asynchttpserver, strformat, asyncnet, httpclient,
  strutils, json, uri
import wsabi/[httputils, args, stats]


type
  Client = object
    ws: WebSocket
    remote: WebSocket
    netAddr: string
    connected: bool

var
  server = newAsyncHttpServer()
  client: Client
  clients = newSeq[Client]()


proc commRemoteClient(c: Client) {.async.} =
  ## Fetch from remote and send to client
  try:
    while c.remote.readyState == Open and c.connected:
      var data = await c.remote.receiveStrPacket()
      await parseStanza(data, client.netAddr, false)
      echo data
      await c.ws.send(data)
  except WebSocketError:
    echo getCurrentExceptionMsg()

proc localServer(req: Request) {.async, gcsafe.} =
  ## Listen on localhost:PORT/ws
  if req.url.path == "/ws":
    var newreq = req

    # TODO: loop through for all remote hosts
    let bits = parseUri(remoteHost[0])
    newreq.headers["host"] = @[fmt"{bits.hostname}:{bits.port}"]
    newreq.headers["hostname"] = @[bits.hostname]

    try:
      client.ws = await newWebSocket(newreq, protocol = "xmpp")
      client.connected = true
      echo "connecting to remote host..."
      client.remote = await newWebSocket(remoteHost[0], protocol = "xmpp")
      let (address, port) = client.remote.tcpSocket.getPeerAddr()
      echo fmt"connected to {address}:{port.int}"
      clients.add client

      while client.ws.readyState == Open:
        client.netAddr = client.ws.tcpSocket.getPeerAddr()[0]
        var
          packet = await client.ws.receiveStrPacket()
          repacket = packet.replace(client.netAddr, by=bits.hostname)
        echo repacket
        await parseStanza(repacket, client.netAddr, true)

        await client.remote.send(repacket)
        echo "sent local packet"
        asyncCheck commRemoteClient(client)
    except WebSocketError:
      echo "socket closed: ", getCurrentExceptionMsg()
#       client.ws.close()
#       client.remote.close()
#      echo "client closed socket: ", e.msg

  ## Health check endpoint 
  ## The 'uri' can be either 'ws' or 'http'.
  if req.url.path == "/check-health":
    if req.reqMethod == HttpPost:
      type 
        HealthCheck = object
          uri: string
          host: string
          port: int
      try:
        let
          hcJson = parseJson(req.body)
          hc = to(hcJson, HealthCheck)

        ## Check if host is a websocket, and try for the xmpp protocol
        if hc.uri in ["ws", "wss"]:
          let hostStr = &"{hc.uri}://{hc.host}:{hc.port}/ws/"
          echo &"connecting to {hostStr}"
          var ws = await newWebSocket(hostStr, protocol = "xmpp")
          echo "sending test packet"
          echo &"""<open xmlns="urn:ietf:params:xml:ns:xmpp-framing" to="{hc.host}" version="1.0" />"""
          await ws.send(&"""<open xmlns="urn:ietf:params:xml:ns:xmpp-framing" to="{hc.host}" version="1.0" />""")

          echo await ws.receiveStrPacket()
          ws.close()
          let r = %*{"status": "OK"}

          await req.makeResponse(Http200, $r)

        if hc.uri in ["http", "https"]:
          let hostStr = &"{hc.uri}://{hc.host}:{hc.port}/"
          echo &"connecting to {hostStr}"
          var tmpClient = newAsyncHttpClient()
          echo await tmpClient.getContent(hostStr)
          let r = %*{"status": "OK"}

          await req.makeResponse(Http200, $r)

        else:
          let r = %*{"status": "error", "msg": &"unknown URI {hc.uri}"}
          await req.makeResponse(Http400, $r)

      except OSError as e:
        echo "unable to reach specified host:port pair"
        let r = %*{
          "status": "error",
          "msg": e.msg
        }
        await req.makeResponse(Http400, $r)

      except KeyError as e:
        let r = %*{"status": "error", "msg": e.msg}
        await req.makeResponse(Http400, $r)

      except WebSocketError as e:
        let r = %*{"status": "error", "msg": e.msg}
        await req.makeResponse(Http500, $r)


when isMainModule:
  parseArgs()
  echo fmt"local server running at ws://127.0.0.1:{localPort}/ws"
  asyncCheck server.serve(Port(localPort), localServer)

  runForever()