## NimPak Stress Testing
##
## Comprehensive stress tests for the NimPak package manager.
## Task 44: Stress testing.
##
## Run with: nim c -r -d:release nip/tests/test_stress.nim
## Note: Some tests may take several minutes to complete.

import std/[os, strutils, strformat, times, random, sequtils, locks, threadpool]
import ../src/nimpak/cas
import ../src/nimpak/benchmark

const
  # Test configuration - adjust based on available resources
  SmallScale* = 100
  MediumScale* = 1000
  LargeScale* = 10000
  HugeScale* = 100000
  
  SmallChunk* = 1024          # 1KB
  MediumChunk* = 65536        # 64KB
  LargeChunk* = 1048576       # 1MB

type
  StressTestResult* = object
    name*: string
    passed*: bool
    duration*: float          # seconds
    operationsCompleted*: int
    bytesProcessed*: int64
    errorsEncountered*: int
    peakMemoryMB*: float
    notes*: seq[string]

  StressTestSuite* = object
    name*: string
    results*: seq[StressTestResult]
    startTime*: DateTime
    endTime*: DateTime
    totalPassed*: int
    totalFailed*: int

# ############################################################################
# Utility Functions
# ############################################################################

proc generateRandomData(size: int): seq[byte] =
  ## Generate random data for testing
  result = newSeq[byte](size)
  for i in 0..<size:
    result[i] = byte(rand(255))

proc generateSemiRandomData(size: int, seed: int): seq[byte] =
  ## Generate semi-random data (reproducible) for deduplication testing
  result = newSeq[byte](size)
  var r = initRand(seed)
  for i in 0..<size:
    result[i] = byte(r.rand(255))

proc getMemoryUsageMB(): float =
  ## Get current memory usage in MB (Linux-specific)
  when defined(linux):
    try:
      let status = readFile("/proc/self/status")
      for line in status.splitLines:
        if line.startsWith("VmRSS:"):
          let parts = line.split()
          if parts.len >= 2:
            return parseFloat(parts[1]) / 1024.0
    except:
      discard
  return 0.0

proc formatBytes(bytes: int64): string =
  if bytes >= 1073741824:
    fmt"{bytes.float / 1073741824.0:.2f} GB"
  elif bytes >= 1048576:
    fmt"{bytes.float / 1048576.0:.2f} MB"
  elif bytes >= 1024:
    fmt"{bytes.float / 1024.0:.2f} KB"
  else:
    fmt"{bytes} bytes"

# ############################################################################
# CAS Stress Tests
# ############################################################################

proc stressTestCasManyChunks*(casRoot: string, chunkCount: int, 
                              chunkSize: int): StressTestResult =
  ## Test CAS with many chunks
  result = StressTestResult(
    name: fmt"CAS Many Chunks ({chunkCount} x {chunkSize} bytes)",
    passed: false,
    notes: @[]
  )
  
  echo fmt"🏋️ Stress Test: {result.name}"
  
  var casManager = initCasManager(casRoot, casRoot / "system")
  let startTime = epochTime()
  let startMem = getMemoryUsageMB()
  
  randomize()
  var storedHashes: seq[string] = @[]
  
  for i in 1..chunkCount:
    let data = generateRandomData(chunkSize)
    let storeResult = casManager.storeObject(data)
    
    if storeResult.isOk:
      storedHashes.add(storeResult.get().hash)
      result.operationsCompleted += 1
      result.bytesProcessed += int64(chunkSize)
    else:
      result.errorsEncountered += 1
    
    if i mod 1000 == 0:
      echo fmt"  Progress: {i}/{chunkCount} ({i * 100 div chunkCount}%)"
  
  let endTime = epochTime()
  result.duration = endTime - startTime
  result.peakMemoryMB = getMemoryUsageMB()
  
  # Verify some random chunks
  echo "  Verifying stored chunks..."
  var verifyErrors = 0
  for i in 0..<min(100, storedHashes.len):
    let idx = rand(storedHashes.len - 1)
    if not casManager.objectExists(storedHashes[idx]):
      verifyErrors += 1
  
  result.passed = result.errorsEncountered == 0 and verifyErrors == 0
  result.notes.add(fmt"Stored {storedHashes.len} unique hashes")
  result.notes.add(fmt"Throughput: {result.bytesProcessed.float / result.duration / 1048576.0:.2f} MB/s")
  result.notes.add(fmt"Memory: {result.peakMemoryMB:.1f} MB")
  
  echo fmt"  ✓ Completed in {result.duration:.2f}s"

