#include "raft.h"

void raft_snapshot(const raft_state_t *s, raft_snapshot_t *snap)
{
    snap->term             = s->term;
    snap->role             = s->role;
    snap->voted_for        = s->voted_for;
    snap->voted_for_term   = s->voted_for_term;
    snap->votes_granted    = s->votes_granted;
    snap->election_decided = s->election_decided;
}

int raft_init(raft_state_t *s, const raft_config_t *cfg)
{
    /* Validate required callbacks and n_peers before touching *s */
    if (cfg->n_peers < 1)
        return -EINVAL;
    /* Upper bound: next_index[]/match_index[]/pre_vote_replies[] are fixed-size
     * RAFT_MAX_PEERS arrays indexed by peer in [0, n_peers).  Reject oversize
     * clusters so the leader-progress and pre-vote loops can never OOB-write. */
    if (cfg->n_peers > RAFT_MAX_PEERS)
        return -EINVAL;
    /* Even cluster sizes (>1) have no fault-tolerance advantage over the next
     * smaller odd size and can mislead operators into thinking split-brain is
     * impossible.  Reject to prevent misconfiguration. */
    if (cfg->n_peers > 1 && cfg->n_peers % 2 == 0)
        return -EINVAL;
    if (!cfg->cbs)
        return -EINVAL;
    if (!cfg->cbs->on_send_vote_request)
        return -EINVAL;
    if (!cfg->cbs->on_become_leader)
        return -EINVAL;
    if (!cfg->cbs->on_become_follower)
        return -EINVAL;
    if (!cfg->cbs->on_persist)
        return -EINVAL;
    if (!cfg->cbs->on_schedule_timeout)
        return -EINVAL;
    if (!cfg->cbs->on_cancel_timeout)
        return -EINVAL;
    if (!cfg->cbs->on_start_election_after_delay)
        return -EINVAL;
    /* Log-replication required callbacks (raft-log-replication Req 7.1–7.4):
     * all four must be non-NULL — the library owns the protocol decisions and
     * delegates the I/O (persist/read-term/send/truncate) through them. */
    if (!cfg->cbs->on_persist_entry)
        return -EINVAL;
    if (!cfg->cbs->on_get_entry_term)
        return -EINVAL;
    if (!cfg->cbs->on_send_append_entries)
        return -EINVAL;
    if (!cfg->cbs->on_truncate_from)
        return -EINVAL;

    s->term               = cfg->initial_term;
    s->role               = RAFT_FOLLOWER;
    s->voted_for          = cfg->initial_voted_for;
    /* A restored vote belongs to the restored term: persistence writes
     * (term, voted_for) as a pair.  Leaving voted_for_term at 0 would let the
     * duplicate-vote check miss after restart, granting a second vote in the
     * same term (Election Safety violation). */
    s->voted_for_term     = (cfg->initial_voted_for != -1) ? cfg->initial_term : 0;
    s->n_peers            = cfg->n_peers;
    s->self_idx           = cfg->self_idx;
    s->votes_granted      = 0;
    s->replies_received   = 0;
    s->election_decided   = 0;
    s->election_timeout_ms = cfg->election_timeout_ms ? cfg->election_timeout_ms : 2000;
    s->my_priority        = cfg->my_priority;
    /* Log length restored from the adapter's recovered WAL length; the library
     * never reads storage (same discipline as initial_term/voted_for). 0 on
     * first boot. (raft-log-replication Req 1.1; design "log_len restored at
     * init from initial_log_len".) */
    s->log_len            = cfg->initial_log_len;
    s->cbs                = cfg->cbs;
    s->ctx                = cfg->ctx;

    return 0;
}

/* Initialize per-peer leader progress on becoming leader (raft.tla BecomeLeader,
 * Req 3.1): next_index[j] = log_len + 1, match_index[j] = 0 for every peer j. */
