// vdm/task/download_task.cpp — the download state machine behind DownloadHandle.
//
// One std::mutex `mu` is the task lock: every state transition, worker start/stop, and
// terminal path runs under it. The curl write-path (seg_data / seg_head) touches only
// per-segment state, guarded by a shared_mutex over the worker map. Segment completions
// and timers are posted to the engine timer thread and run under `mu`. Callbacks to
// DAEMON are collected under `mu` (defer()) and fired only after it is released
// (flush_deferred(), which guards its own queue with a separate mutex).

#include "task/download_task.hpp"

#include <unistd.h>

#include <algorithm>
#include <atomic>
#include <cctype>
#include <cerrno>
#include <condition_variable>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <shared_mutex>
#include <unordered_map>

#include "task/digest.hpp"
#include "vdm/io/sparse_file.hpp"
#include "vdm/io/write_buffer.hpp"
#include "vdm/meta/veloxpart.hpp"
#include "vdm/net/url.hpp"
#include "vdm/segment/segmenter.hpp"

namespace vdm::task {
namespace {

constexpr std::uint64_t kBufFloor = 64u * 1024;
constexpr std::uint64_t kBufCeil = 16u * 1024 * 1024;
constexpr std::uint64_t kNoEnd = ~std::uint64_t{0};

bool is_connection_error(Error e) noexcept {
    switch (e) {
        case Error::connection_reset:
        case Error::timeout:
        case Error::connect_failed:
        case Error::resolve_failed:
        case Error::tls_failed:
            return true;
        default:
            return false;
    }
}

std::chrono::milliseconds backoff_for(int attempt) {
    // 1st retry ~1s, then 2, 4, 8, ... capped at 60s, +/-20% jitter (docs/04 §7).
    int shift = std::clamp(attempt - 1, 0, 6);
    std::uint64_t base = 1000ull << shift;
    base = std::min<std::uint64_t>(base, 60'000);
    std::uint64_t jitter = base / 5;
    std::uint64_t r = static_cast<std::uint64_t>(std::rand()) % (2 * jitter + 1);
    return std::chrono::milliseconds{base - jitter + r};
}

std::string lower(std::string s) {
    std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return std::tolower(c); });
    return s;
}

}  // namespace

// What to do once every worker has drained (see DownloadTaskState::begin_drain_locked).
// A worker that hits one of these outcomes must not tear the task down itself: siblings may
// still be mid-transfer, and cutting them off synchronously — cancel every worker, clear the
// map, proceed — drops their buffered-but-unflushed bytes while segment_completed() already
// counts those bytes as done. That's exactly the class of bug that makes a later resume skip
// real data (seg_data() runs append() without workers_mu, so a sibling's buffer can't be
// flushed safely from here anyway). Instead: cancel the siblings, remember what to do, and
// let each one's own seg_finished (which already flushes on every exit path) run it once the
// worker map is actually empty.
enum class PendingAction { none, verify, fail, auto_pause, demote };

struct SegWorker {
    std::uint32_t seg_index = 0;
    net::Transfer transfer;
    std::unique_ptr<io::WriteBuffer> buf;
    std::uint64_t base_completed = 0;
    std::uint64_t recv = 0;

    long http_status = 0;
    bool needs_auth = false;
    bool wrong_status = false;
    bool range_bad = false;
    bool auth_handshake = false;  // saw a 401/407 and let libcurl resend with credentials
    std::string resp_etag, resp_last_modified;  // captured on a wrong_status 200, for demote
    std::optional<ErrorInfo> flush_error;

    int retries = 0;

    SteadyTime sample_at{};
    std::uint64_t sample_bytes = 0;
    double speed_bps = 0;
};

struct DownloadTaskState : std::enable_shared_from_this<DownloadTaskState> {
    TaskHost &host;
    TaskId id;
    DownloadSpec spec;
    DownloadCallbacks cbs;

    std::mutex mu;
    std::shared_mutex workers_mu;
    std::mutex deferred_mu;
    // Notified (holding mu) whenever seg_finished removes an entry from `workers`; quiesce()
    // waits on it instead of clearing the map itself — see quiesce()'s comment.
    std::condition_variable workers_drained_cv;

    EngineState state = EngineState::probing;
    std::optional<ErrorInfo> last_error;
    std::atomic<bool> retired{false};

    std::string part_path, meta_path;

    net::ProbeResult probe;
    bool have_probe = false;
    bool resumable = false;
    bool probe_needed_auth = false;
    std::optional<std::uint64_t> total_size;
    std::string origin_host;

    std::unique_ptr<segment::Segmenter> seg;
    std::unique_ptr<io::SparseFile> file;
    std::unordered_map<std::uint32_t, std::unique_ptr<SegWorker>> workers;
    std::unordered_map<std::uint32_t, int> retry_counts;
    std::uint32_t slot_target = 0;
    std::uint32_t effective_buffer = 0;
    std::uint32_t requested_segments = 8;
    bool registered = false;

    bool pause_requested = false;
    bool cancel_requested = false;
    bool discard_on_cancel = false;
    bool awaiting_auth = false;
    bool awaiting_decision = false;
    int max_retries = 10;

    // Set by begin_drain_locked while waiting for sibling workers to drain; see
    // PendingAction above.
    PendingAction pending_action = PendingAction::none;
    std::optional<ErrorInfo> pending_error;
    bool pending_auth = false;
    bool pending_decision = false;

    std::atomic<std::int64_t> last_progress_ns{0};
    SteadyTime started_at{};

    std::vector<std::function<void()>> deferred;

    DownloadTaskState(TaskHost &h, TaskId i, DownloadSpec s, DownloadCallbacks c)
        : host(h), id(i), spec(std::move(s)), cbs(std::move(c)) {}

