#!/usr/bin/env julia
#
# bench_stage1.jl — take stage 1 apart and find the slowest component.
#
# bin/bench.jl reports stage 1 as a single number (the wall time of
# `handle_classify_job` under whole-pipeline contention) and bin/bench_model.jl
# takes the *classifier* apart. Neither answers "which part of stage 1 costs the
# most?", because stage 1 is more than the model. Per file it does:
#
#   classify        filesize + two 16-byte reads + a 32x1 forward pass
#     read_features   open, read head, seek, read tail, scale to Float32
#     Lux.apply       the network on a feature vector already in memory
#   move_to         rename spool/<f> -> known/<f> or unknown/<f>
#   enqueue         push a Job reference onto the downstream bounded queue
#   logging         two @info lines ("classified file", "routed to ...")
#
# This script times each of those in isolation, then times the real
# `handle_classify_job` end to end so the parts can be checked against the whole.
# The two knobs that most change the answer get their own sweeps:
#
#   * Logger. The server runs `FlushLogger(ConsoleLogger(stderr))` — it formats
#     and flushes every message. Under redirect (a log file, journald) that is a
#     syscall per line, two lines per file, on the hot path. We time the handler
#     under a null logger, a formatting-but-discarding logger, and the real
#     flushing-to-file logger, so the cost of logging is a subtraction, not a
#     guess. This sweep is what demoted stage 1's per-file lines to `@debug`
#     (see the note in src/worker.jl); the standalone `logging (...)` rows below
#     still price a *formatted* line, i.e. what those lines cost when switched
#     back on with `JULIA_DEBUG=FileServer`, while the handler rows show what the
#     stage pays with them off.
#   * Concurrency. Components that own a lock (the queue's condition, the
#     logger's stream) don't scale, and the ranking at one worker need not be the
#     ranking at sixteen. The `--threads` sweep runs the full handler across
#     worker counts.
#
# Usage:
#   julia --project=. -t auto bin/bench_stage1.jl [options]
#
#   --files N        files per timed pass for consuming benchmarks (default: 2000)
#   --reps N         calls per timed pass for non-consuming benchmarks (default: 20000)
#   --trials N       timed passes; the minimum is reported (default: 5)
#   --size SPEC      corpus file size (default: 64k)
#   --dir PATH       working directory for the corpus (default: a temp dir under data/)
#   --model PATH     classifier artifact (default: $FS_MODEL_PATH or model/classifier.jld2)
#   --threads LIST   worker counts for the concurrency sweep (default: 1,2,4,8,nthreads)
#   --no-threads     skip the concurrency sweep
#   --json PATH      also write the results as JSON
#
# Reported times are the *minimum* over trials: the floor is the signal and
# everything above it is scheduler, page-cache and GC noise.

using FileServer
using Lux
using JSON3
using Logging
using Printf
using Random

const FS = FileServer

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

const DEFAULTS = Dict{String,Any}(
    "files"      => 2000,
    "reps"       => 20_000,
    "trials"     => 5,
    "size"       => "64k",
    "dir"        => nothing,
    "model"      => get(ENV, "FS_MODEL_PATH", "model/classifier.jld2"),
    "threads"    => nothing,
    "no-threads" => false,
    "json"       => nothing,
)

const FLAGS = ("no-threads",)
const INTS  = ("files", "reps", "trials")

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)
    opts = copy(DEFAULTS)
    i = 1
    while i <= length(argv)
        a = argv[i]
        startswith(a, "--") || error("unexpected argument: $a")
        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")
        opts[key] = key in INTS ? parse(Int, argv[i+1]) : argv[i+1]
        i += 2
    end
    return opts
end

# ------------------------------------------------------------------- measurement

# Every timed loop stores its result here. Without a visible side effect the
# compiler is free to hoist a pure call out of the loop and we would be timing an
# empty `for`.
const SINK = Ref{Any}(nothing)