static void raft_init_leader_progress(raft_state_t *s)
{
    for (int j = 0; j < s->n_peers; j++) {
        s->next_index[j]  = s->log_len + 1;
        s->match_index[j] = 0;
    }
}

int raft_start_election(raft_state_t *s)
{
    /* Guard: reject if already CANDIDATE or LEADER */
    if (s->role == RAFT_CANDIDATE || s->role == RAFT_LEADER)
        return -EALREADY;

    /* Increment term and record self-vote before any RPCs (split-brain prevention) */
    s->term++;
    s->voted_for        = s->self_idx;
    s->voted_for_term   = s->term;
    s->votes_granted    = 1;  /* self-vote */
    s->replies_received = 0;
    s->election_decided = 0;

    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "term_update", NULL);
    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "self_vote", NULL);

    /* Persist new term and self-vote before sending any RPCs */
    s->cbs->on_persist(s->ctx, s->term, s->self_idx);

    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "election_start", NULL);

    /* Single-node shortcircuit: win immediately, no RPCs, no timeout */
    if (s->n_peers == 1) {
        s->role             = RAFT_LEADER;
        s->election_decided = 1;
        raft_init_leader_progress(s);
        s->cbs->on_become_leader(s->ctx, s->term);
        return RAFT_ELECTED_IMMEDIATELY;
    }

    /* Multi-node: transition to CANDIDATE */
    s->role = RAFT_CANDIDATE;

    /* Send RequestVote to all peers except self */
    for (int i = 0; i < s->n_peers; i++) {
        if (i == s->self_idx)
            continue;
        s->cbs->on_send_vote_request(s->ctx, i, s->term);
    }

    /* Schedule election timeout watchdog */
    s->cbs->on_schedule_timeout(s->ctx, s->election_timeout_ms);

    return RAFT_OK;
}

int raft_on_vote_request(raft_state_t *s, int from_idx, uint64_t term,
                          uint64_t *reply_term)
{
    /* Step 1: stale term — deny immediately, no state change */
    if (term < s->term) {
        *reply_term = s->term;
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "vote_denied_stale_term", NULL);
        return RAFT_DENY;
    }

    /* Step 2: term upgrade — update term, clear vote, step down to FOLLOWER.
     * If we were LEADER, notify the HA FSM so it updates ha->role.
     * Without this, the HA FSM stays in CTRL_PRIMARY after losing the term,
     * never stepping down to STANDBY → split-brain. */
    if (term > s->term) {
        int was_leader = (s->role == RAFT_LEADER);
        s->term         = term;
        s->voted_for    = -1;
        s->voted_for_term = 0;
        s->role         = RAFT_FOLLOWER;
        if (was_leader) {
            s->cbs->on_become_follower(s->ctx, s->term, from_idx);
        }
    }

    /* reply_term reflects our (possibly updated) current term from here on */
    *reply_term = s->term;

    /* Step 3: role gate — same-term challenger denied by live leader.
     *
     * This only fires when term == s->term (the equal-term case).  For
     * term > s->term, Step 2 already downgraded s->role to RAFT_FOLLOWER
     * before we reach here, so the check below is false and we fall through
     * to grant/deny based on freshness.
     *
     * A RequestVote for the same term as the current leader means the sender
     * started a competing candidacy in a term the leader already won — deny. */
    if (s->role == RAFT_LEADER) {
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "vote_denied_role", NULL);
        return RAFT_DENY;
    }

    /* Step 4: duplicate-vote check — already voted in this term.
     * Idempotent re-grant: if this is the same candidate we already voted for
     * in this term, fall through so we re-persist and re-grant (network retry
     * path).  This is safe: the candidate has already committed to this term
     * and we already durably recorded the vote.
     * Different candidate: deny — standard Raft invariant of at most one vote
     * per node per term. */
    if (s->voted_for != -1 && s->voted_for_term == term) {
        if (s->voted_for != from_idx) {
            if (s->cbs->on_trace)
                s->cbs->on_trace(s->ctx, "vote_denied_already_voted", NULL);
            return RAFT_DENY;
        }
        /* same candidate — fall through to re-persist and re-grant */
    }

    /* Step 5: application-level policy gate (freshness, log lsn, etc.) */
    if (s->cbs->on_vote_eligible != NULL) {
        if (s->cbs->on_vote_eligible(s->ctx, from_idx, term) == 0) {
            if (s->cbs->on_trace)
                s->cbs->on_trace(s->ctx, "vote_denied_freshness", NULL);
            return RAFT_DENY;
        }
    }

    /* Step 6: persist before granting — durability guarantee */
    if (s->cbs->on_persist(s->ctx, term, from_idx) != 0) {
        /* Persist failed — deny without updating voted_for */
        return RAFT_DENY;
    }

    /* Grant: record vote */
    s->voted_for      = from_idx;
    s->voted_for_term = term;
    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "vote_granted", NULL);
    return RAFT_GRANT;
}

