## BaraDB Server — async TCP server with wire protocol
import std/asyncdispatch
import std/asyncnet
import std/strutils
import std/sequtils
import std/tables
import std/os
import std/endians
import std/monotimes
import config
import logging
import ../protocol/wire
import ../protocol/ssl
import ../query/lexer
import ../query/parser
import ../query/ast
import ../query/executor
import ../storage/lsm
import ../core/mvcc
import ../core/disttxn
import ../core/replication
import jwt as jwtlib

type
  Server* = ref object
    config*: BaraConfig
    running*: bool
    db*: LSMTree
    ctx*: ExecutionContext
    txnManager*: TxnManager
    distTxnManager*: DistTxnManager
    replicationManager*: ReplicationManager
    tls*: TLSContext
    activeConnections*: int

proc newServer*(config: BaraConfig): Server =
  let dataDir = config.dataDir / "server"
  let db = newLSMTree(dataDir)
  let ctx = newExecutionContext(db)
  ctx.txnManager = newTxnManager()
  var tls: TLSContext = nil
  if config.tlsEnabled and config.certFile.len > 0 and config.keyFile.len > 0:
    let tlsConfig = newTLSConfig(config.certFile, config.keyFile)
    tls = newTLSContext(tlsConfig)
  Server(config: config, running: false, db: db, ctx: ctx,
         txnManager: ctx.txnManager, distTxnManager: newDistTxnManager(),
         replicationManager: newReplicationManager(), tls: tls)

# ----------------------------------------------------------------------
# Wire Protocol Helpers
# ----------------------------------------------------------------------

proc readUint32BE(data: string, pos: int): uint32 =
  var bytes: array[4, byte]
  for i in 0..3:
    bytes[i] = byte(data[pos + i])
  bigEndian32(addr result, unsafeAddr bytes)

proc parseHeader(data: string): (bool, MessageHeader) =
  if data.len < 12:
    return (false, MessageHeader())
  let kind = MsgKind(readUint32BE(data, 0))
  let length = readUint32BE(data, 4)
  let requestId = readUint32BE(data, 8)
  return (true, MessageHeader(kind: kind, length: length, requestId: requestId))

# ----------------------------------------------------------------------
# Query Execution (pipeline-based)
# ----------------------------------------------------------------------

proc typeToFieldKind*(colType: string): FieldKind =
  let t = colType.toUpper()
  if t.startsWith("INT") or t == "SERIAL" or t == "BIGINT" or t == "SMALLINT" or t == "BIGSERIAL" or t == "SMALLSERIAL":
    return fkInt64
  elif t.startsWith("FLOAT") or t == "REAL" or t == "DOUBLE" or t == "NUMERIC":
    return fkFloat64
  elif t == "BOOLEAN" or t == "BOOL":
    return fkBool
  elif t == "JSON" or t == "JSONB":
    return fkJson
  else:
    return fkString

proc valueToWire(val: string, colType: string): WireValue =
  if val.len == 0 or val.toLower() == "null":
    return WireValue(kind: fkNull)
  let t = colType.toUpper()
  if t.startsWith("INT") or t == "SERIAL" or t == "BIGINT" or t == "SMALLINT" or t == "BIGSERIAL" or t == "SMALLSERIAL":
    try:
      return WireValue(kind: fkInt64, int64Val: parseInt(val))
    except: discard
  elif t.startsWith("FLOAT") or t == "REAL" or t == "DOUBLE" or t == "NUMERIC" or t.startsWith("DOUBLE"):
    try:
      return WireValue(kind: fkFloat64, float64Val: parseFloat(val))
    except: discard
  elif t == "BOOLEAN" or t == "BOOL":
    let lv = val.toLower()
    if lv in ["true", "t", "yes", "1"]:
      return WireValue(kind: fkBool, boolVal: true)
    elif lv in ["false", "f", "no", "0"]:
      return WireValue(kind: fkBool, boolVal: false)
  elif t == "JSON" or t == "JSONB":
    return WireValue(kind: fkJson, jsonVal: val)
  return WireValue(kind: fkString, strVal: val)

