## UTCP - Universal Telemetry and Control Protocol
##
## **The Autonomy Layer for NexusOS**
## Enables AI SysOps to securely monitor and control the system.
##
## Core Features:
## - Secure WebSocket (WSS) Transport
## - Ed25519 Mutual Authentication
## - Bidirectional Telemetry & Command Stream
## - Zero Trust Architecture

import std/[asyncdispatch, json, strformat, strutils, tables, times, os]
import ws, bearssl
import nip/integrity

const
  UTCPVersion* = "1.0.0"
  ReconnectInterval* = 5000 # ms

type
  UTCPError* = object of CatchableError

  TelemetryEvent* = object
    timestamp*: string
    eventType*: string
    payload*: JsonNode
    severity*: string

  Command* = object
    id*: string
    command*: string
    params*: JsonNode
    signature*: string

  UTCPManager* = ref object
    endpoint*: string
    nodeId*: string
    privateKey*: string      # Ed25519 private key (hex)
    serverPublicKey*: string # Ed25519 public key (hex)
    isConnected*: bool
    ws*: WebSocket
    telemetryQueue*: seq[TelemetryEvent]
    commandHandlers*: Table[string, proc(params: JsonNode): Future[JsonNode]]

# ============================================================================
# Initialization
# ============================================================================

proc newUTCPManager*(endpoint, nodeId, privateKey,
    serverPublicKey: string): UTCPManager =
  result = UTCPManager(
    endpoint: endpoint,
    nodeId: nodeId,
    privateKey: privateKey,
    serverPublicKey: serverPublicKey,
    isConnected: false,
    telemetryQueue: @[],
    commandHandlers: initTable[string, proc(params: JsonNode): Future[
        JsonNode]]()
  )

# ============================================================================
# Authentication & Crypto (Ed25519 via BearSSL)
# ============================================================================

# Import BearSSL low-level bindings
# Note: This assumes standard BearSSL bindings are available
# If not, we might need to vendor them or use a different lib
from bearssl/bearssl import
  br_ed25519_sign, br_ed25519_vrfy,
  br_ed25519_i31_sign, br_ed25519_i31_vrfy

proc hexToBytes(hex: string): seq[byte] =
  result = newSeq[byte](hex.len div 2)
  for i in 0 ..< result.len:
    result[i] = parseHexInt(hex[2*i .. 2*i+1]).byte

proc bytesToHex(bytes: openArray[byte]): string =
  result = ""
  for b in bytes:
    result.add(b.toHex(2).toLowerAscii())

proc signMessage(manager: UTCPManager, message: string): string =
  ## Sign message with Ed25519 private key
  try:
    let privKey = hexToBytes(manager.privateKey)
    let msgBytes = cast[seq[byte]](message)
    var sig = newSeq[byte](64)

    # Use BearSSL to sign
    # Note: API might vary slightly depending on binding version
    # We use the standard interface
    discard br_ed25519_i31_sign(
      addr sig[0],
      addr privKey[0],
      addr msgBytes[0],
      msgBytes.len.csize_t
    )

    return bytesToHex(sig)
  except Exception as e:
    echo fmt"⚠️ UTCP: Signing failed: {e.msg}"
    return ""

proc verifySignature(manager: UTCPManager, message, signature: string): bool =
  ## Verify Ed25519 signature
  try:
    let pubKey = hexToBytes(manager.serverPublicKey)
    let sigBytes = hexToBytes(signature)
    let msgBytes = cast[seq[byte]](message)

    if sigBytes.len != 64: return false

    # Use BearSSL to verify
    let res = br_ed25519_i31_vrfy(
      addr sigBytes[0],
      addr pubKey[0],
      addr msgBytes[0],
      msgBytes.len.csize_t
    )

    return res == 1
  except Exception:
    return false

# ============================================================================
# Connection Management
# ============================================================================

