## Run each scenario in its own ThreadSanitizer-instrumented process.
## Success means no race was reported; a crash is not required to expose UB.
import std/[atomics, json, locks, options, os, times]
import ../src/namespaced_logging
import ../src/namespaced_logging/testutil
import ./runtime_allocator

const iterations = 1_000
const collectorIterations = 100
var start: Atomic[bool]
var cloneLock: Lock
var ready: Atomic[int]

proc awaitStart() =
  while not start.load():
    discard

proc serviceWorker(service: LogService) {.thread.} =
  awaitStart()
  for i in 0 ..< iterations:
    let local = service.threadLocalRef()
    let logger = local.getLogger("worker")
    doAssert logger.isEnabled(lvlInfo)

proc independentWorker() {.thread.} =
  awaitStart()
  var service = initLogService()
  service.setThreshold("worker", lvlInfo)
  for i in 0 ..< iterations:
    let local = service.threadLocalRef()
    let logger = local.getLogger("worker")
    doAssert logger.isEnabled(lvlInfo)

type
  AppenderLifetimes = object
    created, destroyed: Atomic[int]
  TrackedAppenderObj = object of LogAppender
    lifetimes: ptr AppenderLifetimes
  TrackedAppender = ref TrackedAppenderObj

proc `=destroy`(app: TrackedAppenderObj) =
  if app.lifetimes != nil:
    app.lifetimes.destroyed.atomicInc()
  `=destroy`(app.namespace)

proc trackedAppender(lifetimes: ptr AppenderLifetimes): TrackedAppender =
  lifetimes.created.atomicInc()
  TrackedAppender(lifetimes: lifetimes)

method clone(app: TrackedAppender): LogAppender {.gcsafe.} =
  trackedAppender(app.lifetimes)

proc lifetimeWorker(source: ptr LogService) {.thread.} =
  # Borrow the original only until this thread owns a local service. The
  # readiness barrier lets the parent release it before we continue.
  let local = source[].threadLocalRef()
  ready.atomicInc()
  awaitStart()
  for i in 0 ..< iterations:
    var handle = local.service
    var copied = handle
    handle = LogService()
    var moved = move(copied)
    moved.setThreshold("worker", lvlInfo)
    let nested = moved.threadLocalRef()
    doAssert nested.getLogger("worker").isEnabled(lvlInfo)
  # Local caches and the last shared owner are destroyed on worker threads.

proc reconfigureWorker(service: LogService) {.thread.} =
  var handle = service
  let local = service.threadLocalRef()
  awaitStart()
  for i in 0 ..< iterations:
    handle.setThreshold("worker", lvlInfo)
    local.ensureFreshness()
    doAssert local.getLogger("worker").isEnabled(lvlInfo)

type
  CustomState = object
    value: int
  StatefulAppender = CustomLogAppender[SharedLogState[CustomState]]

proc ignoreMessage(state: var SharedLogState[CustomState], msg: LogMessage) {.gcsafe, noconv.} =
  discard

proc stateWorker(source: ptr StatefulAppender) {.thread.} =
  awaitStart()
  for i in 0 ..< iterations:
    var copied: StatefulAppender
    withLock cloneLock:
      copied = StatefulAppender(clone(source[]))
    doAssert copied.state.snapshot().value == 42

proc makeFormatter(prefix: string): StatefulLogFormatterFunc =
  result = proc(msg: LogMessage): string {.gcsafe.} = prefix & msg.msg

proc formatterWorker(source: ptr ConsoleLogAppender) {.thread.} =
  awaitStart()
  # This formatter only needs msg. Avoid clock/timezone initialization so
  # sanitizer diagnostics identify closure ownership rather than libc state.
  let msg = LogMessage(msg: "message")
  for i in 0 ..< iterations:
    var copied: ConsoleLogAppender
    withLock cloneLock:
      copied = ConsoleLogAppender(clone(source[]))
    doAssert copied.formatter.format(msg) == "prefix: message"

type
  SharedPayload = object
    entries: seq[string]
    destructions: ptr Atomic[int]
    finalCount: ptr Atomic[int]

proc `=destroy`(value: SharedPayload) =
  if value.destructions != nil:
    value.finalCount[].store(value.entries.len)
    value.destructions[].atomicInc()
  `=destroy`(value.entries)

