#!/usr/bin/env julia
#
# bench.jl — measure end-to-end throughput and server memory for a running FileServer.
#
# Two things make this pipeline awkward to benchmark with off-the-shelf tools
# (ab/hey/wrk), and both shape what this script does:
#
#   1. HTTP latency is not throughput. /upload returns 202 as soon as the bytes
#      are spooled and a reference is enqueued — all four stages run afterwards.
#      So real throughput is the *arrival rate at the terminal sinks*
#      (done/, text_done/, binary/, failed/), not the response rate. We upload a
#      corpus, then poll the sinks until the file count stops moving.
#
#   2. Memory should be flat in file size, and that claim needs checking on two
#      axes. The workers read bounded prefixes (16+16 bytes to classify, 8 KB to
#      sniff, 64 KB to language-detect), and intake streams each upload from the
#      socket to the spool file a chunk at a time (FS_UPLOAD_CHUNK_BYTES, see
#      src/multipart.jl). So peak RSS should track *concurrency*, not file size:
#      sweeping --size at fixed --concurrency should be a flat line, and that is
#      the regression this measures. We sample the server's RSS throughout and
#      report the high-water mark.
#
#      (Before intake was streamed it buffered each upload whole, several times
#      over, and a 256 MiB upload grew RSS by ~700 MiB. If a --size sweep ever
#      slopes upward again, something has started buffering.)
#
# This is an external harness: it makes no assumptions about the server beyond
# the HTTP contract and the on-disk sink layout, and requires no changes to src/.
# It must run on the same machine as the server (it reads sink dirs and /proc).
#
# Usage:
#   julia --project=. -t auto bin/bench.jl [options]
#
#   --url URL          server base URL (default: $FS_URL or http://127.0.0.1:8080)
#   --files N          number of files to upload (default: 200)
#   --size S           size of each generated file, e.g. 4k, 512k, 8m, 1g (default: 64k)
#   --concurrency J    uploads in flight at once (default: 8)
#   --kind K           binary | text | mixed — what to generate (default: binary)
#   --corpus DIR       upload an existing directory instead of generating
#                      (the only way to exercise stage 2: point it at real known files)
#   --keep-corpus      don't delete the generated corpus on exit
#   --pid PID          server pid for memory sampling (default: autodetect)
#   --no-mem           skip memory sampling entirely
#   --sample-ms MS     sink/RSS sampling interval (default: 200)
#   --timeout SEC      give up after this long with no drain progress (default: 120)
#   --json PATH        also write the results as JSON
#   --force            skip the corpus-size safety check
#
# Examples:
#   # throughput: many small files, high concurrency
#   julia --project=. -t auto bin/bench.jl --files 2000 --size 8k --concurrency 32
#
#   # memory: flat in file size? sweep --size with concurrency pinned
#   julia --project=. -t auto bin/bench.jl --files 4 --size 256m --concurrency 1
#   julia --project=. -t auto bin/bench.jl --files 2 --size 1g   --concurrency 1
#   julia --project=. -t auto bin/bench.jl --files 8 --size 1g   --concurrency 8
#
#   # stage 2 (exiftool) with real known files
#   julia --project=. -t auto bin/bench.jl --corpus ../training_set --concurrency 16

using HTTP
using JSON3
using Random

# ---------------------------------------------------------------- option parsing

const DEFAULTS = Dict{String,Any}(
    "url"         => get(ENV, "FS_URL", "http://127.0.0.1:8080"),
    "files"       => 200,
    "size"        => 64 * 1024,
    "concurrency" => 8,
    "kind"        => "binary",
    "corpus"      => nothing,
    "keep-corpus" => false,
    "pid"         => nothing,
    "no-mem"      => false,
    "sample-ms"   => 200,
    "timeout"     => 120,
    "json"        => nothing,
    "force"       => false,
)

const FLAGS = ("keep-corpus", "no-mem", "force")