    // --- deferred callbacks -------------------------------------------------------------
    void defer(std::function<void()> fn) {
        std::lock_guard lk(deferred_mu);
        deferred.push_back(std::move(fn));
    }
    void flush_deferred() {  // MUST be called with `mu` NOT held
        std::vector<std::function<void()>> run;
        {
            std::lock_guard lk(deferred_mu);
            run.swap(deferred);
        }
        for (auto &fn : run)
            fn();
    }

    // caller holds mu
    void transition(EngineState to, std::optional<ErrorInfo> err) {
        if (state == to && !err)
            return;
        EngineState from = state;
        state = to;
        last_error = err;
        auto cb = cbs.on_state;
        defer([cb, from, to, err] {
            if (cb)
                cb(from, to, err);
        });
    }
    [[nodiscard]] std::string current_url() const {
        return have_probe && !probe.effective_url.empty() ? probe.effective_url : spec.url;
    }
    [[nodiscard]] bool has_mirror() const { return !spec.mirrors.empty(); }
    [[nodiscard]] std::uint32_t want_slots() const {
        if (!seg)
            return requested_segments ? requested_segments : 1;
        std::uint32_t target = seg->target_segment_count();
        std::uint32_t incomplete = 0;
        bool any = false;
        for (auto &v : seg->snapshot()) {
            any = true;
            if (v.state != segment::SegState::complete && v.state != segment::SegState::failed)
                ++incomplete;
        }
        if (!any)
            return target;  // fresh: nothing created yet
        if (incomplete == 0)
            return target;  // orphans pending re-split
        return std::clamp<std::uint32_t>(incomplete, 1, target);
    }

    void begin();
    void on_probe_result(Result<net::ProbeResult> r);
    void finish_probe_locked();
    void apply_slot_target(std::uint32_t n);
    void fill_slots_locked();
    void start_worker_locked(std::uint32_t seg_idx);
    void restart_probe(bool with_auth);

    net::DataAction seg_head(std::uint32_t seg_idx, const net::ResponseHead &h);
    net::DataAction seg_data(std::uint32_t seg_idx, ConstByteSpan span);
    void on_transfer_done(std::uint32_t seg_idx, Result<net::TransferStats> r);
    void seg_finished(std::uint32_t seg_idx, Result<net::TransferStats> r);
    void retry_worker(std::uint32_t seg_idx);

    void begin_verify_locked();
    void fail_locked(ErrorInfo e);
    void auto_pause_locked(ErrorInfo e, bool auth, bool decision);
    void demote_to_single_segment_locked();
    void begin_drain_locked(PendingAction action, ErrorInfo e, bool auth, bool decision);
    void finalize_cancel_locked();
    void write_sidecar_locked();
    void emit_progress_if_due();

    void do_pause();
    void do_resume();
    void do_cancel(bool discard);
    void do_provide_auth(std::string u, std::string p, bool remember);
    void do_decide(Decision d);
    void do_refresh_url(std::string url, std::vector<net::HeaderField> headers);
    EngineState snapshot_state();
    Progress snapshot_progress();
    void quiesce();  // engine shutdown: stop everything, fire nothing
};

// ==================================================================================

void DownloadTaskState::begin() {
    bool probe_hint;
    {
        std::unique_lock lk(mu);
        part_path = spec.save_path + ".veloxpart";
        meta_path = part_path + ".meta";
        started_at = std::chrono::steady_clock::now();
        max_retries =
            static_cast<int>(spec.max_retries.value_or(host.config().default_max_retries));
        requested_segments = spec.segments.value_or(host.config().default_segments);
        probe_hint = spec.probe_hint.has_value();
        if (probe_hint) {
            probe = *spec.probe_hint;
            have_probe = true;
            finish_probe_locked();
        } else {
            transition(EngineState::probing, std::nullopt);
        }
    }
    flush_deferred();
    if (!probe_hint)
        restart_probe(false);
}

void DownloadTaskState::restart_probe(bool with_auth) {
    net::ProbeRequest pr;
    pr.url = spec.url;
    pr.headers = spec.headers;
    pr.cookies = spec.cookies;
    pr.referrer = spec.referrer;
    pr.user_agent = spec.user_agent;
    pr.proxy = spec.proxy;
    if (with_auth)
        pr.auth = spec.auth;
    auto wp = weak_from_this();
    host.probe(std::move(pr), [wp](Result<net::ProbeResult> r) {
        if (auto s = wp.lock())
            s->on_probe_result(std::move(r));
    });
}

void DownloadTaskState::on_probe_result(Result<net::ProbeResult> r) {
    {
        std::unique_lock lk(mu);
        if (retired.load() || is_terminal(state))
            return;
        if (!r.has_value()) {
            fail_locked(std::move(r).error());
        } else {
            probe = std::move(r).value();
            have_probe = true;
            if (probe.requires_auth) {
                probe_needed_auth = true;
                auto_pause_locked(
                    ErrorInfo(Error::auth_required, "probe 401/407", probe.http_status), true,
                    false);
            } else {
                finish_probe_locked();
            }
        }
    }
    flush_deferred();
}

