// socket_performance_fixed.wren - Fixed socket performance test
import "socket" for Socket, SocketInstance

foreign class Host {
  foreign static signalDone()
}

class SocketPerformance {
  static findFreePort(startPort) {
    for (port in startPort...(startPort + 100)) {
      var result = SocketInstance.new()
      if (result && !result[0]) {
        var sock = result[1]
        var bindResult = sock.bind("127.0.0.1", port)
        if (bindResult && !bindResult[0]) {
          sock.close()
          return port
        }
        sock.close()
      }
    }
    return 0
  }

  static runEchoTest() {
    System.print("\n=== Echo Performance Test ===")
    
    // Find a free port
    var port = findFreePort(25000)
    if (port == 0) {
      System.print("Could not find free port")
      return
    }
    System.print("Using port %(port)")
    
    // Create server
    var serverResult = SocketInstance.new()
    if (!serverResult || serverResult[0]) {
      System.print("Failed to create server")
      return
    }
    var serverSock = serverResult[1]
    
    // Bind and listen
    var bindResult = serverSock.bind("127.0.0.1", port)
    if (bindResult[0]) {
      System.print("Bind failed: %(bindResult[0])")
      serverSock.close()
      return
    }
    
    var listenResult = serverSock.listen(10)
    if (listenResult[0]) {
      System.print("Listen failed: %(listenResult[0])")
      serverSock.close()
      return
    }
    
    var messageCount = 1000
    var messageSize = 100
    var message = "X" * messageSize
    
    var serverDone = false
    var serverMessages = 0
    var serverBytes = 0
    
    // Server fiber
    var serverFiber = Fiber.new {
      var acceptResult = serverSock.accept()
      if (!acceptResult[0] && acceptResult[1]) {
        var clientSock = SocketInstance.fromFd(acceptResult[1])
        
        while (serverMessages < messageCount) {
          var readResult = clientSock.read(4096)
          if (!readResult[0] && readResult[1] && readResult[1].count > 0) {
            serverBytes = serverBytes + readResult[1].count
            serverMessages = serverMessages + 1
            
            // Echo back
            clientSock.write(readResult[1])
          }
          
          if (serverMessages % 10 == 0) {
            Fiber.yield()
          }
        }
        
        clientSock.close()
      }
      serverDone = true
    }
    
    var clientDone = false
    var clientMessages = 0
    var clientBytes = 0
    
    // Client fiber
    var clientFiber = Fiber.new {
      // Let server start
      for (i in 1..10) Fiber.yield()
      
      var connectResult = SocketInstance.connect("127.0.0.1", port)
      if (connectResult[0]) {
        System.print("Connect failed: %(connectResult[0])")
        clientDone = true
        return
      }
      var clientSock = connectResult[1]
      
      var startTime = System.clock
      
      // Send messages and read echoes
      for (i in 1..messageCount) {
        var writeErr = clientSock.write(message)
        if (!writeErr) {
          clientMessages = clientMessages + 1
          
          // Read echo (with timeout)
          var gotEcho = false
          for (attempt in 1..50) {
            var readResult = clientSock.read(4096)
            if (!readResult[0] && readResult[1] && readResult[1].count > 0) {
              clientBytes = clientBytes + readResult[1].count
              gotEcho = true
              break
            }
            if (attempt % 10 == 0) Fiber.yield()
          }
          
          if (!gotEcho && i > 10) {
            System.print("No echo received for message %(i)")
            break
          }
        }
        
        if (i % 10 == 0) Fiber.yield()
      }
      
      var elapsed = System.clock - startTime
      
      System.print("\nClient Statistics:")
      System.print("  Messages sent: %(clientMessages)")
      System.print("  Bytes received: %(clientBytes)")
      System.print("  Time: %(elapsed)s")
      System.print("  Throughput: %(clientMessages/elapsed) msgs/sec")
      System.print("  Bandwidth: %((clientBytes/elapsed/1024).floor) KB/sec")
      
      clientSock.close()
      clientDone = true
    }
    
    // Run both fibers
    serverFiber.call()
    clientFiber.call()
    
    var iterations = 0
    var maxIterations = 10000
    while ((!serverDone || !clientDone) && iterations < maxIterations) {
      if (!serverDone) serverFiber.call()
      if (!clientDone) clientFiber.call()
      iterations = iterations + 1
      Fiber.yield()
    }
    
    System.print("\nServer Statistics:")
    System.print("  Messages echoed: %(serverMessages)")
    System.print("  Bytes processed: %(serverBytes)")
    System.print("  Event loop iterations: %(iterations)")
    
    serverSock.close()
  }