# Approximate RSS of a freshly started server (Lux + the loaded classifier + the
# language detector, measured on Julia 1.12 / -t auto). Only used to notice that
# a baseline is inflated by a previous run's un-returned GC memory, so a rough
# figure is enough.
const FRESH_RSS_HINT = 950 * 1024^2

"Parse `4k`/`8M`/`1g`/`4096` into a byte count."
function parse_size(s::AbstractString)::Int
    m = match(r"^(\d+(?:\.\d+)?)\s*([kKmMgG]?)[bB]?$", strip(s))
    m === nothing && error("bad --size: $s (expected e.g. 512, 64k, 8m, 1g)")
    mult = Dict('k' => 1024, 'm' => 1024^2, 'g' => 1024^3)
    scale = isempty(m[2]) ? 1 : mult[lowercase(m[2])[1]]
    return round(Int, parse(Float64, m[1]) * scale)
end

function parse_args(argv)::Dict{String,Any}
    opts = copy(DEFAULTS)
    i = 1
    while i <= length(argv)
        a = argv[i]
        startswith(a, "--") || error("unexpected argument: $a (see the header of $(PROGRAM_FILE))")
        key = a[3:end]
        haskey(opts, key) || error("unknown option: $a")
        if key in FLAGS
            opts[key] = true
            i += 1
            continue
        end
        i + 1 <= length(argv) || error("option --$key needs a value")
        val = argv[i+1]
        opts[key] = key == "size"                                  ? parse_size(val) :
                    key in ("files", "concurrency", "sample-ms")   ? parse(Int, val) :
                    key == "timeout"                               ? parse(Float64, val) :
                    key == "pid"                                   ? parse(Int, val) :
                    val
        i += 2
    end
    opts["kind"] in ("binary", "text", "mixed") ||
        error("--kind must be binary, text or mixed (got $(opts["kind"]))")
    opts["files"] >= 1 || error("--files must be >= 1")
    opts["concurrency"] >= 1 || error("--concurrency must be >= 1")
    return opts
end

# ------------------------------------------------------------------------- dirs
#
# Resolved from the same environment variables src/config.jl reads, so a server
# started with custom dirs is benchmarked correctly. Kept as a standalone table
# rather than `using FileServer` so the harness doesn't pay to load Lux.

sinkdirs() = (
    done      = get(ENV, "FS_DONE_DIR",      "data/done"),
    text_done = get(ENV, "FS_TEXT_DONE_DIR", "data/text_done"),
    binary    = get(ENV, "FS_BINARY_DIR",    "data/binary"),
    failed    = get(ENV, "FS_FAILED_DIR",    "data/failed"),
)

stagedirs() = (
    spool   = get(ENV, "FS_SPOOL_DIR",   "data/spool"),
    known   = get(ENV, "FS_KNOWN_DIR",   "data/known"),
    unknown = get(ENV, "FS_UNKNOWN_DIR", "data/unknown"),
    text    = get(ENV, "FS_TEXT_DIR",    "data/text"),
)

"Count work items in `dir`, ignoring the .meta.json sidecars stages 2/4 write."
function count_files(dir::AbstractString)::Int
    isdir(dir) || return 0
    n = 0
    for name in readdir(dir)
        endswith(name, ".meta.json") && continue
        isfile(joinpath(dir, name)) && (n += 1)
    end
    return n
end

counts(dirs) = NamedTuple{keys(dirs)}(map(count_files, values(dirs)))
total(c) = sum(values(c))
deltas(now_, base) = NamedTuple{keys(now_)}(map(-, values(now_), values(base)))

# ----------------------------------------------------------------------- memory

"Read (VmRSS, VmHWM) in bytes for `pid`, or `nothing` if unreadable."
function read_rss(pid::Int)
    rss = hwm = nothing
    try
        for line in eachline("/proc/$pid/status")
            if startswith(line, "VmRSS:")
                rss = parse(Int, split(line)[2]) * 1024
            elseif startswith(line, "VmHWM:")
                hwm = parse(Int, split(line)[2]) * 1024
            end
        end
    catch
        return nothing
    end
    return (rss === nothing || hwm === nothing) ? nothing : (rss, hwm)