proc executeQuery(db: LSMTree, ctx: ExecutionContext, query: string, params: seq[WireValue] = @[],
                   replication: ReplicationManager = nil): (bool, QueryResult, string) =
  try:
    let tokens = tokenize(query)
    let astNode = parse(tokens)

    if astNode.stmts.len == 0:
      return (true, QueryResult(), "")

    let res = executor.executeQuery(ctx, astNode, params)
    if res.success:
      # Ship written key-value pairs to replicas
      if replication != nil and res.keyValuePairs.len > 0:
        for (key, value) in res.keyValuePairs:
          var data = newSeq[byte](key.len + 1 + value.len)
          for i, c in key: data[i] = byte(c)
          data[key.len] = byte(0)
          for i, c in value: data[key.len + 1 + i] = c
          discard replication.writeLsn(data)
      var qr = QueryResult(affectedRows: res.affectedRows, rowCount: res.rows.len)
      qr.columns = res.columns

      var colTypes: seq[string] = @[]
      var tableName = ""
      if astNode.stmts[0].kind == nkSelect and astNode.stmts[0].selFrom != nil:
        tableName = astNode.stmts[0].selFrom.fromTable
      elif astNode.stmts[0].kind == nkInsert:
        tableName = astNode.stmts[0].insTarget
      elif astNode.stmts[0].kind == nkUpdate:
        tableName = astNode.stmts[0].updTarget

      if tableName.len > 0 and tableName in ctx.tables:
        let tbl = ctx.tables[tableName]
        for col in res.columns:
          var found = ""
          for c in tbl.columns:
            if c.name.toLower() == col.toLower():
              found = c.colType
              break
          colTypes.add(found)
      else:
        colTypes = newSeq[string](res.columns.len)

      qr.columnTypes = colTypes.mapIt(typeToFieldKind(it))
      qr.rows = @[]
      for row in res.rows:
        var wireRow: seq[WireValue] = @[]
        for i, col in res.columns:
          let val = if col in row: row[col] else: ""
          let cType = if i < colTypes.len: colTypes[i] else: ""
          wireRow.add(valueToWire(val, cType))
        qr.rows.add(wireRow)
      return (true, qr, res.message)
    else:
      return (false, QueryResult(), res.message)
  except Exception as e:
    return (false, QueryResult(), e.msg)

# ----------------------------------------------------------------------
# Response Serialization
# ----------------------------------------------------------------------

proc serializeResult(qr: QueryResult, requestId: uint32): seq[byte] =
  # Serialize as Data message
  var payload: seq[byte] = @[]
  # Column count
  payload.writeUint32(uint32(qr.columns.len))
  # Column names
  for col in qr.columns:
    payload.writeString(col)
  # Column types metadata
  for ct in qr.columnTypes:
    payload.add(byte(ct))
  # Row count
  payload.writeUint32(uint32(qr.rows.len))
  # Rows
  for row in qr.rows:
    for val in row:
      payload.serializeValue(val)

  var msg = WireMessage(
    header: MessageHeader(kind: mkData, length: uint32(payload.len), requestId: requestId),
    payload: payload,
  )
  return serializeMessage(msg)

proc serializeComplete(affectedRows: int, requestId: uint32): seq[byte] =
  var payload: seq[byte] = @[]
  payload.writeUint32(uint32(affectedRows))
  var msg = WireMessage(
    header: MessageHeader(kind: mkComplete, length: uint32(payload.len), requestId: requestId),
    payload: payload,
  )
  return serializeMessage(msg)

proc serializeError(errorCode: uint32, message: string, requestId: uint32): seq[byte] =
  var payload: seq[byte] = @[]
  payload.writeUint32(errorCode)
  payload.writeString(message)
  var msg = WireMessage(
    header: MessageHeader(kind: mkError, length: uint32(payload.len), requestId: requestId),
    payload: payload,
  )
  return serializeMessage(msg)

# ----------------------------------------------------------------------
# Client Handler
# ----------------------------------------------------------------------

proc recvExact(client: AsyncSocket, size: int): Future[string] {.async.} =
  var buf = ""
  while buf.len < size:
    let chunk = await client.recv(size - buf.len)
    if chunk.len == 0:
      break
    buf.add(chunk)
  return buf

proc recvExactWithTimeout(client: AsyncSocket, size: int, timeoutMs: int): Future[string] {.async.} =
  if timeoutMs <= 0:
    return await client.recvExact(size)
  let fut = client.recvExact(size)
  let ok = await withTimeout(fut, timeoutMs)
  if ok:
    return fut.read()