  static runConcurrentTest() {
    System.print("\n=== Concurrent Connections Test ===")
    
    var port = findFreePort(26000)
    if (port == 0) {
      System.print("Could not find free port")
      return
    }
    System.print("Using port %(port)")
    
    var serverResult = SocketInstance.new()
    if (!serverResult || serverResult[0]) return
    var serverSock = serverResult[1]
    
    serverSock.bind("127.0.0.1", port)
    serverSock.listen(50)
    
    var connectionCount = 20
    var connectionsAccepted = 0
    var messagesSent = 0
    
    // Server fiber - handles all connections
    var serverDone = false
    var serverFiber = Fiber.new {
      var handlers = []
      var attempts = 0
      
      while (connectionsAccepted < connectionCount && attempts < 2000) {
        var acceptResult = serverSock.accept()
        if (!acceptResult[0] && acceptResult[1]) {
          connectionsAccepted = connectionsAccepted + 1
          var clientSock = SocketInstance.fromFd(acceptResult[1])
          
          // Handler for this connection
          var handler = Fiber.new {
            var msg = "Hello from server to connection %(connectionsAccepted)!"
            clientSock.write(msg)
            messagesSent = messagesSent + 1
            
            // Read client response
            for (i in 1..20) {
              var readResult = clientSock.read(256)
              if (!readResult[0] && readResult[1] && readResult[1].count > 0) {
                break
              }
              Fiber.yield()
            }
            
            clientSock.close()
          }
          
          handlers.add(handler)
          handler.call()
        }
        
        // Process active handlers
        var active = []
        for (h in handlers) {
          if (!h.isDone) {
            h.call()
            active.add(h)
          }
        }
        handlers = active
        
        attempts = attempts + 1
        Fiber.yield()
      }
      
      serverDone = true
    }
    
    // Create client fibers
    var clientFibers = []
    var clientsConnected = 0
    var messagesReceived = 0
    
    var startTime = System.clock
    
    for (i in 1..connectionCount) {
      var clientId = i
      var clientFiber = Fiber.new {
        // Stagger connections slightly
        for (j in 1..clientId) Fiber.yield()
        
        var connectResult = SocketInstance.connect("127.0.0.1", port)
        if (!connectResult[0]) {
          clientsConnected = clientsConnected + 1
          var sock = connectResult[1]
          
          // Read server message
          for (attempt in 1..50) {
            var readResult = sock.read(256)
            if (!readResult[0] && readResult[1] && readResult[1].count > 0) {
              messagesReceived = messagesReceived + 1
              // Send response
              sock.write("Client %(clientId) received: %(readResult[1])")
              break
            }
            Fiber.yield()
          }
          
          sock.close()
        }
      }
      clientFibers.add(clientFiber)
    }
    
    // Start all fibers
    serverFiber.call()
    for (cf in clientFibers) cf.call()
    
    // Run event loop
    var iterations = 0
    while (iterations < 5000) {
      if (!serverDone) serverFiber.call()
      
      var active = []
      for (cf in clientFibers) {
        if (!cf.isDone) {
          cf.call()
          active.add(cf)
        }
      }
      clientFibers = active
      
      if (serverDone && clientFibers.count == 0) break
      
      iterations = iterations + 1
      Fiber.yield()
    }
    
    var elapsed = System.clock - startTime
    
    System.print("\nResults:")
    System.print("  Connections attempted: %(connectionCount)")
    System.print("  Clients connected: %(clientsConnected)")
    System.print("  Connections accepted: %(connectionsAccepted)")
    System.print("  Messages sent by server: %(messagesSent)")
    System.print("  Messages received by clients: %(messagesReceived)")
    System.print("  Time: %(elapsed)s")
    System.print("  Connection rate: %(clientsConnected/elapsed) conn/sec")
    System.print("  Event loop iterations: %(iterations)")
    
    serverSock.close()
  }