end

"Find the running server process, or `nothing`."
function detect_pid()
    out = try
        readchomp(`pgrep -f "bin/server.jl"`)
    catch
        return nothing
    end
    pids = parse.(Int, split(out))
    isempty(pids) && return nothing
    length(pids) > 1 && @warn "multiple server processes matched; sampling the first" pids
    return first(pids)
end

"""
Reset the kernel's peak-RSS counter so VmHWM reflects only this run.

Without it, VmHWM carries the high-water mark from startup (model load) or from
an earlier benchmark, which would silently dominate a small run's result.
Requires the server to run as the same user; on failure we say so and fall back
to sampled VmRSS, which can miss a spike between samples.
"""
function reset_peak_rss(pid::Int)::Bool
    try
        write("/proc/$pid/clear_refs", "5")
        return true
    catch
        return false
    end
end

# ------------------------------------------------------------------------ corpus

const WORDS = split("the quick brown fox jumps over a lazy dog while parsing " *
                    "headers and spooling bytes onto disk for later enrichment " *
                    "because throughput matters more than latency here")

"Write one file of exactly `size` bytes, in bounded chunks so the generator
itself never holds a whole 1 GB file in memory."
function write_file(path::AbstractString, size::Int, kind::Symbol, rng)
    chunk = 1024 * 1024
    open(path, "w") do io
        remaining = size
        while remaining > 0
            n = min(chunk, remaining)
            if kind === :binary
                write(io, rand(rng, UInt8, n))
            else
                buf = IOBuffer()
                while buf.size < n
                    print(buf, rand(rng, WORDS), rand(rng) < 0.06 ? ".\n" : " ")
                end
                write(io, take!(buf)[1:n])
            end
            remaining -= n
        end
    end
    return nothing
end

"Generate the corpus and return (dir, paths, total_bytes)."
function make_corpus(opts)
    n, size, kind = opts["files"], opts["size"], opts["kind"]
    totalbytes = n * size
    if totalbytes > 16 * 1024^3 && !opts["force"]
        error("corpus would be $(human(totalbytes)) on disk; pass --force if that's intended")
    end
    dir = mktempdir(; prefix = "fsbench-")
    rng = MersenneTwister(1234)
    paths = String[]
    for i in 1:n
        k = kind == "mixed" ? (isodd(i) ? :binary : :text) :
            kind == "text"  ? :text : :binary
        ext = k === :text ? "txt" : "bin"
        path = joinpath(dir, "bench-$(lpad(i, 6, '0')).$ext")
        write_file(path, size, k, rng)
        push!(paths, path)
    end
    return dir, paths, totalbytes
end

function existing_corpus(dir::AbstractString)
    isdir(dir) || error("--corpus is not a directory: $dir")
    paths = sort(filter(isfile, readdir(dir; join = true)))
    isempty(paths) && error("--corpus directory is empty: $dir")
    return paths, sum(filesize, paths)
end

# ----------------------------------------------------------------------- upload

struct Upload
    status::Int      # HTTP status, or 0 if the request threw
    accepted::Int    # jobs the server actually queued (from the 202/503 body)
    seconds::Float64
end

"POST one file as multipart/form-data and report what the server accepted."
function upload_one(url::String, path::String)::Upload
    t0 = time()
    try
        form = HTTP.Form(["file" => HTTP.Multipart(basename(path), open(path, "r"),
                                                   "application/octet-stream")])
        resp = HTTP.post(url, [], form; status_exception = false, retry = false)
        # 202 and 503 both carry an `accepted` array: a partially-accepted batch
        # still queued those jobs, and they will show up in the sinks.
        acc = try
            length(JSON3.read(String(resp.body)).accepted)
        catch
            resp.status == 202 ? 1 : 0
        end
        return Upload(resp.status, acc, time() - t0)
    catch e
        e isa InterruptException && rethrow()
        @warn "upload failed" file = basename(path) exception = e
        return Upload(0, 0, time() - t0)
    end