void DownloadTaskState::finish_probe_locked() {
    resumable = probe.resumable;
    total_size = probe.total_size;
    origin_host = net::split_url(current_url()).host;

    const std::uint64_t min_seg = host.config().min_segment_bytes;
    const std::uint64_t total = total_size.value_or(0);

    std::vector<segment::ResumedRange> resumed;
    bool do_resume = spec.allow_resume;
    if (do_resume) {
        auto m = meta::read_veloxpart_file(meta_path);
        if (m.has_value() && total > 0 && m.value().total_size == total &&
            (m.value().etag.empty() || m.value().etag == probe.etag) &&
            (m.value().last_modified.empty() || m.value().last_modified == probe.last_modified)) {
            for (auto &s : m.value().segments)
                resumed.push_back({s.start, s.end, s.completed});
        } else {
            do_resume = false;
        }
    }

    file = std::make_unique<io::SparseFile>();
    io::SparseFile::OpenOptions oo;
    oo.total_size = total;
    oo.preallocate = total > 0;
    oo.truncate_existing = !do_resume;
    if (auto o = file->open(part_path, oo); !o.has_value()) {
        fail_locked(std::move(o).error());
        return;
    }

    if (do_resume && !resumed.empty())
        seg = std::make_unique<segment::Segmenter>(total, requested_segments, resumed, resumable,
                                                   min_seg);
    else
        seg = std::make_unique<segment::Segmenter>(total, requested_segments, resumable, min_seg);

    const std::uint32_t cap = resumable ? std::min<std::uint32_t>(requested_segments, 32) : 1;
    std::uint64_t want_buf = spec.buffer_bytes.value_or(host.config().default_buffer_bytes);
    want_buf = std::clamp<std::uint64_t>(want_buf, kBufFloor, kBufCeil);
    std::uint64_t max_per = std::max<std::uint64_t>(
        kBufFloor, host.config().max_total_buffer_bytes / std::max(1u, cap));
    effective_buffer = static_cast<std::uint32_t>(std::min(want_buf, max_per));

    auto wp = weak_from_this();
    host.budget().register_task(id, {origin_host, cap, resumable}, [wp](std::uint32_t n) {
        auto s = wp.lock();
        if (!s)
            return;
        s->host.schedule(std::chrono::steady_clock::now(), [wp, n] {
            if (auto s2 = wp.lock())
                s2->apply_slot_target(n);
        });
    });
    registered = true;
    host.limiter().attach_task(id, std::nullopt);

    transition(EngineState::connecting, std::nullopt);
    host.budget().set_want(id, want_slots());
}

// Start workers up to `slot_target`, given whatever the budget currently confirms. Shared
// by apply_slot_target() (the budget's async callback, whenever the computed target
// actually changes) and demote_to_single_segment_locked() -- the demoted target can
// legitimately equal what it was before the rebuild (e.g. a download already down to its
// last segment), in which case SegmentBudget::set_want() no-ops and the async callback
// never fires, so nothing else would ever start the replacement worker.
void DownloadTaskState::fill_slots_locked() {
    while (workers.size() < slot_target) {
        auto s = seg->assign_slot();
        if (!s) {
            host.budget().set_want(id, static_cast<std::uint32_t>(workers.size()));
            break;
        }
        if (!host.budget().confirm_slot(id)) {
            seg->set_segment_state(*s, segment::SegState::idle);
            break;
        }
        start_worker_locked(*s);
    }
    if (state == EngineState::connecting && !workers.empty())
        transition(EngineState::downloading, std::nullopt);
}

void DownloadTaskState::apply_slot_target(std::uint32_t n) {
    {
        std::unique_lock lk(mu);
        if (retired.load() || is_terminal(state) || pause_requested || cancel_requested ||
            awaiting_auth || awaiting_decision || pending_action != PendingAction::none || !seg)
            return;
        slot_target = n;
        fill_slots_locked();
    }
    flush_deferred();
}

void DownloadTaskState::start_worker_locked(std::uint32_t seg_idx) {
    auto w = std::make_unique<SegWorker>();
    w->seg_index = seg_idx;
    const std::uint64_t sstart = seg->segment_start(seg_idx);
    const std::uint64_t completed = seg->segment_completed(seg_idx);
    const std::uint64_t send = seg->segment_end(seg_idx);
    w->base_completed = completed;
    w->sample_at = std::chrono::steady_clock::now();
    if (auto it = retry_counts.find(seg_idx); it != retry_counts.end())
        w->retries = it->second;

    auto wp = weak_from_this();
    w->buf = std::make_unique<io::WriteBuffer>(
        sstart + completed, effective_buffer,
        [wp](std::uint64_t off, ConstByteSpan sp) -> Result<void> {
            if (auto s = wp.lock())
                return s->file->write_at(off, sp);
            return ErrorInfo(Error::canceled, "task gone");
        });

    net::Request req;
    req.url = current_url();
    req.headers = spec.headers;
    req.cookies = spec.cookies;
    req.referrer = spec.referrer;
    req.user_agent = spec.user_agent;
    req.proxy = spec.proxy;
    req.auth = spec.auth;
    req.follow_redirects = true;
    req.accept_encoding = false;
    req.low_speed_bytes_per_sec = 1024;
    req.low_speed_secs = 30;
    // Only range-request a resumable source; a non-resumable single segment is a plain GET.
    if (resumable && total_size && send != kNoEnd)
        req.range = net::ByteRange{sstart + completed, send};
    if (completed > 0 && (!probe.etag.empty() || !probe.last_modified.empty()))
        req.headers.push_back({"If-Range", !probe.etag.empty() ? probe.etag : probe.last_modified});

    net::TransferCallbacks tc;
    tc.on_head = [wp, seg_idx](const net::ResponseHead &h) {
        if (auto s = wp.lock())
            return s->seg_head(seg_idx, h);
        return net::DataAction::abort;
    };
    tc.on_data = [wp, seg_idx](ConstByteSpan sp) {
        if (auto s = wp.lock())
            return s->seg_data(seg_idx, sp);
        return net::DataAction::abort;
    };
    tc.on_finished = [wp, seg_idx](Result<net::TransferStats> r) {
        if (auto s = wp.lock())
            s->on_transfer_done(seg_idx, std::move(r));
    };

    auto *raw = w.get();
    {
        std::unique_lock wl(workers_mu);
        workers[seg_idx] = std::move(w);
    }
    raw->transfer = host.http().start(std::move(req), std::move(tc));
    seg->set_segment_state(seg_idx, segment::SegState::connecting);
}