proc sharedStateWorker(source: ptr SharedLogState[SharedPayload]) {.thread.} =
  let state = source[]
  ready.atomicInc()
  awaitStart()
  for i in 0 ..< iterations:
    var copied = state
    let moved = move(copied)
    moved.withState(proc(value: var SharedPayload) =
      value.entries.add("entry")
      doAssert value.entries[^1] == "entry")

type
  FactoryState = ref object
    count: int
    self: FactoryState
    output: SharedLogState[seq[int]]

proc makeState(output: SharedLogState[seq[int]]): FactoryState =
  result = FactoryState(output: output)
  result.self = result

proc logState(state: var FactoryState, msg: LogMessage) {.gcsafe, noconv.} =
  inc state.count
  doAssert state.self == state
  let count = state.count
  state.output.withState(proc(values: var seq[int]) = values.add(count))

proc factoryWorker(source: ptr LogService) {.thread.} =
  block:
    let local = source[].threadLocalRef()
    ready.atomicInc()
    awaitStart()
    let logger = local.getLogger("worker")
    let stamp = fromUnix(0).utc
    for i in 0 ..< iterations:
      logger.log(LogMessage(scope: "worker", level: lvlInfo, ts: stamp, msg: "message"))
  GC_fullCollect()

proc timestampWorker() {.thread.} =
  awaitStart()
  for i in 0 ..< iterations:
    for msg in [initLogMessage("worker", lvlInfo, "message"),
        initLogMessage("worker", lvlInfo, %*{"msg": "message"})]:
      doAssert msg.msg == "message"
      doAssert msg.ts.timezone.name == "Etc/UTC"
      doAssert msg.ts.utcOffset == 0
      doAssert not msg.ts.isDst
      doAssert parseJson(formatJsonStructuredLog(msg))["ts"].getStr.len > 0

proc collectorWorker(collector: LoggedMessages) {.thread.} =
  # Construct every managed field in this worker, including the timezone.
  let stamp = fromUnix(0).utc
  awaitStart()
  for i in 0 ..< collectorIterations:
    collector.add(LogMessage(scope: "worker", level: lvlInfo, ts: stamp,
      msg: "message", err: some[ref Exception](newException(ValueError, "error")),
      additionalData: %*{"number": i}))
    var snapshot = collector.get()
    doAssert snapshot.len > 0
    doAssert snapshot[^1].err.get.msg == "error"
    snapshot[0].msg = "private copy"
    doAssert collector.get()[0].msg == "message"

type
  FormatterSeed = object
    prefix: string
    destroyed: ptr Atomic[int]
  FormatterProbeObj = object
    count: int
    destroyed: ptr Atomic[int]
  FormatterProbe = ref FormatterProbeObj

proc `=destroy`(probe: FormatterProbeObj) =
  if probe.destroyed != nil: probe.destroyed[].atomicInc()

proc makeCountedFormatter(seed: FormatterSeed): StatefulLogFormatterFunc =
  let probe = FormatterProbe(destroyed: seed.destroyed)
  result = proc(msg: LogMessage): string {.gcsafe.} =
    inc probe.count
    seed.prefix & $probe.count & msg.msg

proc formatterLifetimeWorker(source: ptr LogFormatter) {.thread.} =
  block:
    var formatter: LogFormatter
    withLock cloneLock:
      formatter = clone(source[])
    ready.atomicInc()
    awaitStart()
    doAssert formatter.format(LogMessage(msg: "!")) == "count 1!"
    for i in 0 ..< iterations:
      let copied = clone(formatter)
      doAssert copied.format(LogMessage(msg: "!")) == "count 1!"
      doAssert copied.format(LogMessage(msg: "!")) == "count 2!"
  GC_fullCollect()