end

"""
Upload every path, at most `concurrency` in flight.

A bounded set of worker tasks pulling from a shared index keeps exactly
`concurrency` requests in flight for the whole run — unlike batching, where each
batch stalls on its slowest (largest) file and the real concurrency sags.
"""
function upload_all(url::String, paths::Vector{String}, concurrency::Int)
    endpoint = string(rstrip(url, '/'), "/upload")
    results = Vector{Upload}(undef, length(paths))
    next = Threads.Atomic{Int}(1)
    @sync for _ in 1:min(concurrency, length(paths))
        Threads.@spawn while true
            i = Threads.atomic_add!(next, 1)
            i > length(paths) && break
            results[i] = upload_one(endpoint, paths[i])
        end
    end
    return results
end

# ------------------------------------------------------------------- formatting

function human(bytes::Real)
    b = Float64(bytes)
    for unit in ("B", "KiB", "MiB", "GiB", "TiB")
        (abs(b) < 1024 || unit == "TiB") && return "$(round(b; digits = 2)) $unit"
        b /= 1024
    end
end

fmt(x::Real, digits::Int = 2) = string(round(Float64(x); digits = digits))

function percentile(sorted::Vector{Float64}, p::Float64)
    isempty(sorted) && return NaN
    idx = clamp(ceil(Int, p * length(sorted)), 1, length(sorted))
    return sorted[idx]
end

# ------------------------------------------------------------------------- main