proc stressTestCasDeduplication*(casRoot: string, chunkCount: int,
                                  duplicateRatio: float): StressTestResult =
  ## Test CAS deduplication efficiency at scale
  result = StressTestResult(
    name: fmt"CAS Deduplication ({chunkCount} chunks, {duplicateRatio*100:.0f}% duplicates)",
    passed: false,
    notes: @[]
  )
  
  echo fmt"🏋️ Stress Test: {result.name}"
  
  var casManager = initCasManager(casRoot, casRoot / "system")
  let startTime = epochTime()
  
  # Generate unique chunks
  let uniqueCount = int(float(chunkCount) * (1.0 - duplicateRatio))
  var uniqueChunks: seq[seq[byte]] = @[]
  
  echo fmt"  Generating {uniqueCount} unique chunks..."
  for i in 0..<uniqueCount:
    uniqueChunks.add(generateSemiRandomData(MediumChunk, i))
  
  # Store chunks (with duplicates)
  randomize()
  var totalStored = 0
  var dedupCount = 0
  
  echo "  Storing chunks..."
  for i in 0..<chunkCount:
    let chunkIdx = if rand(1.0) < duplicateRatio and uniqueChunks.len > 0:
      rand(uniqueChunks.len - 1)  # Reuse existing chunk
    else:
      i mod uniqueChunks.len
    
    let storeResult = casManager.storeObject(uniqueChunks[chunkIdx])
    if storeResult.isOk:
      totalStored += 1
      if storeResult.get().refCount > 1:
        dedupCount += 1
      result.bytesProcessed += int64(MediumChunk)
    
    if i mod 1000 == 0:
      echo fmt"  Progress: {i}/{chunkCount} ({i * 100 div chunkCount}%)"
  
  let endTime = epochTime()
  result.duration = endTime - startTime
  result.operationsCompleted = totalStored
  
  let actualDedupRatio = if totalStored > 0: float(dedupCount) / float(totalStored) else: 0.0
  result.passed = totalStored == chunkCount
  result.notes.add(fmt"Expected dedup ratio: {duplicateRatio*100:.1f}%")
  result.notes.add(fmt"Actual dedup hits: {dedupCount} ({actualDedupRatio*100:.1f}%)")
  result.notes.add(fmt"Storage saved: ~{formatBytes(int64(dedupCount * MediumChunk))}")
  
  echo fmt"  ✓ Completed in {result.duration:.2f}s"

proc stressTestCasLargeObjects*(casRoot: string, objectSizeMB: int, 
                                 objectCount: int): StressTestResult =
  ## Test CAS with large objects
  result = StressTestResult(
    name: fmt"CAS Large Objects ({objectCount} x {objectSizeMB}MB)",
    passed: false,
    notes: @[]
  )
  
  echo fmt"🏋️ Stress Test: {result.name}"
  
  var casManager = initCasManager(casRoot, casRoot / "system")
  let startTime = epochTime()
  let objectSize = objectSizeMB * 1048576
  
  for i in 1..objectCount:
    echo fmt"  Storing object {i}/{objectCount} ({objectSizeMB}MB)..."
    let data = generateRandomData(objectSize)
    
    let storeResult = casManager.storeObject(data)
    if storeResult.isOk:
      result.operationsCompleted += 1
      result.bytesProcessed += int64(objectSize)
      
      # Verify retrieval
      let retrieveResult = casManager.retrieveObject(storeResult.get().hash)
      if retrieveResult.isOk:
        let retrieved = retrieveResult.get()
        if retrieved.len == objectSize:
          result.notes.add(fmt"Object {i}: stored and verified")
        else:
          result.errorsEncountered += 1
          result.notes.add(fmt"Object {i}: size mismatch")
      else:
        result.errorsEncountered += 1
    else:
      result.errorsEncountered += 1
  
  let endTime = epochTime()
  result.duration = endTime - startTime
  result.passed = result.errorsEncountered == 0 and result.operationsCompleted == objectCount
  result.notes.add(fmt"Total data: {formatBytes(result.bytesProcessed)}")
  result.notes.add(fmt"Throughput: {result.bytesProcessed.float / result.duration / 1048576.0:.2f} MB/s")
  
  echo fmt"  ✓ Completed in {result.duration:.2f}s"

