#!/usr/bin/env julia
#
# Train the file-type novelty classifier and write the artifact the server loads.
#
#   julia --project=. bin/train.jl <positives_dir> [negatives_dir]
#
# positives_dir  every file in it (>=32 bytes) is a "known" example (class 1).
# negatives_dir  optional grab-bag of *other* real file types used as "unknown"
#                (class 2). If omitted/empty, unknown examples are all random.
#
# Negatives total ~1:1 with positives, split 50/50 between uniform-random byte
# vectors and grab-bag files (falling back to all-random when no grab-bag).
# One fixed seed (FS_TRAIN_SEED, default 42) drives negatives, the train/val
# split, and weight init, so the committed artifact is exactly regenerable.
#
# The model is trained on an 80% split; validation accuracy on the held-out 20%
# is reported so you can see whether it actually learned. The 80%-trained model
# is what gets saved (no retrain-on-all for this test).

using Lux
using JLD2
using Optimisers
using Zygote
using Random
using Statistics

include(joinpath(@__DIR__, "..", "src", "model.jl"))  # build_model(), read_features(), FEATURE_DIM, MIN_FILE_BYTES

const EPOCHS = 200
const LEARNING_RATE = 1f-3
const VAL_FRACTION = 0.20

# --- data assembly -----------------------------------------------------------

"List regular files directly under `dir` that are large enough to featurize."
function eligible_files(dir::AbstractString)
    (isempty(dir) || !isdir(dir)) && return String[]
    paths = String[]
    for name in readdir(dir; join=true)
        isfile(name) && filesize(name) >= MIN_FILE_BYTES && push!(paths, name)
    end
    return paths
end

"32×N Float32 feature matrix from a list of file paths (all assumed eligible)."
function feature_matrix(paths::Vector{String})
    isempty(paths) && return Matrix{Float32}(undef, FEATURE_DIM, 0)
    return reduce(hcat, (read_features(p) for p in paths))
end

# logsoftmax over the class dimension (rows), numerically stabilized. Local so
# the trainer needs no extra dependency for the loss.
function logsoftmax(x::AbstractMatrix)
    m = maximum(x; dims=1)
    shifted = x .- m
    return shifted .- log.(sum(exp.(shifted); dims=1))
end

logitcrossentropy(logits, onehot) = mean(-sum(onehot .* logsoftmax(logits); dims=1))

accuracy(logits, onehot) = mean(vec(map(i -> i[1], argmax(logits; dims=1))) .==
                                vec(map(i -> i[1], argmax(onehot; dims=1))))

function main()
    if isempty(ARGS)
        println(stderr, "usage: julia --project=. bin/train.jl <positives_dir> [negatives_dir]")
        exit(2)
    end
    positives_dir = ARGS[1]
    negatives_dir = length(ARGS) >= 2 ? ARGS[2] : ""
    seed = parse(Int, get(ENV, "FS_TRAIN_SEED", "42"))
    rng = Random.MersenneTwister(seed)

    pos_paths = eligible_files(positives_dir)
    P = length(pos_paths)
    P == 0 && error("no eligible (>= $MIN_FILE_BYTES byte) files found in positives dir: $positives_dir")

    grabbag = eligible_files(negatives_dir)

    # Negatives ~1:1 with positives, 50/50 random vs grab-bag (all-random if the
    # grab-bag is empty). Grab-bag files are sampled with replacement so we can
    # always hit the target count from a small pile.
    n_neg = P
    n_grab = isempty(grabbag) ? 0 : n_neg ÷ 2
    n_rand = n_neg - n_grab

    Xpos = feature_matrix(pos_paths)
    Xrand = rand(rng, Float32, FEATURE_DIM, n_rand)            # uniform [0,1) ≈ random bytes / 255
    grab_paths = n_grab == 0 ? String[] : rand(rng, grabbag, n_grab)
    Xgrab = feature_matrix(grab_paths)

    X = hcat(Xpos, Xrand, Xgrab)
    N = size(X, 2)
    # One-hot targets: known = [1,0], unknown = [0,1].
    Y = zeros(Float32, 2, N)
    Y[1, 1:P] .= 1f0            # positives -> known
    Y[2, P+1:end] .= 1f0        # all negatives -> unknown

    @info "assembled dataset" positives=P negatives=n_neg random_neg=n_rand grabbag_neg=n_grab total=N grabbag_pool=length(grabbag)

    # Seeded shuffle + 80/20 split.
    perm = randperm(rng, N)
    X, Y = X[:, perm], Y[:, perm]
    n_val = round(Int, VAL_FRACTION * N)
    n_train = N - n_val
    Xtr, Ytr = X[:, 1:n_train], Y[:, 1:n_train]
    Xval, Yval = X[:, n_train+1:end], Y[:, n_train+1:end]

    # --- train (full-batch, fixed epochs) ------------------------------------
    model = build_model()
    ps, st = Lux.setup(rng, model)
    opt_state = Optimisers.setup(Optimisers.Adam(LEARNING_RATE), ps)

    for epoch in 1:EPOCHS
        (loss, st), back = Zygote.pullback(ps) do p
            logits, st_new = model(Xtr, p, st)
            logitcrossentropy(logits, Ytr), st_new
        end
        grads = back((one(loss), nothing))[1]
        opt_state, ps = Optimisers.update(opt_state, ps, grads)
        if epoch == 1 || epoch % 20 == 0 || epoch == EPOCHS
            tr_logits, _ = model(Xtr, ps, st)
            @info "epoch" epoch loss=loss train_acc=accuracy(tr_logits, Ytr)
        end
    end

    # --- report + save -------------------------------------------------------
    val_logits, _ = model(Xval, ps, st)
    @info "validation" n_val=n_val val_acc=(n_val == 0 ? NaN : accuracy(val_logits, Yval))

    out = get(ENV, "FS_MODEL_PATH", "model/classifier.jld2")
    mkpath(dirname(out))
    jldsave(out; ps=ps, st=st)
    @info "saved model artifact" path=out
end

main()