"""
    best_of(pass, prepare; trials) -> ns_per_op

Run `pass()` `trials` times and report the fastest, in nanoseconds per operation
(`pass` returns the number of operations it performed). `prepare()` runs before
each pass and is *not* timed — that is where a consuming benchmark puts the file
back where it started. `pass` comes first so callers can pass it as a `do` block.

The first pass is thrown away: it pays Julia's JIT compilation, which on calls
this small is orders of magnitude more than the thing being measured.
"""
function best_of(pass, prepare; trials::Int)
    best = Inf
    for t in 0:trials
        prepare()
        GC.gc()
        t0 = time_ns()
        n = pass()
        dt = Float64(time_ns() - t0)
        t == 0 && continue          # warm-up: compiled, not measured
        best = min(best, dt / n)
    end
    return best
end

noop() = nothing

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

function human_time(ns::Real)
    ns < 1_000      && return @sprintf("%.0f ns", ns)
    ns < 1_000_000  && return @sprintf("%.2f µs", ns / 1e3)
    ns < 1e9        && return @sprintf("%.2f ms", ns / 1e6)
    return @sprintf("%.2f s", ns / 1e9)
end

function human_rate(r::Real)
    r >= 1e6 && return @sprintf("%.2fM/s", r / 1e6)
    r >= 1e3 && return @sprintf("%.1fk/s", r / 1e3)
    return @sprintf("%.0f/s", r)
end

rate(ns::Real) = 1e9 / max(ns, 1e-9)

rule(n = 78) = println("-" ^ n)

function header(title)
    println()
    println(title)
    rule()
end

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

"""
Write a file of exactly `size` random bytes, in bounded chunks.

Content is random rather than zeros so the classifier sees a realistic input —
and so the filesystem can't cheat with a sparse file, which would make the tail
`seek` unrepresentatively fast.
"""
function write_file(path::AbstractString, size::Int, rng)
    chunk = 1024 * 1024
    open(path, "w") do io
        remaining = size
        while remaining > 0
            n = min(chunk, remaining)
            write(io, rand(rng, UInt8, n))
            remaining -= n
        end
    end
    return path
end

"""
    make_corpus(cfg, n, size, rng) -> Vector{Job}

Create `n` spooled files and the `Job` references a stage-1 worker would dequeue
for them — the exact input `handle_classify_job` sees.
"""
function make_corpus(cfg::FS.Config, n::Int, size::Int, rng)
    jobs = FS.Job[]
    for i in 1:n
        id, path = FS.spool_path(cfg, @sprintf("bench-%06d.bin", i))
        write_file(path, size, rng)
        push!(jobs, FS.Job(id, basename(path), path, size, time()))
    end
    return jobs
end

"""
    respool!(cfg, jobs)

Put every corpus file back in `spool/`, wherever the last pass left it (known/,
unknown/, or already home). This is the untimed `prepare` step for benchmarks
that consume their input by moving it.
"""
function respool!(cfg::FS.Config, jobs::Vector{FS.Job})
    for job in jobs
        isfile(job.path) && continue
        for dir in (cfg.known_dir, cfg.unknown_dir, cfg.failed_dir)
            candidate = joinpath(dir, basename(job.path))
            if isfile(candidate)
                mv(candidate, job.path; force = true)
                break
            end
        end
    end
    return nothing
end

"Drain a queue without blocking, so the next pass starts from empty."
function drain!(q::FS.ChannelQueue)
    while length(q) > 0
        FS.dequeue!(q)
    end
    return nothing
end

# --------------------------------------------------------------------- loggers

"""
    with_logger_named(name, path, f)

Run `f` under one of the three loggers the cost of logging is bracketed by:

  * `:null`    — `NullLogger`: the `@info` macro's own overhead, nothing else.
  * `:format`  — `ConsoleLogger` to `devnull`: message formatting and key/value
                 interpolation, but no I/O.
  * `:flush`   — `FlushLogger(ConsoleLogger(io))` to a real file: what
                 `FileServer.run` installs, under the redirect it was written for.
  * `:debug`   — the same, at `Debug` level: the stage's per-file lines are
                 `@debug`, so this is the equivalent of running the server with
                 `JULIA_DEBUG=FileServer` and the only setting under which they
                 are emitted at all.
"""
function with_logger_named(f, which::Symbol, path::AbstractString)
    if which === :null
        return with_logger(f, NullLogger())
    elseif which === :format
        return with_logger(f, ConsoleLogger(devnull))
    elseif which === :flush || which === :debug
        level = which === :debug ? Logging.Debug : Logging.Info
        return open(path, "w") do io
            with_logger(f, FS.FlushLogger(ConsoleLogger(io, level)))
        end
    end
    error("unknown logger: $which")