int raft_on_pre_vote_request(raft_state_t *s, int from_idx, uint64_t cand_term)
{
    /* PreVote is a side-effect-free straw poll: it answers "would you vote for
     * me in cand_term?" WITHOUT changing any persistent state.  Unlike
     * raft_on_vote_request, this function never adopts cand_term (even when
     * higher — that adoption is the very thing PreVote exists to prevent), never
     * touches term / voted_for / voted_for_term / role, and never calls
     * on_persist.  The dispatcher applies the freshness floor before this call,
     * exactly as on the real-vote path. */

    /* Step 1: stale term — deny (mirrors raft_on_vote_request Step 1).  A
     * candidate proposing a term we already passed could never win a real
     * election, so it cannot win a PreVote either.  No term adoption: even when
     * cand_term > s->term we only evaluate, never adopt. */
    if (cand_term < s->term) {
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "pre_vote_denied_stale_term", NULL);
        return RAFT_DENY;
    }

    /* Step 2: application-level eligibility gate — reuse the EXISTING
     * on_vote_eligible callback (primary-deny / fresh-heartbeat stickiness /
     * catch-up deny), identical to raft_on_vote_request Step 5.  If NULL, every
     * Raft-eligible candidate is granted. */
    if (s->cbs->on_vote_eligible != NULL) {
        if (s->cbs->on_vote_eligible(s->ctx, from_idx, cand_term) == 0) {
            if (s->cbs->on_trace)
                s->cbs->on_trace(s->ctx, "pre_vote_denied_eligibility", NULL);
            return RAFT_DENY;
        }
    }

    /* Grant — record NOTHING.  No voted_for write (so a later real vote in this
     * term is still grantable), no persist, no term/role change. */
    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "pre_vote_granted", NULL);
    return RAFT_GRANT;
}

int raft_begin_pre_vote(raft_state_t *s)
{
    /* Pre-image of raft_start_election: probe peers + tally a self-grant, but
     * STRIP every persistent side effect.  Per the header INVARIANT (Req 1.2,
     * 1.3, 6.2) this function MUST NOT touch s->term, s->voted_for,
     * s->voted_for_term, or s->role, and MUST NOT call on_persist.  The node
     * stays FOLLOWER throughout the round; CANDIDATE happens later in the
     * unchanged raft_start_election once the PreVote round is WON. */

    /* Prospective term is a LOCAL value only — never assigned to s->term and
     * never persisted (this is the whole point of PreVote). */
    uint64_t prospective_term = s->term + 1;

    /* Reset and arm the round tally on the raft state.  Self always grants. */
    s->pre_vote_in_progress = 1;
    s->pre_votes_granted    = 1;  /* self-grant */
    memset(s->pre_vote_replies, 0, sizeof(s->pre_vote_replies));

    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "pre_vote_start", NULL);

    /* Single-node: no peers, the self-grant is already a majority.  Win
     * immediately and let the caller proceed to raft_start_election. */
    if (s->n_peers == 1)
        return RAFT_PRE_VOTE_WON;

    /* Probe every peer except self with the prospective term (NULL-checked:
     * if the callback is unset, the caller owns the transport). */
    if (s->cbs->on_send_pre_vote_request) {
        for (int i = 0; i < s->n_peers; i++) {
            if (i == s->self_idx)
                continue;
            s->cbs->on_send_pre_vote_request(s->ctx, i, prospective_term);
        }
    }

    return RAFT_OK;
}