// --- curl write path -------------------------------------------------------------------

net::DataAction DownloadTaskState::seg_head(std::uint32_t seg_idx, const net::ResponseHead &h) {
    SegWorker *w = nullptr;
    {
        std::shared_lock lk(workers_mu);
        auto it = workers.find(seg_idx);
        if (it == workers.end())
            return net::DataAction::abort;
        w = it->second.get();
    }
    w->http_status = h.status;
    if (h.status == 401 || h.status == 407) {
        // libcurl with CURLAUTH_ANY does a challenge round-trip: the first response is a
        // 401/407, then it resends with credentials. Don't kill that handshake — let it
        // proceed once and wait for the real status on the next header block. Only when we
        // have nothing to try, or we already tried, is this a genuine "auth required".
        const bool have_creds =
            spec.auth.scheme != net::AuthScheme::none && !spec.auth.username.empty();
        if (have_creds && !w->auth_handshake) {
            w->auth_handshake = true;
            return net::DataAction::proceed;
        }
        w->needs_auth = true;
        return net::DataAction::abort;
    }
    // A 200 where we sent a Range is only "the file changed under us" when we are running
    // a resumable, multi-segment transfer. A non-resumable single-segment GET legitimately
    // gets a 200 (the source has no Range support).
    if (resumable && total_size && *total_size > 0 && h.status == 200) {
        w->wrong_status = true;
        if (auto v = h.headers.get("ETag"))
            w->resp_etag.assign(*v);
        if (auto v = h.headers.get("Last-Modified"))
            w->resp_last_modified.assign(*v);
        return net::DataAction::abort;
    }
    if (h.status == 416) {
        w->range_bad = true;
        return net::DataAction::abort;
    }
    if (h.status >= 400)
        return net::DataAction::abort;
    seg->set_segment_state(seg_idx, segment::SegState::downloading);
    return net::DataAction::proceed;
}

net::DataAction DownloadTaskState::seg_data(std::uint32_t seg_idx, ConstByteSpan span) {
    SegWorker *w = nullptr;
    {
        std::shared_lock lk(workers_mu);
        auto it = workers.find(seg_idx);
        if (it == workers.end())
            return net::DataAction::abort;
        w = it->second.get();
    }
    if (span.empty())
        return net::DataAction::proceed;
    // Body of libcurl's pre-auth 401/407 response — discard it; the real body follows the
    // resend under a 2xx header.
    if (w->http_status == 401 || w->http_status == 407)
        return net::DataAction::proceed;

    auto wait = host.limiter().acquire(id, span.size());
    if (wait.count() > 0) {
        net::Transfer t = w->transfer;
        host.schedule(std::chrono::steady_clock::now() + wait, [t]() mutable { t.resume(); });
        return net::DataAction::pause;
    }

    if (auto r = w->buf->append(span); !r.has_value()) {
        w->flush_error = std::move(r).error();
        return net::DataAction::abort;
    }
    w->recv += span.size();
    seg->advance(seg_idx, w->base_completed + w->recv);

    auto now = std::chrono::steady_clock::now();
    auto dt = std::chrono::duration<double>(now - w->sample_at).count();
    if (dt >= 0.5) {
        double inst = static_cast<double>(w->recv - w->sample_bytes) / dt;
        w->speed_bps = w->speed_bps == 0 ? inst : 0.7 * w->speed_bps + 0.3 * inst;
        w->sample_at = now;
        w->sample_bytes = w->recv;
    }
    emit_progress_if_due();
    return net::DataAction::proceed;
}

void DownloadTaskState::on_transfer_done(std::uint32_t seg_idx, Result<net::TransferStats> r) {
    auto wp = weak_from_this();
    host.schedule(std::chrono::steady_clock::now(), [wp, seg_idx, r = std::move(r)]() mutable {
        if (auto s = wp.lock())
            s->seg_finished(seg_idx, std::move(r));
    });
}