proc main() =
  doAssert paramCount() == 1
  initLock(cloneLock)
  defer: deinitLock(cloneLock)
  case paramStr(1)
  of "independent":
    var workers: array[4, Thread[void]]
    for i in 0 ..< workers.len:
      createThread(workers[i], independentWorker)
    start.store(true)
    joinThreads(workers)
  of "service":
    var service = initLogService()
    service.setThreshold("worker", lvlInfo)
    service.addAppender(initConsoleLogAppender())
    var workers: array[4, Thread[LogService]]
    for i in 0 ..< workers.len:
      createThread(workers[i], serviceWorker, service)
    # No parent-side copying/configuration changes overlap the workers.
    # The parent retains service and all thread argument storage until join.
    start.store(true)
    joinThreads(workers)
  of "custom-state":
    var source = initCustomLogAppender(initSharedLogState(CustomState(value: 42)), ignoreMessage)
    var workers: array[4, Thread[ptr StatefulAppender]]
    for i in 0 ..< workers.len:
      createThread(workers[i], stateWorker, addr source)
    start.store(true)
    joinThreads(workers)
  of "service-lifetime":
    var lifetimes: AppenderLifetimes
    var service = initLogService()
    block:
      let original = trackedAppender(addr lifetimes)
      service.addAppender(original)
    var workers: array[4, Thread[ptr LogService]]
    for i in 0 ..< workers.len:
      createThread(workers[i], lifetimeWorker, addr service)
    while ready.load() != workers.len:
      discard
    # Thread arguments contain only a borrowed pointer, so there are no hidden
    # owning handles in Thread.data keeping the control block alive here.
    service = LogService()
    start.store(true)
    joinThreads(workers)
    doAssert lifetimes.created.load() > workers.len
    doAssert lifetimes.destroyed.load() == lifetimes.created.load()
  of "service-reconfigure":
    var service = initLogService()
    service.addAppender(initConsoleLogAppender())
    var workers: array[4, Thread[LogService]]
    for i in 0 ..< workers.len:
      createThread(workers[i], reconfigureWorker, service)
    start.store(true)
    joinThreads(workers)
  of "formatter":
    var source = initConsoleLogAppender()
    source.formatter = initLogFormatter("prefix: ", makeFormatter)
    var workers: array[4, Thread[ptr ConsoleLogAppender]]
    for i in 0 ..< workers.len:
      createThread(workers[i], formatterWorker, addr source)
    start.store(true)
    joinThreads(workers)
  of "shared-state-lifetime":
    var destructions, finalCount: Atomic[int]
    var state = initSharedLogState(SharedPayload(
      destructions: addr destructions, finalCount: addr finalCount))
    var workers: array[4, Thread[ptr SharedLogState[SharedPayload]]]
    for i in 0 ..< workers.len:
      createThread(workers[i], sharedStateWorker, addr state)
    while ready.load() != workers.len: discard
    state = default(typeof(state))
    start.store(true)
    joinThreads(workers)
    doAssert destructions.load() == 1
    doAssert finalCount.load() == workers.len * iterations
  of "factory-state":
    let output = initSharedLogState(newSeq[int]())
    var service = initLogService()
    service.addAppender(initFactoryCustomLogAppender(output, makeState, logState))
    var workers: array[4, Thread[ptr LogService]]
    for i in 0 ..< workers.len:
      createThread(workers[i], factoryWorker, addr service)
    while ready.load() != workers.len: discard
    service = LogService()
    start.store(true)
    joinThreads(workers)
    let values = output.snapshot()
    doAssert values.len == workers.len * iterations
    var counts: array[iterations + 1, int]
    for value in values: inc counts[value]
    for i in 1 .. iterations: doAssert counts[i] == workers.len
  of "collector":
    let collector = initLoggedMessages()
    var workers: array[4, Thread[LoggedMessages]]
    for i in 0 ..< workers.len:
      createThread(workers[i], collectorWorker, collector)
    start.store(true)
    joinThreads(workers)
    doAssert collector.get().len == workers.len * collectorIterations
  of "formatter-lifetime":
    var destroyed: Atomic[int]
    var formatter = initLogFormatter(
      FormatterSeed(prefix: "count ", destroyed: addr destroyed), makeCountedFormatter)
    doAssert formatter.format(LogMessage(msg: "!")) == "count 1!"
    var workers: array[4, Thread[ptr LogFormatter]]
    for i in 0 ..< workers.len:
      createThread(workers[i], formatterLifetimeWorker, addr formatter)
    while ready.load() != workers.len: discard
    formatter = nil
    start.store(true)
    joinThreads(workers)
    GC_fullCollect()
    doAssert destroyed.load() == 1 + workers.len * (iterations + 1)
  of "timestamp":
    var workers: array[4, Thread[void]]
    for i in 0 ..< workers.len:
      createThread(workers[i], timestampWorker)
    start.store(true)
    joinThreads(workers)
  of "allocator":
    runAllocatorRegression()
  else:
    raise newException(ValueError, "unknown thread safety scenario")

# The pointer-based scenarios borrow immutable source appenders whose owner
# outlives all workers. This isolates clone/destruction races from LogService
# copying and from ordinary unsynchronized application-state mutation.
main()