end

# ------------------------------------------------------------------ components

"""
    component_rows(cfg, clf, jobs, opts) -> Vector

Time each piece of stage 1 on its own. Non-consuming pieces (`filesize`,
`read_features`, `Lux.apply`, `classify`, the log lines) run `reps` times over
the corpus; consuming pieces (`move_to`, the full handler) run once per corpus
file with an untimed reset between passes.
"""
function component_rows(cfg::FS.Config, clf::FS.Classifier, jobs::Vector{FS.Job}, opts)
    reps, trials = opts["reps"], opts["trials"]
    nfiles = length(jobs)
    paths = [j.path for j in jobs]
    rows = []

    # A feature vector already in memory, so the inference row measures the net
    # and not the disk read in front of it.
    feats = FS.read_features(paths[1])
    x = reshape(feats, FS.FEATURE_DIM, 1)

    logger = ConsoleLogger(devnull)   # components other than the log rows: quiet

    # --- filesize: the stat() read_features does before touching the bytes
    push!(rows, (; name = "filesize (stat)", part = "classify",
        ns = best_of(noop; trials) do
            @inbounds for i in 1:reps
                SINK[] = filesize(paths[(i - 1) % nfiles + 1])
            end
            reps
        end))

    # --- read_features: open + head read + seek + tail read + scale
    push!(rows, (; name = "read_features", part = "classify",
        ns = best_of(noop; trials) do
            @inbounds for i in 1:reps
                SINK[] = FS.read_features(paths[(i - 1) % nfiles + 1])
            end
            reps
        end))

    # --- Lux.apply: the network alone, features already in memory
    push!(rows, (; name = "Lux.apply (1x32)", part = "classify",
        ns = best_of(noop; trials) do
            for _ in 1:reps
                SINK[] = Lux.apply(clf.model, x, clf.ps, clf.st)
            end
            reps
        end))

    # --- classify: read_features + apply + argmax, what the handler calls
    push!(rows, (; name = "classify (total)", part = "classify",
        ns = best_of(noop; trials) do
            @inbounds for i in 1:reps
                SINK[] = FS.classify(clf, paths[(i - 1) % nfiles + 1])
            end
            reps
        end))

    # --- move_to: the rename out of spool/. Consuming: reset before each pass.
    push!(rows, (; name = "move_to (rename)", part = "route",
        ns = best_of(() -> respool!(cfg, jobs); trials) do
            @inbounds for job in jobs
                SINK[] = FS.move_to(cfg.unknown_dir, job)
            end
            nfiles
        end))

    # --- enqueue: lock, push, notify on an uncontended, non-full queue
    q = FS.ChannelQueue(nfiles + 1)
    stats = FS.StageStats()
    push!(rows, (; name = "enqueue_blocking!", part = "route",
        ns = best_of(() -> drain!(q); trials) do
            @inbounds for job in jobs
                SINK[] = FS.enqueue_blocking!(q, job, stats;
                                              retry_seconds = FS.ROUTE_ENQUEUE_RETRY_SECONDS)
            end
            nfiles
        end))
    drain!(q)

    # --- the two @info lines, under each of the three loggers
    job1 = jobs[1]
    logfile = joinpath(cfg.spool_dir, "..", "bench_stage1.log")
    for (which, label) in ((:null, "logging (NullLogger)"),
                           (:format, "logging (format only)"),
                           (:flush, "logging (flush→file)"))
        ns = with_logger_named(which, logfile) do
            best_of(noop; trials) do
                for _ in 1:reps
                    @info "classified file" worker=1 id=job1.id name=job1.original_name size=job1.size classification=:unknown
                    @info "routed to content triage" worker=1 id=job1.id dest=job1.path
                end
                reps
            end
        end
        push!(rows, (; name = label, part = "log", ns))
    end
    rm(logfile; force = true)

    return rows, logger