void DownloadTaskState::seg_finished(std::uint32_t seg_idx, Result<net::TransferStats> r) {
    std::unique_lock lk(mu);

    std::unique_ptr<SegWorker> w;
    {
        std::unique_lock wl(workers_mu);
        auto it = workers.find(seg_idx);
        if (it == workers.end()) {
            lk.unlock();
            flush_deferred();
            return;
        }
        w = std::move(it->second);
        workers.erase(it);
    }
    workers_drained_cv.notify_all();  // quiesce() may be waiting for `workers` to empty out
    if (retired.load()) {  // engine shutting down / already terminal — no more callbacks
        if (w->buf)
            (void)w->buf->flush();
        lk.unlock();
        return;
    }

    // The bytes may all have arrived even though the connection closed dirty
    // (content-length-mismatch's honest-length lie, flaky-reset's tail, a proxy RST after
    // the last byte). If the segment is fully covered, that's a success.
    if (seg && !cancel_requested && !pause_requested && !w->needs_auth && !w->wrong_status &&
        !w->flush_error && !w->range_bad) {
        const std::uint64_t len = seg->segment_end(seg_idx) - seg->segment_start(seg_idx) + 1;
        if (len != 0 && seg->segment_completed(seg_idx) >= len) {
            r = Result<net::TransferStats>(net::TransferStats{});
        }
    }
    auto release_slot = [&] {
        if (registered)
            host.budget().release_slot(id);
    };
    auto done = [&] {
        lk.unlock();
        flush_deferred();
    };

    if (cancel_requested) {
        if (w->buf)
            (void)w->buf->flush();
        release_slot();
        if (workers.empty())
            finalize_cancel_locked();
        return done();
    }
    if (pause_requested) {
        if (w->buf)
            (void)w->buf->flush();
        seg->advance(seg_idx, w->base_completed + w->recv);
        // Hand the range back unassigned so resume's assign_slot() picks it up instead of
        // treating it as still-held and splitting it (which would orphan its front half).
        seg->release_segment(seg_idx);
        release_slot();
        if (workers.empty()) {
            (void)file->sync();
            write_sidecar_locked();
            transition(EngineState::paused, std::nullopt);
        }
        return done();
    }
    if (pending_action != PendingAction::none) {
        // A sibling already decided the task is finishing (verify / fail / auto-pause /
        // demote); this worker's own outcome no longer matters. Drain it like every other
        // exit path: flush its buffer so segment_completed() stays true to disk, then hand
        // off to whichever worker finds the map empty.
        if (w->buf)
            (void)w->buf->flush();  // best-effort: we're already tearing down for another
                                     // reason, and the pending action doesn't depend on
                                     // this segment reaching any particular state.
        seg->advance(seg_idx, w->base_completed + w->recv);
        release_slot();
        if (workers.empty()) {
            PendingAction action = std::exchange(pending_action, PendingAction::none);
            ErrorInfo e = pending_error.value_or(ErrorInfo(Error::internal, ""));
            bool auth = pending_auth, decision = pending_decision;
            switch (action) {
                case PendingAction::verify:
                    begin_verify_locked();
                    break;
                case PendingAction::fail:
                    fail_locked(e);
                    break;
                case PendingAction::auto_pause:
                    auto_pause_locked(e, auth, decision);
                    break;
                case PendingAction::demote:
                    demote_to_single_segment_locked();
                    break;
                case PendingAction::none:
                    break;
            }
        }
        return done();
    }
    if (w->needs_auth) {
        release_slot();
        auto_pause_locked(ErrorInfo(Error::auth_required, "401/407", w->http_status), true, false);
        return done();
    }
    if (w->wrong_status) {
        release_slot();
        // A 200 where 206 was expected is ambiguous: the file really changed (ask, don't
        // corrupt — docs/04 §5), or the server just stopped honouring Range for this
        // connection while it's still the same file (docs/04 §7: demote to 1 segment and
        // continue). ETag/Last-Modified from the 200 itself, compared against what the
        // probe recorded, is the only signal that tells them apart. Prefer ETag strictly
        // when both sides have one — same rule If-Range itself uses — and only fall back to
        // Last-Modified when there's no ETag to compare; a coarse (often second-resolution)
        // Last-Modified that happens to match is weak evidence next to a mismatching ETag.
        bool same_file;
        if (!probe.etag.empty() && !w->resp_etag.empty())
            same_file = probe.etag == w->resp_etag;
        else if (!probe.last_modified.empty() && !w->resp_last_modified.empty())
            same_file = probe.last_modified == w->resp_last_modified;
        else
            same_file = false;  // no validator to compare -> can't prove it, ask
        if (same_file) {
            demote_to_single_segment_locked();
        } else {
            auto_pause_locked(ErrorInfo(Error::server_file_changed, "200 where 206 expected"),
                              false, true);
        }
        return done();
    }
    if (w->flush_error) {
        ErrorInfo e = *w->flush_error;
        release_slot();
        if (e.code == Error::disk_full) {
            auto_pause_locked(e, false, false);
        } else {
            fail_locked(e);
        }
        return done();
    }
    if (w->range_bad) {
        // 416 mid-download means our range metadata is stale (docs/04 §7): re-probe and
        // re-split rather than retrying the same now-invalid range until it exhausts.
        release_slot();
        auto_pause_locked(ErrorInfo(Error::range_not_satisfiable, "416"), false, true);
        return done();
    }

    if (!r.has_value()) {
        ErrorInfo e = std::move(r).error();
        // Persist what this attempt received before the connection dropped. seg_data() has
        // already advanced `completed` past these bytes, but they are still only in the
        // WriteBuffer (nothing forces a flush until the buffer fills or the segment ends).
        // Without this, a retry resumes from `completed` and skips an unwritten hole — for
        // a small file that never fills the buffer, every reset loses its whole payload.
        if (w->buf) {
            if (auto fl = w->buf->flush(); !fl.has_value()) {
                release_slot();
                fail_locked(std::move(fl).error());
                return done();
            }
        }
        const bool conn = is_connection_error(e.code);
        segment::FailAction fa = seg->on_failed(seg_idx, conn, has_mirror());
        if (w->recv > 0)
            w->retries = 0;  // progress this attempt: not a stuck segment
        if (fa == segment::FailAction::requeue) {
            retry_counts.erase(seg_idx);
            release_slot();
            host.budget().set_want(id, want_slots());
        } else if (++w->retries > max_retries) {
            ErrorInfo x(Error::max_retries_exhausted, e.to_string());
            x.cause = e.code;
            release_slot();
            fail_locked(x);
        } else {
            retry_counts[seg_idx] = w->retries;
            seg->set_segment_state(seg_idx, segment::SegState::stalled);
            release_slot();  // give the slot back while we back off
            int attempt = w->retries;
            auto wp = weak_from_this();
            host.schedule(std::chrono::steady_clock::now() + backoff_for(attempt), [wp, seg_idx] {
                if (auto s = wp.lock())
                    s->retry_worker(seg_idx);
            });
            if (workers.empty())
                transition(EngineState::retry_wait, std::nullopt);
        }
        return done();
    }

    // success
    if (w->buf) {
        if (auto f = w->buf->flush(); !f.has_value()) {
            release_slot();
            fail_locked(std::move(f).error());
            return done();
        }
    }
    retry_counts.erase(seg_idx);
    const bool may_steal = (workers.size() + 1) <= slot_target && !pause_requested;
    auto cont = seg->on_complete(seg_idx, may_steal);
    (void)file->sync();
    write_sidecar_locked();

    if (cont)
        start_worker_locked(*cont);  // slot-neutral steal
    else
        release_slot();

    if (seg->all_complete())
        begin_verify_locked();  // drain-aware: defers if other workers are still live
    return done();
}

