## Signature Management for Nexus Formats
##
## This module implements Ed25519 signing and verification for NPK, NIP, and NEXTER formats.
## It handles key generation, storage, and cryptographic operations.
##
## Key Storage Structure:
## ~/.local/share/nexus/keys/
##   ├── private/          # Private keys (0600 permissions)
##   │   └── <key_id>.key
##   ├── public/           # Public keys (0644 permissions)
##   │   └── <key_id>.pub
##   └── trusted/          # Trusted public keys for verification
##       └── <key_id>.pub
##

import std/[os, strutils, json, base64, tables, times, sets]
import ed25519
import ../nip/types

type
  SignatureManager* = object
    keysPath*: string
    privateKeysPath*: string
    publicKeysPath*: string
    trustedKeysPath*: string
    trustedKeys*: Table[string, PublicKey]

  KeyId* = string
  
  SignatureError* = object of NimPakError

  KeyPairInfo* = object
    id*: KeyId
    publicKey*: string  # Base64 encoded
    privateKey*: string # Base64 encoded (only when generating/exporting private)
    created*: DateTime

const
  KeyExtension = ".key"
  PubExtension = ".pub"

# Helper functions for conversion
proc toArray32(data: seq[byte]): array[32, byte] =
  if data.len != 32:
    raise newException(ValueError, "Invalid length for 32-byte array: " & $data.len)
  for i in 0..<32:
    result[i] = data[i]

proc toArray64(data: seq[byte]): array[64, byte] =
  if data.len != 64:
    raise newException(ValueError, "Invalid length for 64-byte array: " & $data.len)
  for i in 0..<64:
    result[i] = data[i]

proc toSeq(arr: array[32, byte]): seq[byte] =
  result = newSeq[byte](32)
  for i in 0..<32:
    result[i] = arr[i]

proc toSeq(arr: array[64, byte]): seq[byte] =
  result = newSeq[byte](64)
  for i in 0..<64:
    result[i] = arr[i]

proc encodeKey(key: array[32, byte]): string =
  base64.encode(key.toSeq)

proc encodeSig(sig: array[64, byte]): string =
  base64.encode(sig.toSeq)

proc decodeKey(data: string): array[32, byte] =
  let decodedStr = base64.decode(data)
  var bytes = newSeq[byte](decodedStr.len)
  for i in 0..<decodedStr.len:
    bytes[i] = decodedStr[i].byte
  result = toArray32(bytes)

proc decodeSig(data: string): array[64, byte] =
  let decodedStr = base64.decode(data)
  var bytes = newSeq[byte](decodedStr.len)
  for i in 0..<decodedStr.len:
    bytes[i] = decodedStr[i].byte
  result = toArray64(bytes)

proc decodePrivateKey(data: string): array[64, byte] =
  let decodedStr = base64.decode(data)
  var bytes = newSeq[byte](decodedStr.len)
  for i in 0..<decodedStr.len:
    bytes[i] = decodedStr[i].byte
  result = toArray64(bytes)

proc initSignatureManager*(rootPath: string): SignatureManager =
  ## Initialize signature manager with specified root path
  result.keysPath = rootPath / "keys"
  result.privateKeysPath = result.keysPath / "private"
  result.publicKeysPath = result.keysPath / "public"
  result.trustedKeysPath = result.keysPath / "trusted"
  result.trustedKeys = initTable[string, PublicKey]()

  # Create directories
  createDir(result.keysPath)
  createDir(result.privateKeysPath)
  createDir(result.publicKeysPath)
  createDir(result.trustedKeysPath)

  # Set secure permissions for private keys directory
  setFilePermissions(result.privateKeysPath, {fpUserRead, fpUserWrite, fpUserExec})

proc loadTrustedKeys*(manager: var SignatureManager) =
  ## Load all trusted public keys from disk
  for kind, path in walkDir(manager.trustedKeysPath):
    if kind == pcFile and path.endsWith(PubExtension):
      try:
        let keyId = path.extractFilename.changeFileExt("")
        let content = readFile(path).strip()
        let key = decodeKey(content)
        manager.trustedKeys[keyId] = key
      except Exception as e:
        echo "Warning: Failed to load trusted key ", path, ": ", e.msg