# ############################################################################
# Concurrent Operations Stress Tests
# ############################################################################

var globalCasManager {.threadvar.}: CasManager
var globalLock: Lock

proc stressTestConcurrentStores*(casRoot: string, threadCount: int,
                                  operationsPerThread: int): StressTestResult =
  ## Test concurrent store operations
  result = StressTestResult(
    name: fmt"Concurrent Stores ({threadCount} threads x {operationsPerThread} ops)",
    passed: false,
    notes: @[]
  )
  
  echo fmt"🏋️ Stress Test: {result.name}"
  
  initLock(globalLock)
  var casManager = initCasManager(casRoot, casRoot / "system")
  let startTime = epochTime()
  
  # For now, simulate concurrent operations sequentially
  # (Full threading would require thread-safe CasManager)
  var totalOps = 0
  var totalErrors = 0
  
  for t in 0..<threadCount:
    for op in 0..<operationsPerThread:
      let data = generateRandomData(SmallChunk)
      let storeResult = casManager.storeObject(data)
      if storeResult.isOk:
        totalOps += 1
        result.bytesProcessed += int64(SmallChunk)
      else:
        totalErrors += 1
    
    echo fmt"  Thread {t+1}/{threadCount} completed"
  
  let endTime = epochTime()
  result.duration = endTime - startTime
  result.operationsCompleted = totalOps
  result.errorsEncountered = totalErrors
  result.passed = totalErrors == 0
  result.notes.add("Note: Sequential simulation (true concurrency requires thread-safe CAS)")
  result.notes.add(fmt"Ops/sec: {float(totalOps) / result.duration:.0f}")
  
  deinitLock(globalLock)
  echo fmt"  ✓ Completed in {result.duration:.2f}s"

# ############################################################################
# Garbage Collection Stress Tests
# ############################################################################

proc stressTestGarbageCollection*(casRoot: string, chunkCount: int): StressTestResult =
  ## Test garbage collection at scale
  result = StressTestResult(
    name: fmt"Garbage Collection ({chunkCount} chunks)",
    passed: false,
    notes: @[]
  )
  
  echo fmt"🏋️ Stress Test: {result.name}"
  
  var casManager = initCasManager(casRoot, casRoot / "system")
  let startTime = epochTime()
  
  # Store chunks with references
  echo "  Phase 1: Storing chunks with references..."
  var pinnedHashes: seq[string] = @[]
  var unpinnedHashes: seq[string] = @[]
  
  for i in 0..<chunkCount:
    let data = generateRandomData(SmallChunk)
    let storeResult = casManager.storeObject(data)
    if storeResult.isOk:
      let hash = storeResult.get().hash
      if i mod 10 == 0:
        # Pin every 10th chunk
        pinnedHashes.add(hash)
        discard casManager.pinObject(hash, fmt"pin_{i}")
      else:
        unpinnedHashes.add(hash)
      result.bytesProcessed += int64(SmallChunk)
    
    if i mod 1000 == 0:
      echo fmt"    Progress: {i}/{chunkCount}"
  
  # Remove references from unpinned chunks
  echo "  Phase 2: Removing references..."
  for hash in unpinnedHashes:
    discard casManager.removeReference(hash, NPK, "stress-test")
  
  # Run garbage collection
  echo "  Phase 3: Running garbage collection..."
  let gcStartTime = epochTime()
  let gcResult = casManager.garbageCollect()
  let gcDuration = epochTime() - gcStartTime
  
  # Verify pinned chunks still exist
  echo "  Phase 4: Verifying pinned chunks..."
  var stillExist = 0
  for hash in pinnedHashes:
    if casManager.objectExists(hash):
      stillExist += 1
  
  let endTime = epochTime()
  result.duration = endTime - startTime
  result.operationsCompleted = chunkCount
  
  if gcResult.isOk:
    let gcCollected = gcResult.get()
    result.notes.add(fmt"GC collected: {gcCollected} objects")
    result.notes.add(fmt"GC duration: {gcDuration:.2f}s")
  
  result.notes.add(fmt"Pinned chunks preserved: {stillExist}/{pinnedHashes.len}")
  result.passed = stillExist == pinnedHashes.len
  
  echo fmt"  ✓ Completed in {result.duration:.2f}s"