void DownloadTaskState::retry_worker(std::uint32_t seg_idx) {
    {
        std::unique_lock lk(mu);
        if (retired.load() || is_terminal(state) || pause_requested || cancel_requested ||
            pending_action != PendingAction::none || !seg)
            return;
        if (workers.count(seg_idx))
            return;
        if (!host.budget().confirm_slot(id))
            return;  // budget is full; try again later
        if (state == EngineState::retry_wait)
            transition(EngineState::connecting, std::nullopt);
        start_worker_locked(seg_idx);
        if (!workers.empty() && state == EngineState::connecting)
            transition(EngineState::downloading, std::nullopt);
    }
    flush_deferred();
}

void DownloadTaskState::begin_verify_locked() {
    if (!workers.empty()) {
        begin_drain_locked(PendingAction::verify, ErrorInfo(Error::internal, ""), false, false);
        return;
    }
    transition(EngineState::assembling, std::nullopt);
    transition(EngineState::verifying, std::nullopt);
    (void)file->sync();
    (void)file->close();

    std::optional<std::string> got;
    if (spec.checksum) {
        auto h = hash_file(part_path, spec.checksum->algo);
        if (!h.has_value()) {
            fail_locked(std::move(h).error());
            return;
        }
        got = h.value();
        if (*got != lower(spec.checksum->hex)) {
            fail_locked(ErrorInfo(Error::checksum_mismatch,
                                  "want " + lower(spec.checksum->hex) + " got " + *got));
            return;
        }
    }
    if (::rename(part_path.c_str(), spec.save_path.c_str()) != 0) {
        fail_locked(ErrorInfo(Error::io_error, std::string("rename: ") + std::strerror(errno)));
        return;
    }
    ::unlink(meta_path.c_str());
    if (registered)
        host.budget().deregister_task(id);
    host.limiter().detach_task(id);
    registered = false;
    transition(EngineState::complete, std::nullopt);

    DownloadOutcome o;
    o.final_path = spec.save_path;
    o.bytes = seg->downloaded();
    o.sha256_hex =
        (spec.checksum && spec.checksum->algo == Checksum::Algo::sha256) ? got : std::nullopt;
    o.elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
        std::chrono::steady_clock::now() - started_at);
    auto cb = cbs.on_finished;
    defer([cb, o] {
        if (cb)
            cb(Result<DownloadOutcome>(o));
    });
    retired.store(true);
    TaskHost *h = &host;
    TaskId tid = id;
    defer([h, tid] { h->task_retired(tid); });
}

void DownloadTaskState::fail_locked(ErrorInfo e) {
    if (!workers.empty()) {
        begin_drain_locked(PendingAction::fail, std::move(e), false, false);
        return;
    }
    if (file)
        (void)file->close();
    if (seg)
        write_sidecar_locked();
    if (registered)
        host.budget().deregister_task(id);
    host.limiter().detach_task(id);
    registered = false;
    transition(EngineState::failed, e);
    auto cb = cbs.on_finished;
    defer([cb, e] {
        if (cb)
            cb(Result<DownloadOutcome>(ErrorInfo(e)));
    });
    retired.store(true);
    TaskHost *h = &host;
    TaskId tid = id;
    defer([h, tid] { h->task_retired(tid); });
}

void DownloadTaskState::auto_pause_locked(ErrorInfo e, bool auth, bool decision) {
    if (!workers.empty()) {
        begin_drain_locked(PendingAction::auto_pause, std::move(e), auth, decision);
        return;
    }
    awaiting_auth = auth;
    awaiting_decision = decision;
    if (file)
        (void)file->sync();
    if (seg)
        write_sidecar_locked();
    if (registered)
        host.budget().set_want(id, 0);
    transition(EngineState::paused, e);
    if (auth) {
        AuthChallenge ac;
        ac.host = origin_host;
        auto cb = cbs.on_auth_required;
        defer([cb, ac] {
            if (cb)
                cb(ac);
        });
    }
    if (decision) {
        DecisionRequest dr;
        dr.kind = e.code == Error::server_file_changed
                      ? DecisionRequest::Kind::server_file_changed
                      : DecisionRequest::Kind::range_metadata_stale;
        dr.detail = e.context;
        auto cb = cbs.on_decision_needed;
        defer([cb, dr] {
            if (cb)
                cb(dr);
        });
    }
}

// Cancel every live worker and remember what to do once they've all drained through
// seg_finished's PendingAction branch (see the enum's comment). Never clears `workers`
// itself — each worker removes itself, flushed, when its own transfer actually completes.
void DownloadTaskState::begin_drain_locked(PendingAction action, ErrorInfo e, bool auth,
                                           bool decision) {
    pending_action = action;
    pending_error = std::move(e);
    pending_auth = auth;
    pending_decision = decision;
    std::shared_lock wl(workers_mu);
    for (auto &[idx, w] : workers)
        w->transfer.cancel();
}

// The 200-where-206-expected we just saw carried the same ETag/Last-Modified the probe
// recorded: same file, the server (or this connection) just doesn't honour Range. Rebuild
// as a single non-resumable segment covering the whole file and keep going with a plain
// GET. It re-transfers bytes we may already have — there's no way to ask a Range-blind
// server for a suffix — but it never truncates or discards what's on disk, and a source
// that hasn't changed serves identical bytes, so the result is still byte-correct.
void DownloadTaskState::demote_to_single_segment_locked() {
    if (!workers.empty()) {
        begin_drain_locked(PendingAction::demote, ErrorInfo(Error::internal, ""), false, false);
        return;
    }
    resumable = false;
    seg = std::make_unique<segment::Segmenter>(total_size.value_or(0), 1, /*resumable=*/false,
                                               host.config().min_segment_bytes);
    transition(EngineState::connecting, std::nullopt);
    if (registered) {
        // set_want() alone is not enough: if the demoted target happens to equal what it
        // was before the rebuild (e.g. this was already the last live segment), it's a
        // no-op and the async budget callback never fires. Drive slot assignment directly.
        slot_target = want_slots();
        host.budget().set_want(id, slot_target);
        fill_slots_locked();
    }
}

