## Replication — sync and async replication between nodes
import std/tables
import std/sets
import std/locks
import std/net
import std/strutils
import std/nativesockets
import std/monotimes
import std/asyncdispatch


type
  ReplicationMode* = enum
    rmSync    # synchronous — wait for all replicas
    rmAsync   # asynchronous — fire and forget
    rmSemiSync # semi-sync — wait for at least N replicas

  ReplicaState* = enum
    rsConnecting
    rsStreaming
    rsLagging
    rsDisconnected

  Replica* = ref object
    id*: string
    host*: string
    port*: int
    state*: ReplicaState
    lastAckLsn*: uint64
    lagBytes*: int
    lagTime*: int64  # nanoseconds
    lastSeen*: int64
    connected*: bool

  ReplicationManager* = ref object
    lock: Lock
    mode*: ReplicationMode
    replicas*: Table[string, Replica]
    currentLsn*: uint64
    syncReplicaCount*: int  # for semi-sync
    pendingAcks*: Table[uint64, HashSet[string]]  # lsn -> replica_ids waiting
    appliedLsn*: uint64

proc newReplica*(id: string, host: string, port: int): Replica =
  Replica(
    id: id, host: host, port: port,
    state: rsConnecting, lastAckLsn: 0,
    lagBytes: 0, lagTime: 0, lastSeen: 0,
    connected: false,
  )

proc newReplicationManager*(mode: ReplicationMode = rmAsync,
                            syncCount: int = 1): ReplicationManager =
  new(result)
  initLock(result.lock)
  result.mode = mode
  result.replicas = initTable[string, Replica]()
  result.currentLsn = 0
  result.syncReplicaCount = syncCount
  result.pendingAcks = initTable[uint64, HashSet[string]]()
  result.appliedLsn = 0

proc addReplica*(rm: ReplicationManager, replica: Replica) =
  acquire(rm.lock)
  rm.replicas[replica.id] = replica
  release(rm.lock)

proc removeReplica*(rm: ReplicationManager, id: string) =
  acquire(rm.lock)
  rm.replicas.del(id)
  release(rm.lock)

proc connectReplica*(rm: ReplicationManager, id: string) =
  acquire(rm.lock)
  if id in rm.replicas:
    rm.replicas[id].state = rsStreaming
    rm.replicas[id].connected = true
  release(rm.lock)

proc connectWithTimeout(sock: Socket, host: string, port: Port, timeoutMs: int): bool =
  ## Non-blocking connect with timeout to avoid hanging on unreachable hosts.
  sock.getFd.setBlocking(false)
  try:
    sock.connect(host, port)
    sock.getFd.setBlocking(true)
    return true
  except OSError:
    var fds = @[sock.getFd]
    if selectWrite(fds, timeoutMs) <= 0:
      return false
    sock.getFd.setBlocking(true)
    return true

proc shipToReplica(replica: Replica, lsn: uint64, data: seq[byte]): bool =
  ## Send replication data to a replica via TCP.
  ## Protocol: "REP <lsn> <dataLen>\n<data>"
  ## Response: "ACK <lsn>\n" on success
  try:
    var sock = newSocket()
    if not connectWithTimeout(sock, replica.host, Port(replica.port), 500):
      sock.close()
      return false
    let header = "REP " & $lsn & " " & $data.len & "\n"
    sock.send(header)
    if data.len > 0:
      sock.send(cast[string](data))
    var response = ""
    sock.readLine(response)
    sock.close()
    let parts = response.strip().split(" ")
    return parts.len >= 2 and parts[0] == "ACK"
  except:
    return false

proc writeLsn*(rm: ReplicationManager, data: seq[byte]): uint64 =
  acquire(rm.lock)
  inc rm.currentLsn
  let lsn = rm.currentLsn

  var replicasToShip: seq[Replica]
  for id, replica in rm.replicas:
    if replica.connected and replica.host.len > 0 and replica.port > 0:
      replicasToShip.add(replica)

  case rm.mode
  of rmAsync:
    release(rm.lock)
    for replica in replicasToShip:
      discard shipToReplica(replica, lsn, data)
    return lsn
  of rmSync:
    if replicasToShip.len > 0:
      rm.pendingAcks[lsn] = initHashSet[string]()
      for replica in replicasToShip:
        rm.pendingAcks[lsn].incl(replica.id)
    release(rm.lock)
    var ackCount = 0
    for replica in replicasToShip:
      if shipToReplica(replica, lsn, data):
        inc ackCount
    if replicasToShip.len > 0 and ackCount < replicasToShip.len:
      # Not all replicas acked — log but still return LSN (caller decides)
      when defined(debug):
        echo "Replication sync: only ", ackCount, "/", replicasToShip.len, " replicas acked for LSN ", lsn
    return lsn
  of rmSemiSync:
    if replicasToShip.len > 0:
      rm.pendingAcks[lsn] = initHashSet[string]()
      var count = 0
      for replica in replicasToShip:
        if count < rm.syncReplicaCount:
          rm.pendingAcks[lsn].incl(replica.id)
          inc count
    release(rm.lock)
    var ackCount = 0
    for replica in replicasToShip:
      if shipToReplica(replica, lsn, data):
        inc ackCount
        if ackCount >= rm.syncReplicaCount:
          break
    if replicasToShip.len > 0 and ackCount == 0 and rm.syncReplicaCount > 0:
      when defined(debug):
        echo "Replication semi-sync: no replicas acked for LSN ", lsn
    return lsn