  static runLatencyTest() {
    System.print("\n=== Latency Test (ping-pong) ===")
    
    var port = findFreePort(27000)
    if (port == 0) return
    
    var serverResult = SocketInstance.new()
    if (!serverResult || serverResult[0]) return
    var serverSock = serverResult[1]
    
    serverSock.bind("127.0.0.1", port)
    serverSock.listen(1)
    
    var pingCount = 100
    var serverDone = false
    
    // Server: respond to pings with pongs
    var serverFiber = Fiber.new {
      var acceptResult = serverSock.accept()
      if (!acceptResult[0] && acceptResult[1]) {
        var clientSock = SocketInstance.fromFd(acceptResult[1])
        
        for (i in 1..pingCount) {
          // Wait for ping
          var gotPing = false
          for (attempt in 1..100) {
            var readResult = clientSock.read(256)
            if (!readResult[0] && readResult[1] && readResult[1].count > 0) {
              // Send pong immediately
              clientSock.write("pong")
              gotPing = true
              break
            }
            Fiber.yield()
          }
          if (!gotPing) break
        }
        
        clientSock.close()
      }
      serverDone = true
    }
    
    var clientDone = false
    var latencies = []
    
    // Client: send pings and measure round-trip time
    var clientFiber = Fiber.new {
      for (i in 1..10) Fiber.yield()
      
      var connectResult = SocketInstance.connect("127.0.0.1", port)
      if (connectResult[0]) {
        clientDone = true
        return
      }
      var clientSock = connectResult[1]
      
      for (i in 1..pingCount) {
        var pingStart = System.clock
        
        // Send ping
        clientSock.write("ping")
        
        // Wait for pong
        var gotPong = false
        for (attempt in 1..100) {
          var readResult = clientSock.read(256)
          if (!readResult[0] && readResult[1] && readResult[1].count > 0) {
            var latency = (System.clock - pingStart) * 1000 // Convert to ms
            latencies.add(latency)
            gotPong = true
            break
          }
          Fiber.yield()
        }
        
        if (!gotPong) break
      }
      
      clientSock.close()
      clientDone = true
    }
    
    // Run test
    serverFiber.call()
    clientFiber.call()
    
    while (!serverDone || !clientDone) {
      if (!serverDone) serverFiber.call()
      if (!clientDone) clientFiber.call()
      Fiber.yield()
    }
    
    // Calculate statistics
    if (latencies.count > 0) {
      var sum = 0
      var min = latencies[0]
      var max = latencies[0]
      
      for (lat in latencies) {
        sum = sum + lat
        if (lat < min) min = lat
        if (lat > max) max = lat
      }
      
      var avg = sum / latencies.count
      
      System.print("\nLatency Statistics (%(latencies.count) samples):")
      System.print("  Average: %(avg) ms")
      System.print("  Min: %(min) ms")
      System.print("  Max: %(max) ms")
      System.print("  Ping rate: %(1000/avg) pings/sec possible")
    }
    
    serverSock.close()
  }
}

var mainFiber = Fiber.new {
  System.print("=== Socket Performance Test Suite ===")
  System.print("With fiber performance: 2.6M yields/sec (0.38μs/yield)")
  
  SocketPerformance.runEchoTest()
  SocketPerformance.runConcurrentTest()
  SocketPerformance.runLatencyTest()
  
  System.print("\n=== All Tests Complete ===")
  Host.signalDone()
  
  while(true) { Fiber.yield() }
}