void DownloadTaskState::finalize_cancel_locked() {
    if (file)
        (void)file->close();
    if (registered)
        host.budget().deregister_task(id);
    host.limiter().detach_task(id);
    registered = false;
    if (discard_on_cancel) {
        ::unlink(part_path.c_str());
        ::unlink(meta_path.c_str());
    } else if (seg) {
        write_sidecar_locked();
    }
    transition(EngineState::cancelled, std::nullopt);
    auto cb = cbs.on_finished;
    defer([cb] {
        if (cb)
            cb(Result<DownloadOutcome>(ErrorInfo(Error::canceled)));
    });
    retired.store(true);
    TaskHost *h = &host;
    TaskId tid = id;
    defer([h, tid] { h->task_retired(tid); });
}

void DownloadTaskState::write_sidecar_locked() {
    if (!seg || !total_size)
        return;
    meta::VeloxPart vp;
    vp.total_size = *total_size;
    vp.urls = {spec.url};
    if (!current_url().empty() && current_url() != spec.url)
        vp.urls.push_back(current_url());
    for (auto &m : spec.mirrors)
        vp.urls.push_back(m);
    vp.etag = probe.etag;
    vp.last_modified = probe.last_modified;
    vp.content_type = probe.mime;
    std::uint64_t dl = 0;
    for (auto &v : seg->snapshot()) {
        if (v.state == segment::SegState::failed)
            continue;
        vp.segments.push_back({v.start, v.end, v.completed});
        dl += v.completed;
    }
    vp.downloaded = dl;
    (void)meta::write_veloxpart_file(meta_path, vp, true);
}

void DownloadTaskState::emit_progress_if_due() {
    const auto now = std::chrono::steady_clock::now();
    const std::int64_t now_ns = now.time_since_epoch().count();
    std::int64_t prev = last_progress_ns.load(std::memory_order_relaxed);
    if (now_ns - prev < 250'000'000)
        return;
    if (!last_progress_ns.compare_exchange_strong(prev, now_ns))
        return;

    Progress p;
    {
        std::shared_lock lk(workers_mu);
        double agg = 0;
        for (auto &[idx, w] : workers) {
            agg += w->speed_bps;
            SegmentProgress sp;
            sp.index = idx;
            sp.speed_bps = static_cast<std::uint64_t>(w->speed_bps);
            p.segments.push_back(sp);
        }
        p.speed_bps = static_cast<std::uint64_t>(agg);
        p.effective_segments = static_cast<std::uint32_t>(workers.size());
    }
    if (seg) {
        p.downloaded = seg->downloaded();
        for (auto &v : seg->snapshot())
            for (auto &sp : p.segments)
                if (sp.index == v.index) {
                    sp.start = v.start;
                    sp.end = v.end;
                    sp.completed = v.completed;
                    sp.state = v.state;
                }
    }
    p.total = total_size;
    p.effective_buffer_bytes = effective_buffer;
    if (p.speed_bps > 0 && total_size && *total_size > p.downloaded)
        p.eta_seconds = static_cast<std::uint32_t>((*total_size - p.downloaded) / p.speed_bps);
    auto cb = cbs.on_progress;
    if (cb)
        cb(p);
}

// --- handle-facing --------------------------------------------------------------------

void DownloadTaskState::do_pause() {
    {
        std::unique_lock lk(mu);
        if (state == EngineState::paused || is_terminal(state))
            return;
        pause_requested = true;
        if (registered)
            host.budget().set_want(id, 0);
        if (workers.empty()) {
            if (file)
                (void)file->sync();
            write_sidecar_locked();
            transition(EngineState::paused, std::nullopt);
        } else {
            std::shared_lock wl(workers_mu);
            for (auto &[idx, w] : workers)
                w->transfer.cancel();
        }
    }
    flush_deferred();
}

void DownloadTaskState::do_resume() {
    {
        std::unique_lock lk(mu);
        if (state != EngineState::paused || awaiting_auth || awaiting_decision)
            return;
        pause_requested = false;
        transition(EngineState::connecting, std::nullopt);
        if (registered)
            host.budget().set_want(id, want_slots());
    }
    flush_deferred();
}

void DownloadTaskState::do_cancel(bool discard) {
    {
        std::unique_lock lk(mu);
        if (is_terminal(state))
            return;
        cancel_requested = true;
        discard_on_cancel = discard;
        if (workers.empty()) {
            finalize_cancel_locked();
        } else {
            std::shared_lock wl(workers_mu);
            for (auto &[idx, w] : workers)
                w->transfer.cancel();
        }
    }
    flush_deferred();
}

void DownloadTaskState::do_provide_auth(std::string u, std::string p, bool) {
    bool reprobe = false;
    {
        std::unique_lock lk(mu);
        if (!awaiting_auth)
            return;
        spec.auth = net::AuthConfig{net::AuthScheme::any, std::move(u), std::move(p)};
        awaiting_auth = false;
        reprobe = probe_needed_auth;
        probe_needed_auth = false;
        transition(EngineState::connecting, std::nullopt);
        if (!reprobe && registered)
            host.budget().set_want(id, want_slots());
    }
    flush_deferred();
    if (reprobe)
        restart_probe(true);
}