proc ackLsn*(rm: ReplicationManager, replicaId: string, lsn: uint64) =
  acquire(rm.lock)
  if replicaId in rm.replicas:
    rm.replicas[replicaId].lastAckLsn = max(rm.replicas[replicaId].lastAckLsn, lsn)
    rm.replicas[replicaId].lagBytes = int(rm.currentLsn - lsn)

  if lsn in rm.pendingAcks:
    rm.pendingAcks[lsn].excl(replicaId)
    if rm.pendingAcks[lsn].len == 0:
      rm.pendingAcks.del(lsn)
      rm.appliedLsn = max(rm.appliedLsn, lsn)

  release(rm.lock)

proc isFullyAcked*(rm: ReplicationManager, lsn: uint64): bool =
  acquire(rm.lock)
  result = lsn notin rm.pendingAcks
  release(rm.lock)

proc minAckLsn*(rm: ReplicationManager): uint64 =
  acquire(rm.lock)
  result = rm.appliedLsn
  release(rm.lock)

proc connectedReplicaCount*(rm: ReplicationManager): int =
  acquire(rm.lock)
  result = 0
  for id, replica in rm.replicas:
    if replica.connected:
      inc result
  release(rm.lock)

proc totalReplicaCount*(rm: ReplicationManager): int =
  acquire(rm.lock)
  result = rm.replicas.len
  release(rm.lock)

proc maxLag*(rm: ReplicationManager): int =
  acquire(rm.lock)
  result = 0
  for id, replica in rm.replicas:
    if replica.lagBytes > result:
      result = replica.lagBytes
  release(rm.lock)

proc replicaStatus*(rm: ReplicationManager): seq[(string, ReplicaState, int)] =
  acquire(rm.lock)
  result = @[]
  for id, replica in rm.replicas:
    result.add((id, replica.state, replica.lagBytes))
  release(rm.lock)

proc switchMode*(rm: ReplicationManager, mode: ReplicationMode) =
  acquire(rm.lock)
  rm.mode = mode
  release(rm.lock)

# ---------------------------------------------------------------------------
# Health check and reconnection
# ---------------------------------------------------------------------------

proc healthCheck*(rm: ReplicationManager) =
  acquire(rm.lock)
  for id, replica in rm.replicas:
    if replica.connected:
      # Probe connection by sending a heartbeat
      var sock = newSocket()
      if not connectWithTimeout(sock, replica.host, Port(replica.port), 1000):
        replica.connected = false
        replica.state = rsDisconnected
      else:
        sock.send("PING\n")
        var response = ""
        try:
          sock.readLine(response)
          if response.strip() != "PONG":
            replica.connected = false
            replica.state = rsDisconnected
        except:
          replica.connected = false
          replica.state = rsDisconnected
        sock.close()
  release(rm.lock)

proc reconnectReplica*(rm: ReplicationManager, id: string): bool =
  acquire(rm.lock)
  result = false
  if id in rm.replicas:
    let replica = rm.replicas[id]
    if not replica.connected and replica.host.len > 0 and replica.port > 0:
      var sock = newSocket()
      if connectWithTimeout(sock, replica.host, Port(replica.port), 2000):
        replica.connected = true
        replica.state = rsStreaming
        replica.lastSeen = getMonoTime().ticks()
        result = true
      sock.close()
  release(rm.lock)

proc reconnectAll*(rm: ReplicationManager): int =
  acquire(rm.lock)
  result = 0
  for id, replica in rm.replicas:
    if not replica.connected and replica.host.len > 0 and replica.port > 0:
      var sock = newSocket()
      if connectWithTimeout(sock, replica.host, Port(replica.port), 2000):
        replica.connected = true
        replica.state = rsStreaming
        replica.lastSeen = getMonoTime().ticks()
        inc result
      sock.close()
  release(rm.lock)

proc startHealthCheck*(rm: ReplicationManager, intervalMs: int = 5000) {.async.} =
  while true:
    await sleepAsync(intervalMs)
    rm.healthCheck()

proc startReconnectionLoop*(rm: ReplicationManager, intervalMs: int = 10000) {.async.} =
  while true:
    await sleepAsync(intervalMs)
    discard rm.reconnectAll()