int raft_on_pre_vote_reply(raft_state_t *s, int peer_idx, int granted)
{
    /* Pure straw-poll tally — mirrors raft_on_vote_reply's quorum and
     * all-replied-no-quorum logic, but STRIPS every side effect.  Per the
     * header INVARIANT (Req 1.4, 1.5, 6.2) this function MUST NOT fire
     * on_become_leader / on_become_follower, MUST NOT touch s->term /
     * s->voted_for / s->voted_for_term / s->role, and MUST NOT call on_persist.
     * It only writes the pre_vote_* tally fields and returns a status the
     * controller (task 4.1) acts on.  There is no term in a pre-vote reply, so
     * no term adoption is even possible here. */

    /* Bounds guard: ignore replies from an out-of-range peer (mirrors the
     * per-peer dedup array bound).  A self reply (peer_idx == self_idx) is
     * never expected — self is granted by raft_begin_pre_vote — and would also
     * be filtered by the already-set flag, but the range check alone is the
     * defensive boundary. */
    if (peer_idx < 0 || peer_idx >= s->n_peers || peer_idx >= RAFT_MAX_PEERS) {
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "pre_vote_reply_out_of_range", NULL);
        return RAFT_PRE_VOTE_PENDING;
    }

    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "pre_vote_reply_in", NULL);

    /* Per-peer idempotency: a duplicate reply from a peer we already counted is
     * ignored (no double-count).  Stay undecided. */
    if (s->pre_vote_replies[peer_idx]) {
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "pre_vote_reply_dup", NULL);
        return RAFT_PRE_VOTE_PENDING;
    }
    s->pre_vote_replies[peer_idx] = 1;

    /* Accumulate grant.  pre_votes_granted already includes the self-grant set
     * by raft_begin_pre_vote. */
    if (granted)
        s->pre_votes_granted++;

    /* Quorum mirrors raft_on_vote_reply: n_peers is the TOTAL cluster size
     * including self. */
    int quorum = (s->n_peers / 2) + 1;
    if (s->pre_votes_granted >= quorum) {
        s->pre_vote_in_progress = 0;
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "pre_vote_won", NULL);
        return RAFT_PRE_VOTE_WON;
    }

    /* All non-self peers have replied without reaching quorum → lost.  Count
     * the set per-peer flags (a self reply, if it ever arrived, would have been
     * range-filtered out of n_peers - 1 expectation only if self_idx is in
     * range, so compare against n_peers - 1 replied flags). */
    int replied = 0;
    for (int i = 0; i < s->n_peers; i++)
        replied += s->pre_vote_replies[i];
    if (replied == s->n_peers - 1) {
        s->pre_vote_in_progress = 0;
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "pre_vote_lost", NULL);
        return RAFT_PRE_VOTE_LOST;
    }

    /* Still awaiting replies. */
    return RAFT_PRE_VOTE_PENDING;
}