end

"""
    handler_rows(cfg, jobs, opts) -> Vector

Time the real `handle_classify_job` end to end under each logger. The difference
between the rows is the cost logging adds to a file; the `:flush` row is what the
running server actually pays.
"""
function handler_rows(cfg::FS.Config, jobs::Vector{FS.Job}, opts)
    trials = opts["trials"]
    nfiles = length(jobs)
    logfile = joinpath(dirname(cfg.spool_dir), "bench_stage1.log")
    known   = FS.ChannelQueue(nfiles + 1)
    unknown = FS.ChannelQueue(nfiles + 1)
    stats   = FS.StageStats()
    rows = []
    for (which, label) in ((:null, "handle_classify_job (NullLogger)"),
                           (:format, "handle_classify_job (format only)"),
                           (:flush, "handle_classify_job (flush→file)"),
                           (:debug, "handle_classify_job (JULIA_DEBUG)"))
        ns = with_logger_named(which, logfile) do
            best_of(() -> (respool!(cfg, jobs); drain!(known); drain!(unknown)); trials) do
                @inbounds for job in jobs
                    FS.handle_classify_job(job, cfg, 1, known, unknown, stats)
                end
                nfiles
            end
        end
        push!(rows, (; name = label, part = "total", ns))
    end
    respool!(cfg, jobs); drain!(known); drain!(unknown)
    rm(logfile; force = true)
    return rows
end

"""
    thread_rows(cfg, jobs, opts) -> Vector

Run the full handler across worker counts, under the server's real logger. A
component that owns a lock — the queue's condition variable, the logger's
stream — stops scaling here even though it looked cheap single-threaded, so this
is where the single-thread ranking gets checked against the deployed one.
"""
function thread_rows(cfg::FS.Config, jobs::Vector{FS.Job}, opts)
    trials = opts["trials"]
    nfiles = length(jobs)
    counts = opts["threads"] === nothing ?
        unique([1; 2; 4; 8; Threads.nthreads()]) :
        [parse(Int, s) for s in split(String(opts["threads"]), ",")]
    counts = sort(unique(filter(k -> 1 <= k <= Threads.nthreads(), counts)))

    logfile = joinpath(dirname(cfg.spool_dir), "bench_stage1.log")
    known   = FS.ChannelQueue(nfiles + 1)
    unknown = FS.ChannelQueue(nfiles + 1)
    stats   = FS.StageStats()
    rows = []
    base = 0.0
    for k in counts
        ns = with_logger_named(:flush, logfile) do
            best_of(() -> (respool!(cfg, jobs); drain!(known); drain!(unknown)); trials) do
                # Static split: each task takes a contiguous slice, so the only
                # sharing between workers is the state the server also shares —
                # the classifier, the queues, the logger, the filesystem.
                chunk = cld(nfiles, k)
                @sync for t in 1:k
                    lo = (t - 1) * chunk + 1
                    hi = min(t * chunk, nfiles)
                    lo > hi && continue
                    Threads.@spawn begin
                        @inbounds for i in lo:hi
                            FS.handle_classify_job(jobs[i], cfg, t, known, unknown, stats)
                        end
                    end
                end
                nfiles
            end
        end
        r = rate(ns) # files/sec aggregate (ns is already per file, wall-clock)
        k == counts[1] && (base = r)
        push!(rows, (; workers = k, ns, files_per_sec = r, speedup = r / base))
    end
    respool!(cfg, jobs); drain!(known); drain!(unknown)
    rm(logfile; force = true)
    return rows
end

# ------------------------------------------------------------------- reporting