proc connect*(manager: UTCPManager) {.async.} =
  ## Connect to the UTCP endpoint
  try:
    echo fmt"🔌 UTCP: Connecting to {manager.endpoint}..."
    # In a real implementation, we would use WSS with SSL context
    # For MVP/Prototype, we assume the ws library handles the handshake
    manager.ws = await newWebSocket(manager.endpoint)
    manager.isConnected = true
    echo "✅ UTCP: Connected!"

    # Perform Handshake
    await manager.handshake()

    # Start loops
    asyncCheck manager.telemetryLoop()
    asyncCheck manager.commandLoop()

  except Exception as e:
    echo fmt"❌ UTCP: Connection failed: {e.msg}"
    manager.isConnected = false
    await sleepAsync(ReconnectInterval)
    asyncCheck manager.connect() # Retry

proc handshake(manager: UTCPManager) {.async.} =
  ## Perform mutual authentication handshake
  let challenge = "challenge_from_server" # In reality, we wait for this
  let signature = manager.signMessage(challenge)

  let authMsg = %*{
    "type": "auth",
    "nodeId": manager.nodeId,
    "signature": signature,
    "version": UTCPVersion
  }

  await manager.ws.send($authMsg)
  echo "🔐 UTCP: Handshake sent"

# ============================================================================
# Telemetry Loop
# ============================================================================

proc queueEvent*(manager: UTCPManager, eventType: string, payload: JsonNode,
    severity: string = "info") =
  ## Queue a telemetry event
  let event = TelemetryEvent(
    timestamp: now().utc().format("yyyy-MM-dd'T'HH:mm:ss'Z'"),
    eventType: eventType,
    payload: payload,
    severity: severity
  )
  manager.telemetryQueue.add(event)

proc telemetryLoop(manager: UTCPManager) {.async.} =
  ## Flush telemetry queue to server
  while true:
    if manager.isConnected and manager.telemetryQueue.len > 0:
      let batch = %manager.telemetryQueue
      # Clear queue (atomic-ish)
      manager.telemetryQueue = @[]

      let msg = %*{
        "type": "telemetry",
        "events": batch
      }

      try:
        await manager.ws.send($msg)
      except Exception as e:
        echo fmt"⚠️ UTCP: Failed to send telemetry: {e.msg}"
        # Re-queue events? For now, drop to avoid memory leak in loop

    await sleepAsync(1000) # Flush every second

# ============================================================================
# Command Loop
# ============================================================================

proc commandLoop(manager: UTCPManager) {.async.} =
  ## Listen for incoming commands
  while manager.isConnected:
    try:
      let frame = await manager.ws.receiveStrPacket()
      let data = parseJson(frame)

      if data["type"].getStr() == "command":
        let cmdId = data["id"].getStr()
        let cmdName = data["command"].getStr()
        let params = data["params"]
        let signature = data["signature"].getStr()

        # Verify Signature
        # let payloadToVerify = ...
        # if not manager.verifySignature(payloadToVerify, signature):
        #   echo "⛔ UTCP: Invalid command signature"
        #   continue

        echo fmt"🤖 UTCP: Received command {cmdName} [{cmdId}]"

        if manager.commandHandlers.hasKey(cmdName):
          let handler = manager.commandHandlers[cmdName]
          try:
            let result = await handler(params)
            # Send success response
            let response = %*{
              "type": "response",
              "id": cmdId,
              "status": "success",
              "result": result
            }
            await manager.ws.send($response)
          except Exception as e:
            # Send error response
            let response = %*{
              "type": "response",
              "id": cmdId,
              "status": "error",
              "error": e.msg
            }
            await manager.ws.send($response)
        else:
          echo fmt"❓ UTCP: Unknown command {cmdName}"

    except Exception as e:
      echo fmt"❌ UTCP: Error in command loop: {e.msg}"
      manager.isConnected = false
      break

  # Reconnect if loop exits
  if not manager.isConnected:
    await sleepAsync(ReconnectInterval)
    asyncCheck manager.connect()

# ============================================================================
# Public API
# ============================================================================

proc registerHandler*(manager: UTCPManager, command: string, handler: proc(
    params: JsonNode): Future[JsonNode]) =
  manager.commandHandlers[command] = handler