int raft_on_vote_reply(raft_state_t *s, int from_idx, int granted, uint64_t term)
{
    (void)from_idx;

    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "vote_reply_in", NULL);

    /* Step 1: term upgrade — higher term means we lost the election */
    if (term > s->term) {
        s->term            = term;
        s->role            = RAFT_FOLLOWER;
        s->election_decided = 1;
        s->cbs->on_become_follower(s->ctx, s->term, -1);
        s->cbs->on_cancel_timeout(s->ctx);
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "became_follower", NULL);
        return RAFT_BECAME_FOLLOWER;
    }

    /* Step 2: no-op guard — election already decided */
    if (s->election_decided) {
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "vote_reply_noop", NULL);
        return RAFT_PENDING;
    }

    /* Step 3: count this reply */
    s->replies_received++;

    /* Step 4: accumulate grant */
    if (granted)
        s->votes_granted++;

    /* Step 5: quorum check */
    int quorum = (s->n_peers / 2) + 1;
    if (s->votes_granted >= quorum) {
        s->role             = RAFT_LEADER;
        s->election_decided = 1;
        raft_init_leader_progress(s);
        s->cbs->on_become_leader(s->ctx, s->term);
        s->cbs->on_cancel_timeout(s->ctx);
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "quorum_reached", NULL);
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "became_leader", NULL);
        return RAFT_BECAME_LEADER;
    }

    /* Step 6: all-replied-no-quorum */
    if (s->replies_received == s->n_peers - 1 && s->votes_granted < quorum) {
        s->role             = RAFT_FOLLOWER;
        s->election_decided = 1;
        s->cbs->on_become_follower(s->ctx, s->term, -1);
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "became_follower", NULL);
        return RAFT_BECAME_FOLLOWER;
    }

    /* Step 7: still waiting */
    return RAFT_PENDING;
}

void raft_on_election_timeout(raft_state_t *s)
{
    if (s->election_decided) {
        /* Req 6.3: cancel any pending timer and be a no-op */
        s->cbs->on_cancel_timeout(s->ctx);
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "election_timeout_noop", NULL);
        return;
    }

    s->role             = RAFT_FOLLOWER;
    s->election_decided = 1;
    s->cbs->on_become_follower(s->ctx, s->term, -1);
    s->cbs->on_cancel_timeout(s->ctx);
    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "election_timeout", NULL);
}

void raft_on_leader_lost(raft_state_t *s)
{
    s->voted_for = -1;
    s->role      = RAFT_FOLLOWER;
    s->cbs->on_start_election_after_delay(s->ctx,
        (uint32_t)s->my_priority * s->election_timeout_ms);
    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "leader_lost", NULL);
}

int raft_check_and_start_election(raft_state_t *s)
{
    if (s->voted_for != -1 && s->voted_for != s->self_idx) {
        if (s->cbs->on_trace)
            s->cbs->on_trace(s->ctx, "jitter_skip", NULL);
        return RAFT_SKIP_ELECTION;
    }

    /* Precondition: raft_on_leader_lost sets role=FOLLOWER, so -EALREADY is unreachable here */
    int rc = raft_start_election(s);
    (void)rc;
    return RAFT_ELECTION_STARTED;
}

int raft_log_append(raft_state_t *s, const void *entry, size_t len,
                    uint64_t *out_index)
{
    /* Role gate (Req 2.2): only a leader may grow the log.  Reject without any
     * log mutation or callback (raft.tla ClientRequest is enabled only at the
     * leader). */
    if (s->role != RAFT_LEADER)
        return -EPERM;

    /* Assign the next 1-based index (== WAL LSN) and stamp the leader's current
     * term, mirroring raft.tla ClientRequest's [term ↦ currentTerm, value ↦ v]
     * (Req 2.1, 2.3, 2.4). */
    uint64_t index = s->log_len + 1;
    uint64_t term  = s->term;

    /* Persist-before-act (Req 2.3, 2.5): durably write the entry BEFORE
     * advancing log_len.  On failure, propagate the error verbatim and leave the
     * log untouched (no phantom entry). */
    int rc = s->cbs->on_persist_entry(s->ctx, index, term, entry, len);
    if (rc != 0)
        return rc;

    /* Commit the append only after a successful persist. */
    s->log_len = index;
    if (out_index)
        *out_index = index;

    if (s->cbs->on_trace)
        s->cbs->on_trace(s->ctx, "log_append", NULL);

    return RAFT_OK;
}