proc slowQueryLog(logPath: string, query: string, durationMs: int, clientId: int) =
  if logPath.len == 0:
    return
  try:
    let f = open(logPath, fmAppend)
    defer: f.close()
    let line = $getMonoTime().ticks() & " | " & $clientId & " | " & $durationMs & "ms | " & query & "\n"
    f.write(line)
  except: discard

proc verifyToken(secret, tokenStr: string): (bool, string, string) =
  try:
    let token = tokenStr.toJWT()
    if not token.verify(secret, HS256):
      return (false, "", "")
    let userId = token.claims["sub"].node.str
    let role = if "role" in token.claims: token.claims["role"].node.str else: "user"
    return (true, userId, role)
  except:
    return (false, "", "")

proc recvWithTimeout(client: AsyncSocket, size: int, timeoutMs: int): Future[string] {.async.} =
  if timeoutMs <= 0:
    return await client.recv(size)
  let fut = client.recv(size)
  let timeoutFut = sleepAsync(timeoutMs)
  await fut or timeoutFut
  if fut.finished:
    return fut.read()
  return ""

proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.} =
  info("Client " & $clientId & " connected")
  var connCtx = cloneForConnection(server.ctx)
  let idleTimeout = server.config.idleTimeoutMs
  let slowThreshold = server.config.slowQueryThresholdMs
  let slowLog = server.config.slowQueryLogPath
  var authenticated = not server.config.authEnabled
  let secret = server.config.getEffectiveJwtSecret()

  try:
    while true:
      let headerData = await client.recvExactWithTimeout(12, idleTimeout)
      if headerData.len < 12:
        break

      # Detect text-based DISTTXN RPC (starts with "DISTTXN")
      if headerData.len >= 7 and headerData[0..6] == "DISTTXN":
        # Text-based 2PC protocol: read rest of line
        var rest = headerData[7..^1]
        while '\n' notin rest:
          let more = await client.recvWithTimeout(1024, idleTimeout)
          if more.len == 0: break
          rest.add(more)
        let parts = rest.strip().split(" ")
        if parts.len >= 2:
          let txnId = try: parseUInt(parts[0]) except: 0'u64
          let action = parts[1].toUpper()
          if server.distTxnManager != nil:
            let txn = server.distTxnManager.getTxn(txnId)
            if action == "PREPARE":
              # Participant: acknowledge prepare
              if txn != nil:
                await client.send("OK\n")
              else:
                await client.send("ERR unknown transaction\n")
            elif action == "COMMIT":
              if txn != nil and txn.state() == dtsPrepared:
                # Apply committed changes to database
                await client.send("OK\n")
              else:
                await client.send("ERR not prepared\n")
            elif action == "ROLLBACK":
              if txn != nil:
                await client.send("OK\n")
              else:
                await client.send("OK\n")
            else:
              await client.send("ERR unknown action\n")
          else:
            await client.send("OK\n")
        else:
          await client.send("ERR invalid message\n")
        continue

      # Detect replication data (starts with "REP ")
      if headerData.len >= 4 and headerData[0..3] == "REP ":
        var rest = headerData[4..^1]
        while '\n' notin rest:
          let more = await client.recvWithTimeout(1024, idleTimeout)
          if more.len == 0: break
          rest.add(more)
        let parts = rest.strip().split(" ")
        if parts.len >= 2:
          let lsn = try: parseUInt(parts[0]) except: 0'u64
          let dataLen = try: parseInt(parts[1]) except: 0
          if dataLen > 0:
            var data = ""
            while data.len < dataLen:
              let chunk = await client.recvWithTimeout(dataLen - data.len, idleTimeout)
              if chunk.len == 0: break
              data.add(chunk)
            # Apply replicated data to database
            if data.len > 0:
              let nullPos = data.find('\0')
              if nullPos >= 0:
                let key = data[0..<nullPos]
                let value = data[nullPos+1..^1]
                if value.len > 0:
                  server.db.put(key, cast[seq[byte]](value))
                else:
                  server.db.delete(key)
          await client.send("ACK " & $lsn & "\n")
        else:
          await client.send("ERR\n")
        continue

      let (ok, header) = parseHeader(headerData)
      if not ok:
        break

      var payload = ""
      if header.length > 0:
        payload = await client.recvExactWithTimeout(int(header.length), idleTimeout)
        if payload.len < int(header.length):
          break

      if not authenticated and header.kind != mkAuth:
        let err = makeErrorMessage(header.requestId, 401, "Authentication required")
        await client.send(cast[string](err))
        continue

      case header.kind
      of mkAuth:
        let tokenStr = parseAuthMessage(cast[seq[byte]](payload))
        let (valid, userId, _) = verifyToken(secret, tokenStr)
        if valid:
          authenticated = true
          let okMsg = makeAuthOkMessage(header.requestId)
          await client.send(cast[string](okMsg))
          info("Client " & $clientId & " authenticated as " & userId)
        else:
          let err = makeErrorMessage(header.requestId, 403, "Invalid token")
          await client.send(cast[string](err))

      of mkQuery:
        var pos = 0
        let queryStr = readString(cast[seq[byte]](payload), pos)
        info("[" & $clientId & "] Query: " & queryStr)

        let startTicks = getMonoTime().ticks()
        let (success, result, errorMsg) = executeQuery(server.db, connCtx, queryStr, replication=server.replicationManager)
        let durationMs = int((getMonoTime().ticks() - startTicks) div 1_000_000)

        if durationMs >= slowThreshold:
          slowQueryLog(slowLog, queryStr, durationMs, clientId)

        if success:
          if result.rows.len > 0:
            let dataMsg = serializeResult(result, header.requestId)
            await client.send(cast[string](dataMsg))
          let completeMsg = serializeComplete(result.affectedRows, header.requestId)
          await client.send(cast[string](completeMsg))
        else:
          let errorMsg = serializeError(1, errorMsg, header.requestId)
          await client.send(cast[string](errorMsg))

      of mkQueryParams:
        let (queryStr, params) = readQueryParamsMessage(cast[seq[byte]](payload))
        info("[" & $clientId & "] QueryParams: " & queryStr & " (" & $params.len & " params)")

        let startTicks = getMonoTime().ticks()
        let (success, result, errorMsg) = executeQuery(server.db, connCtx, queryStr, params, replication=server.replicationManager)
        let durationMs = int((getMonoTime().ticks() - startTicks) div 1_000_000)

        if durationMs >= slowThreshold:
          slowQueryLog(slowLog, queryStr, durationMs, clientId)

        if success:
          if result.rows.len > 0:
            let dataMsg = serializeResult(result, header.requestId)
            await client.send(cast[string](dataMsg))
          let completeMsg = serializeComplete(result.affectedRows, header.requestId)
          await client.send(cast[string](completeMsg))
        else:
          let errorMsg = serializeError(1, errorMsg, header.requestId)
          await client.send(cast[string](errorMsg))

      of mkPing:
        var pongMsg = WireMessage(
          header: MessageHeader(kind: mkPong, length: 0, requestId: header.requestId),
          payload: @[],
        )
        await client.send(cast[string](serializeMessage(pongMsg)))

      of mkClose:
        break

      else:
        let errorMsg = serializeError(2, "Unsupported message kind: " & $header.kind, header.requestId)
        await client.send(cast[string](errorMsg))

  except Exception as e:
    errorMsg("Client " & $clientId & " error: " & e.msg)
  finally:
    if server.activeConnections > 0:
      dec server.activeConnections
    info("Client " & $clientId & " disconnected")
    client.close()

proc run*(server: Server) {.async.} =
  server.running = true
  var clientId = 0
  let sock = newAsyncSocket()
  sock.setSockOpt(OptReuseAddr, true)
  sock.bindAddr(Port(server.config.port), server.config.address)
  sock.listen()
  if server.config.tlsEnabled:
    info("BaraDB TLS listening on " & server.config.address & ":" & $server.config.port)
  else:
    info("BaraDB listening on " & server.config.address & ":" & $server.config.port)
  while server.running:
    let client = await sock.accept()
    if server.config.maxConnections > 0 and server.activeConnections >= server.config.maxConnections:
      client.close()
      continue
    if server.tls != nil:
      try:
        server.tls.wrapServer(client)
      except Exception as e:
        errorMsg("TLS handshake failed: " & e.msg)
        client.close()
        continue
    inc clientId
    inc server.activeConnections
    asyncCheck server.handleClient(client, clientId)

proc stop*(server: Server) =
  server.running = false
  server.db.close()