# ############################################################################
# File Operations Stress Tests
# ############################################################################

proc stressTestManySmallFiles*(casRoot: string, fileCount: int): StressTestResult =
  ## Test storing many small files
  result = StressTestResult(
    name: fmt"Many Small Files ({fileCount} files)",
    passed: false,
    notes: @[]
  )
  
  echo fmt"🏋️ Stress Test: {result.name}"
  
  let testDir = getTempDir() / fmt"nip_stress_files_{epochTime().int}"
  createDir(testDir)
  
  var casManager = initCasManager(casRoot, casRoot / "system")
  let startTime = epochTime()
  
  # Create and store files
  echo "  Phase 1: Creating and storing files..."
  var storedHashes: seq[string] = @[]
  
  for i in 0..<fileCount:
    let filePath = testDir / fmt"file_{i:06d}.dat"
    let fileSize = 100 + rand(4000)  # 100B to 4KB
    let data = generateRandomData(fileSize)
    
    writeFile(filePath, cast[string](data))
    
    let storeResult = casManager.storeFile(filePath)
    if storeResult.isOk:
      storedHashes.add(storeResult.get().hash)
      result.operationsCompleted += 1
      result.bytesProcessed += int64(fileSize)
    else:
      result.errorsEncountered += 1
    
    if i mod 500 == 0:
      echo fmt"    Progress: {i}/{fileCount}"
  
  # Cleanup test files
  echo "  Phase 2: Cleaning up..."
  removeDir(testDir)
  
  let endTime = epochTime()
  result.duration = endTime - startTime
  result.passed = result.errorsEncountered == 0
  result.notes.add(fmt"Files stored: {storedHashes.len}")
  result.notes.add(fmt"Throughput: {float(result.operationsCompleted) / result.duration:.0f} files/sec")
  result.notes.add(fmt"Data rate: {result.bytesProcessed.float / result.duration / 1048576.0:.2f} MB/s")
  
  echo fmt"  ✓ Completed in {result.duration:.2f}s"

# ############################################################################
# Stress Test Suite Runner
# ############################################################################