function print_components(rows, total_ns)
    @printf("%-34s %-9s %12s %10s %9s\n", "component", "part", "per file", "rate", "% total")
    rule()
    for r in rows
        @printf("%-34s %-9s %12s %10s %8.1f%%\n", r.name, r.part, human_time(r.ns),
                human_rate(rate(r.ns)), 100 * r.ns / total_ns)
    end
end

function print_threads(rows)
    @printf("%-9s %12s %12s %9s\n", "workers", "per file", "throughput", "speedup")
    rule()
    for r in rows
        @printf("%-9d %12s %12s %8.2fx\n", r.workers, human_time(r.ns),
                human_rate(r.files_per_sec), r.speedup)
    end
end

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

function main(argv)
    opts = parse_args(argv)
    modelpath = String(opts["model"])
    isfile(modelpath) || (println(stderr, "model artifact not found: $modelpath"); return 1)

    size = parse_size(String(opts["size"]))
    nfiles, trials = opts["files"], opts["trials"]

    root = opts["dir"] === nothing ?
        mktempdir(pwd(); prefix = "bench_stage1_") : String(opts["dir"])
    owned = opts["dir"] === nothing
    cfg = FS.Config(
        spool_dir   = joinpath(root, "spool"),
        known_dir   = joinpath(root, "known"),
        unknown_dir = joinpath(root, "unknown"),
        failed_dir  = joinpath(root, "failed"),
        model_path  = modelpath,
    )
    for d in (cfg.spool_dir, cfg.known_dir, cfg.unknown_dir, cfg.failed_dir)
        mkpath(d)
    end

    clf = FS.load_classifier(modelpath)
    FS.CLASSIFIER[] = clf     # handle_classify_job reads the global, as in the server

    println("stage-1 component benchmark")
    rule()
    @printf("%-22s %s\n", "julia threads",  Threads.nthreads())
    @printf("%-22s %s\n", "model",          modelpath)
    @printf("%-22s %s\n", "corpus",         "$(nfiles) files x $(size) B in $(root)")
    @printf("%-22s %s\n", "reps / trials",  "$(opts["reps"]) / $(trials)")

    rng = MersenneTwister(0x5741524d)
    jobs = make_corpus(cfg, nfiles, size, rng)

    try
        comps, _ = component_rows(cfg, clf, jobs, opts)
        handlers = handler_rows(cfg, jobs, opts)

        # The denominator is the handler as the server actually runs it: the real
        # flushing logger at Info level, one worker. Percentages are shares of
        # that, so they are directly comparable and the parts can be checked
        # against the whole. The JULIA_DEBUG row is deliberately *not* the
        # baseline — it is the opt-in configuration, and letting it set the scale
        # would make every other component look free.
        total = only(r.ns for r in handlers if r.name == "handle_classify_job (flush→file)")

        header("Components (single worker)")
        print_components(comps, total)

        header("Whole handler, by logger")
        print_components(handlers, total)

        # Stage 1's own log lines are `@debug`, so what the deployed handler pays
        # for them is the disabled-macro cost, not a formatted line.
        accounted = sum(r.ns for r in comps if r.name in
            ("classify (total)", "move_to (rename)", "enqueue_blocking!", "logging (NullLogger)"))
        println()
        @printf("accounted: %s of %s (%.0f%%); unaccounted overhead %s\n",
                human_time(accounted), human_time(total), 100 * accounted / total,
                human_time(max(total - accounted, 0)))

        threads = nothing
        if !opts["no-threads"] && Threads.nthreads() > 1
            threads = thread_rows(cfg, jobs, opts)
            header("Full handler across workers (server logger)")
            print_threads(threads)
        end

        if opts["json"] !== nothing
            open(String(opts["json"]), "w") do io
                JSON3.write(io, (;
                    julia_threads = Threads.nthreads(),
                    file_size = size, files = nfiles, reps = opts["reps"], trials,
                    components = comps, handlers, threads,
                ))
            end
            println("\nwrote ", opts["json"])
        end
    finally
        owned && rm(root; recursive = true, force = true)
    end
    return 0
end

exit(main(ARGS))