proc generateKeyPair*(manager: SignatureManager, keyId: string): KeyPairInfo =
  ## Generate a new Ed25519 key pair and save it
  let seed = ed25519.seed()
  let kp = ed25519.createKeypair(seed)
  
  let pubStr = encodeKey(kp.publicKey)
  let privStr = base64.encode(kp.privateKey.toSeq) # Private key is 64 bytes in this lib

  # Save private key
  let privPath = manager.privateKeysPath / keyId & KeyExtension
  writeFile(privPath, privStr)
  setFilePermissions(privPath, {fpUserRead, fpUserWrite})

  # Save public key
  let pubPath = manager.publicKeysPath / keyId & PubExtension
  writeFile(pubPath, pubStr)

  # Also save to trusted keys by default for own keys? 
  # Maybe not automatically, but useful for testing.
  # For now, we just return the info.

  result = KeyPairInfo(
    id: keyId,
    publicKey: pubStr,
    privateKey: privStr,
    created: now()
  )

proc sign*(manager: SignatureManager, data: string, keyId: string): string =
  ## Sign data with the specified private key
  let privPath = manager.privateKeysPath / keyId & KeyExtension
  if not fileExists(privPath):
    raise newException(SignatureError, "Private key not found: " & keyId)

  let privStr = readFile(privPath).strip()
  let privKey = decodePrivateKey(privStr)
  
  # We need the public key to reconstruct the KeyPair for the library
  # The library's KeyPair type is tuple[publicKey: PublicKey, privateKey: PrivateKey]
  # But wait, ed25519.sign takes KeyPair.
  # The private key file should ideally contain both or we need to derive public from private.
  # The `ed25519` lib `createKeypair` takes a seed.
  # But `PrivateKey` type is 64 bytes.
  # Usually Ed25519 private key is 32 bytes seed + 32 bytes public key = 64 bytes.
  # Let's check if we can extract public key from private key.
  # The library doesn't seem to have `getPublicKey(privateKey)`.
  # However, if we stored the 64-byte private key, the last 32 bytes MIGHT be the public key?
  # Standard Ed25519 private key is 32 bytes seed. 
  # The library defines PrivateKey as array[64, byte].
  # Let's assume we need to load the public key too.
  
  let pubPath = manager.publicKeysPath / keyId & PubExtension
  var pubKey: PublicKey
  if fileExists(pubPath):
    let pubStr = readFile(pubPath).strip()
    pubKey = decodeKey(pubStr)
  else:
    # Fallback: try to extract from private key if it follows standard layout
    # But safer to require public key or store it together.
    # For now, let's assume public key file exists.
    raise newException(SignatureError, "Public key not found for signing: " & keyId)

  let kp: KeyPair = (publicKey: pubKey, privateKey: privKey)
  let signature = ed25519.sign(data, kp)
  
  result = encodeSig(signature)

proc verify*(manager: SignatureManager, data: string, signature: string, keyId: string): bool =
  ## Verify signature using a trusted public key
  if not manager.trustedKeys.hasKey(keyId):
    # Try to load it if not in memory
    let trustedPath = manager.trustedKeysPath / keyId & PubExtension
    if fileExists(trustedPath):
      let content = readFile(trustedPath).strip()
      let key = decodeKey(content)
      # We can't modify manager here easily if it's not var, but we can return verification
      let sigBytes = decodeSig(signature)
      return ed25519.verify(data, sigBytes, key)
    else:
      return false
  
  let pubKey = manager.trustedKeys[keyId]
  let sigBytes = decodeSig(signature)
  return ed25519.verify(data, sigBytes, pubKey)

proc trustKey*(manager: SignatureManager, keyId: string) =
  ## Mark a public key as trusted (copy to trusted directory)
  let pubPath = manager.publicKeysPath / keyId & PubExtension
  let trustedPath = manager.trustedKeysPath / keyId & PubExtension
  
  if not fileExists(pubPath):
    raise newException(SignatureError, "Public key not found: " & keyId)
    
  copyFile(pubPath, trustedPath)

proc revokeKey*(manager: SignatureManager, keyId: string) =
  ## Revoke a trusted key (remove from trusted directory)
  let trustedPath = manager.trustedKeysPath / keyId & PubExtension
  if fileExists(trustedPath):
    removeFile(trustedPath)