void DownloadTaskState::do_decide(Decision d) {
    {
        std::unique_lock lk(mu);
        if (!awaiting_decision)
            return;
        awaiting_decision = false;
        if (d == Decision::abort) {
            // Surface the reason the decision was actually asked for (range_metadata_stale
            // sets last_error to range_not_satisfiable, server_file_changed to itself), not
            // a hardcoded label that would misreport a 416 as a changed file.
            fail_locked(last_error.value_or(ErrorInfo(Error::server_file_changed, "user aborted")));
        } else {
            if (d == Decision::restart) {
                ::unlink(part_path.c_str());
                ::unlink(meta_path.c_str());
                spec.allow_resume = false;
                io::SparseFile::OpenOptions oo;
                oo.total_size = total_size.value_or(0);
                oo.preallocate = oo.total_size > 0;
                oo.truncate_existing = true;
                file = std::make_unique<io::SparseFile>();
                (void)file->open(part_path, oo);
                seg = std::make_unique<segment::Segmenter>(total_size.value_or(0),
                                                           requested_segments, resumable,
                                                           host.config().min_segment_bytes);
            }
            transition(EngineState::connecting, std::nullopt);
            if (registered)
                host.budget().set_want(id, want_slots());
        }
    }
    flush_deferred();
}

void DownloadTaskState::do_refresh_url(std::string url, std::vector<net::HeaderField> headers) {
    {
        std::unique_lock lk(mu);
        if (is_terminal(state))
            return;
        spec.url = std::move(url);
        if (!headers.empty())
            spec.headers = std::move(headers);
        std::shared_lock wl(workers_mu);
        for (auto &[idx, w] : workers)
            w->transfer.cancel();
    }
    flush_deferred();
    auto wp = weak_from_this();
    net::ProbeRequest pr;
    pr.url = spec.url;
    pr.headers = spec.headers;
    pr.auth = spec.auth;
    pr.proxy = spec.proxy;
    host.probe(std::move(pr), [wp](Result<net::ProbeResult> r) {
        auto s = wp.lock();
        if (!s)
            return;
        std::unique_lock lk(s->mu);
        if (s->retired.load() || is_terminal(s->state))
            return;
        if (r.has_value()) {
            s->probe.effective_url = r.value().effective_url;
            s->probe.etag = r.value().etag;
            s->probe.last_modified = r.value().last_modified;
        }
        if (s->registered)
            s->host.budget().set_want(s->id, s->want_slots());
        lk.unlock();
        s->flush_deferred();
    });
}

void DownloadTaskState::quiesce() {
    std::unique_lock lk(mu);
    retired.store(true);
    {
        std::shared_lock wl(workers_mu);
        for (auto &[idx, w] : workers)
            w->transfer.cancel();
    }
    // Do not clear `workers` here: transfer.cancel() only requests the HttpClient worker
    // thread stop the transfer, asynchronously -- it does not wait for that to happen. A
    // worker's SegWorker (and its WriteBuffer) may still be in active use by a curl write
    // callback running on that other thread right now. Clearing the map out from under it
    // was a real, ASan-caught heap-use-after-free (ring buffer freed here while
    // SparseFile::write_at() on the HttpClient worker thread was still writing through it).
    // Every worker removes and flushes itself, safely, via seg_finished once HttpClient
    // actually confirms the transfer has stopped (same path every other exit uses; see the
    // `retired` branch there) -- just wait for that to happen for all of them. Bounded by
    // however long a cancelled curl transfer takes to unwind, not user-controllable.
    workers_drained_cv.wait(lk, [this] { return workers.empty(); });
}

EngineState DownloadTaskState::snapshot_state() {
    std::lock_guard lk(mu);
    return state;
}

Progress DownloadTaskState::snapshot_progress() {
    Progress p;
    std::lock_guard lk(mu);
    if (seg) {
        p.downloaded = seg->downloaded();
        for (auto &v : seg->snapshot()) {
            SegmentProgress sp;
            sp.index = v.index;
            sp.start = v.start;
            sp.end = v.end;
            sp.completed = v.completed;
            sp.state = v.state;
            p.segments.push_back(sp);
        }
    }
    p.total = total_size;
    {
        std::shared_lock wl(workers_mu);
        p.effective_segments = static_cast<std::uint32_t>(workers.size());
    }
    p.effective_buffer_bytes = effective_buffer;
    return p;
}

// ==================================================================================

std::shared_ptr<DownloadTaskState> create_task(TaskHost &host, TaskId id, DownloadSpec spec,
                                               DownloadCallbacks cbs) {
    auto s = std::make_shared<DownloadTaskState>(host, id, std::move(spec), std::move(cbs));
    s->begin();
    return s;
}

void quiesce_task(const std::shared_ptr<DownloadTaskState> &s) {
    if (s)
        s->quiesce();
}

// --- DownloadHandle bodies ----------------------------------------------------------

TaskId DownloadHandle::id() const noexcept {
    return state_ ? state_->id : TaskId{};
}
void DownloadHandle::pause() {
    if (state_)
        state_->do_pause();
}
void DownloadHandle::resume() {
    if (state_)
        state_->do_resume();
}
void DownloadHandle::cancel(bool discard_partial) {
    if (state_)
        state_->do_cancel(discard_partial);
}
void DownloadHandle::provide_auth(std::string u, std::string p, bool remember) {
    if (state_)
        state_->do_provide_auth(std::move(u), std::move(p), remember);
}
void DownloadHandle::decide(Decision d) {
    if (state_)
        state_->do_decide(d);
}
void DownloadHandle::refresh_url(std::string url, std::vector<net::HeaderField> headers) {
    if (state_)
        state_->do_refresh_url(std::move(url), std::move(headers));
}
EngineState DownloadHandle::state() const {
    return state_ ? state_->snapshot_state() : EngineState::failed;
}
Progress DownloadHandle::progress() const {
    return state_ ? state_->snapshot_progress() : Progress{};
}

}  // namespace vdm::task