function main(argv)
    opts = parse_args(argv)
    url  = string(opts["url"])

    # Fail fast and clearly if there's no server, rather than reporting a run of zeros.
    try
        HTTP.get(string(rstrip(url, '/'), "/health"); retry = false, readtimeout = 5)
    catch e
        println(stderr, "cannot reach $url/health — is the server running?")
        println(stderr, "  start it with: julia --project=. -t auto bin/server.jl")
        return 1
    end

    sinks, stages = sinkdirs(), stagedirs()

    # Leftovers mid-pipeline would land in the sinks during our window and be
    # counted as our throughput, so say so up front rather than quietly skewing.
    pending = total(counts(stages))
    pending > 0 && @warn "pipeline is not idle: $pending file(s) in the stage dirs; " *
                         "throughput will include their completions"

    # --- corpus
    generated = opts["corpus"] === nothing
    corpusdir, paths, corpusbytes = if generated
        print("generating corpus: $(opts["files"]) × $(human(opts["size"])) ($(opts["kind"]))… ")
        t = time()
        d, p, b = make_corpus(opts)
        println("done in $(fmt(time() - t))s → $d")
        d, p, b
    else
        p, b = existing_corpus(String(opts["corpus"]))
        println("using corpus: $(length(p)) file(s), $(human(b)) from $(opts["corpus"])")
        String(opts["corpus"]), p, b
    end

    try
        pid = opts["no-mem"] ? nothing : something(opts["pid"], detect_pid(), Some(nothing))
        if pid === nothing && !opts["no-mem"]
            @warn "could not find the server process; skipping memory (pass --pid PID)"
        end

        baseline_rss = nothing
        peak_reset = false
        if pid !== nothing
            r = read_rss(pid)
            r === nothing && (@warn "cannot read /proc/$pid/status; skipping memory"; pid = nothing)
            if pid !== nothing
                baseline_rss = r[1]
                peak_reset = reset_peak_rss(pid)
                peak_reset || @warn "could not reset the peak-RSS counter (/proc/$pid/clear_refs); " *
                                    "reporting sampled RSS only"
            end
        end

        if Threads.nthreads() == 1 && opts["size"] > 64 * 1024^2
            @warn "running with 1 thread and large files: the sampler shares a thread with " *
                  "blocking file reads, so the RSS curve will be coarse. Prefer -t auto."
        end

        base_sinks = counts(sinks)
        interval = opts["sample-ms"] / 1000

        # --- sampler: RSS curve + stage depths, for bottleneck attribution.
        stop = Threads.Atomic{Bool}(false)
        rss_samples = Float64[]
        depth_max = Dict(k => 0 for k in keys(stages))
        sampler = Threads.@spawn begin
            while !stop[]
                if pid !== nothing
                    r = read_rss(pid)
                    r !== nothing && push!(rss_samples, Float64(r[1]))
                end
                d = counts(stages)
                for k in keys(d)
                    depth_max[k] = max(depth_max[k], getfield(d, k))
                end
                sleep(interval)
            end
        end

        # --- intake
        println("uploading $(length(paths)) file(s) at concurrency $(opts["concurrency"])…")
        t_start = time()
        ups = upload_all(url, paths, opts["concurrency"])
        t_intake_end = time()

        accepted = sum(u.accepted for u in ups)
        n_202 = count(u -> u.status == 202, ups)
        n_503 = count(u -> u.status == 503, ups)
        n_err = count(u -> !(u.status in (202, 503)), ups)
        intake_secs = t_intake_end - t_start
        lat = sort([u.seconds for u in ups])

        println("  intake: $accepted job(s) accepted in $(fmt(intake_secs))s " *
                "($(fmt(accepted / max(intake_secs, 1e-9))) files/s, " *
                "$(fmt(corpusbytes / max(intake_secs, 1e-9) / 1024^2)) MiB/s)")
        n_503 > 0 && println("  backpressure: $n_503 request(s) got 503 (intake queue full)")
        n_err > 0 && println("  errors: $n_err request(s) failed or returned an unexpected status")

        # --- drain: poll the terminal sinks until they stop moving.
        println("draining (polling sinks every $(opts["sample-ms"])ms)…")
        completed = 0
        t_last_progress = time()
        t_last_completion = t_intake_end
        timed_out = false
        while completed < accepted
            sleep(interval)
            c = total(deltas(counts(sinks), base_sinks))
            if c > completed
                completed = c
                t_last_completion = time()
                t_last_progress = t_last_completion
            elseif time() - t_last_progress > opts["timeout"]
                timed_out = true
                break
            end
        end
        stop[] = true
        wait(sampler)

        sink_delta = deltas(counts(sinks), base_sinks)
        e2e_secs = t_last_completion - t_start

        final_rss = pid === nothing ? nothing : read_rss(pid)
        peak_rss = if final_rss !== nothing && peak_reset
            final_rss[2]                      # kernel VmHWM: catches spikes between samples
        elseif !isempty(rss_samples)
            maximum(rss_samples)
        else
            nothing
        end

        # --- report
        println()
        println("=" ^ 68)
        println("corpus            $(length(paths)) file(s), $(human(corpusbytes)) total, " *
                "$(human(corpusbytes / length(paths))) avg")
        println("concurrency       $(opts["concurrency"])")
        println()
        println("INTAKE (HTTP 202 — bytes spooled, not processed)")
        println("  accepted        $accepted of $(length(paths))   [202: $n_202, 503: $n_503, error: $n_err]")
        println("  wall            $(fmt(intake_secs))s")
        println("  rate            $(fmt(accepted / max(intake_secs, 1e-9))) files/s, " *
                "$(fmt(corpusbytes / max(intake_secs, 1e-9) / 1024^2)) MiB/s")
        println("  latency         p50 $(fmt(percentile(lat, 0.5) * 1000, 1))ms  " *
                "p95 $(fmt(percentile(lat, 0.95) * 1000, 1))ms  " *
                "max $(fmt(percentile(lat, 1.0) * 1000, 1))ms")
        println()
        println("END-TO-END (files reaching a terminal sink)")
        println("  completed       $completed of $accepted accepted" * (timed_out ? "   ** TIMED OUT **" : ""))
        println("  wall            $(fmt(e2e_secs))s  (first upload → last completion)")
        println("  throughput      $(fmt(completed / max(e2e_secs, 1e-9))) files/s, " *
                "$(fmt(corpusbytes / max(e2e_secs, 1e-9) / 1024^2)) MiB/s")
        println("  sinks           done $(sink_delta.done)  text_done $(sink_delta.text_done)  " *
                "binary $(sink_delta.binary)  failed $(sink_delta.failed)")
        println("  peak depth      spool $(depth_max[:spool])  known $(depth_max[:known])  " *
                "unknown $(depth_max[:unknown])  text $(depth_max[:text])")
        println("                  (the stage that backs up is the bottleneck)")
        println()
        if peak_rss !== nothing
            println("SERVER MEMORY (pid $pid)")
            println("  baseline RSS    $(human(baseline_rss))")
            println("  peak RSS        $(human(peak_rss))" *
                    (peak_reset ? "  (kernel VmHWM, reset at start)" : "  (sampled — may miss spikes)"))
            println("  growth          $(human(peak_rss - baseline_rss))")
            println("  per in-flight   $(human((peak_rss - baseline_rss) / opts["concurrency"])) " *
                    "at $(human(corpusbytes / length(paths))) avg file size")
            println("                  (should not grow with file size — intake streams to disk)")
            # Julia's GC returns memory to the OS lazily, so a second run on the
            # same process starts from an inflated baseline and under-reports
            # growth. Absolute peak is the number to trust across runs.
            baseline_rss > 1.3 * FRESH_RSS_HINT &&
                println("  ! baseline is well above a fresh start ($(human(FRESH_RSS_HINT))): the GC " *
                        "has not\n                    returned memory from earlier work. Compare " *
                        "absolute peak, or restart\n                    the server for a clean growth figure.")
        else
            println("SERVER MEMORY     not sampled")
        end
        println("=" ^ 68)

        sink_delta.failed > 0 &&
            println("\nnote: $(sink_delta.failed) file(s) landed in $(sinks.failed) — check the server log.")
        timed_out &&
            println("\nnote: drain stalled with $(accepted - completed) file(s) outstanding. " *
                    "Check the server log and the stage dirs; raise --timeout if the pipeline is just slow.")

        if opts["json"] !== nothing
            result = (
                url, concurrency = opts["concurrency"], kind = opts["kind"],
                files = length(paths), corpus_bytes = corpusbytes,
                avg_file_bytes = corpusbytes / length(paths),
                intake = (; accepted, n_202, n_503, n_err, seconds = intake_secs,
                            files_per_sec = accepted / max(intake_secs, 1e-9),
                            p50_ms = percentile(lat, 0.5) * 1000,
                            p95_ms = percentile(lat, 0.95) * 1000,
                            max_ms = percentile(lat, 1.0) * 1000),
                end_to_end = (; completed, seconds = e2e_secs, timed_out,
                                files_per_sec = completed / max(e2e_secs, 1e-9),
                                mib_per_sec = corpusbytes / max(e2e_secs, 1e-9) / 1024^2,
                                sinks = sink_delta, peak_stage_depth = depth_max),
                memory = (; pid, baseline_rss, peak_rss, peak_is_kernel_hwm = peak_reset,
                            growth = peak_rss === nothing ? nothing : peak_rss - baseline_rss,
                            samples = rss_samples),
            )
            open(String(opts["json"]), "w") do io
                JSON3.write(io, result)
            end
            println("\nwrote $(opts["json"])")
        end

        return timed_out ? 1 : 0
    finally
        if generated && !opts["keep-corpus"]
            rm(corpusdir; recursive = true, force = true)
        elseif generated
            println("\nkept corpus: $corpusdir")
        end
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    exit(main(ARGS))
end