proc runStressTestSuite*(casRoot: string, scale: string = "small"): StressTestSuite =
  ## Run the complete stress test suite
  result = StressTestSuite(
    name: fmt"NimPak Stress Tests ({scale} scale)",
    startTime: now()
  )
  
  # Determine scale factors
  let (chunkCount, fileCount, largeSize) = case scale.toLowerAscii():
    of "small": (SmallScale, 100, 10)
    of "medium": (MediumScale, 500, 50)
    of "large": (LargeScale, 2000, 100)
    of "huge": (HugeScale, 10000, 500)
    else: (SmallScale, 100, 10)
  
  echo "\n" & "=" .repeat(70)
  echo fmt"  NimPak Stress Test Suite - {scale.toUpperAscii()} Scale"
  echo "=" .repeat(70)
  let timeStr = result.startTime.format("yyyy-MM-dd HH:mm:ss")
  echo fmt"Start time: {timeStr}"
  echo fmt"Scale: chunks={chunkCount}, files={fileCount}"
  echo ""
  
  # Create test directory
  let testCasRoot = casRoot / fmt"stress_test_{epochTime().int}"
  createDir(testCasRoot)
  
  # Run tests
  var testNum = 1
  
  echo fmt"\n[{testNum}] Many Chunks Test"
  result.results.add(stressTestCasManyChunks(testCasRoot / "many_chunks", 
                                              chunkCount, SmallChunk))
  testNum += 1
  
  echo fmt"\n[{testNum}] Deduplication Test"
  result.results.add(stressTestCasDeduplication(testCasRoot / "dedup", 
                                                 chunkCount, 0.7))
  testNum += 1
  
  if scale in ["medium", "large", "huge"]:
    echo fmt"\n[{testNum}] Large Objects Test"
    result.results.add(stressTestCasLargeObjects(testCasRoot / "large", 
                                                  largeSize, 3))
    testNum += 1
  
  echo fmt"\n[{testNum}] Concurrent Operations Test"
  result.results.add(stressTestConcurrentStores(testCasRoot / "concurrent", 
                                                 4, chunkCount div 4))
  testNum += 1
  
  echo fmt"\n[{testNum}] Many Small Files Test"
  result.results.add(stressTestManySmallFiles(testCasRoot / "files", fileCount))
  testNum += 1
  
  echo fmt"\n[{testNum}] Garbage Collection Test"
  result.results.add(stressTestGarbageCollection(testCasRoot / "gc", 
                                                  chunkCount div 2))
  testNum += 1
  
  # Cleanup
  echo "\nCleaning up test data..."
  removeDir(testCasRoot)
  
  result.endTime = now()
  
  # Calculate totals
  for r in result.results:
    if r.passed:
      result.totalPassed += 1
    else:
      result.totalFailed += 1
  
  # Print summary
  echo "\n" & "=" .repeat(70)
  echo "  STRESS TEST SUMMARY"
  echo "=" .repeat(70)
  echo fmt"Duration: {(result.endTime - result.startTime).inSeconds}s"
  echo fmt"Passed: {result.totalPassed}/{result.results.len}"
  echo fmt"Failed: {result.totalFailed}"
  echo ""
  
  for r in result.results:
    let status = if r.passed: "✅ PASS" else: "❌ FAIL"
    echo fmt"{status} {r.name} ({r.duration:.2f}s)"
    for note in r.notes:
      echo fmt"       {note}"

# ############################################################################
# Main Entry Point
# ############################################################################

when isMainModule:
  import std/parseopt
  
  var scale = "small"
  var casRoot = getTempDir() / "nip_stress_test"
  
  var p = initOptParser()
  while true:
    p.next()
    case p.kind
    of cmdEnd: break
    of cmdShortOption, cmdLongOption:
      case p.key
      of "scale", "s": scale = p.val
      of "cas-root": casRoot = p.val
      of "help", "h":
        echo """
NimPak Stress Testing

Usage: test_stress [options]

Options:
  -s, --scale=SCALE    Test scale: small, medium, large, huge (default: small)
  --cas-root=PATH      CAS root directory (default: /tmp/nip_stress_test)
  -h, --help           Show this help

Scales:
  small   100 chunks, 100 files (quick validation)
  medium  1000 chunks, 500 files (~1 minute)
  large   10000 chunks, 2000 files (~5 minutes)
  huge    100000 chunks, 10000 files (~30+ minutes)
"""
        quit(0)
    of cmdArgument:
      discard
  
  createDir(casRoot)
  let suite = runStressTestSuite(casRoot, scale)
  
  if suite.totalFailed > 0:
    quit(1)