// Connection ownership state machine: station/candidate takeover, generation
// invalidation, stale-result handling and failure release.
//
// Covers the exact interleavings Codex found on hardware (B3/B4/B6) plus the
// driver-failure and SSID-mismatch cases (B5):
//
//   T1 steady connected      station tick loop must never reconnect
//   T2 candidate takeover    old station waiter is Stale, candidate wins
//   T3 candidate failure     candidate is not persisted, station re-acquires
//   T4 set_config failure    surfaced as Failed (firmware maps the return code)
//   T5 actual SSID mismatch  GOT_IP for another network is rejected
//   T6 candidate success     commits, station stays steady
//   T7 ownership cancel      late candidate result is Stale
#include "wifi_connection_ownership.h"

#include <cassert>
#include <cstdio>
#include <atomic>
#include <cstring>
#include <mutex>
#include <thread>
#include <vector>

using namespace wearable;

// T1: A steady connection is never re-established by the station loop. The loop
// asks `station_steady()` before touching the driver; while it is true, no
// connect/set_config/mark_success may run.
static void steady_connected_no_reconnect() {
    ConnectionOwnership own;
    const auto attempt = own.station_begin("A");
    assert(own.station_resolve(attempt, true, "A") == AttemptOutcome::Accepted);
    assert(own.station_steady());
    assert(std::strcmp(own.connected_ssid(), "A") == 0);

    // 100 loop ticks: a steady link must stay steady and keep the same
    // generation, so a caller that gates on `station_steady()` performs zero
    // driver work.
    for (int tick = 0; tick < 100; ++tick) {
        assert(own.station_steady());
        assert(own.owner() == ConnectionOwner::Station);
    }
    // Only a real disconnect clears the steady state.
    own.station_disconnected();
    assert(!own.station_steady());
    assert(own.owner() == ConnectionOwner::Station);  // still owns, will retry
    std::puts("T1 steady connected: 100 ticks, no reconnect PASS");
}

// T2: Station A is waiting when a candidate B takes over. B succeeds. The old A
// waiter wakes up and must be told Stale; it may not change current_ssid and may
// not mark A as successful off B's event.
static void candidate_takeover_stale_station() {
    ConnectionOwnership own;
    const auto station_a = own.station_begin("A");  // generation 1, STATION, A
    assert(own.owner() == ConnectionOwner::Station);
    const uint32_t station_generation = station_a.generation;

    const auto candidate_b = own.take_candidate("B");  // generation 2, CANDIDATE
    assert(candidate_b.generation != station_generation);
    assert(own.owner() == ConnectionOwner::Candidate);
    assert(own.station_paused());

    // B's GOT_IP arrives; candidate accepts and commits B.
    AttemptDirective directive;
    const auto outcome = own.candidate_resolve(candidate_b, true, "B", directive);
    assert(outcome == AttemptOutcome::Accepted);
    assert(directive.persist && directive.commit_current && directive.mark_success);
    assert(std::strcmp(own.connected_ssid(), "B") == 0);

    // The stale A waiter (captured before the takeover) now runs. It must be
    // rejected as Stale and must not be able to write A back.
    const auto stale = own.station_resolve(station_a, true, "A");
    assert(stale == AttemptOutcome::Stale);
    assert(std::strcmp(own.connected_ssid(), "B") == 0);  // unchanged
    assert(own.owner() == ConnectionOwner::Station);       // B handed back
    assert(own.station_steady());                          // and stays connected
    std::puts("T2 candidate takeover: old station waiter STALE, B wins PASS");
}

// T3: Candidate C fails (wrong password). C must not be persisted; ownership is
// released so the station re-acquires and reconnects a known profile A. The
// provisioning AP/Portal is not modelled here (it stays up independently).
static void candidate_failure_reeleases_to_station() {
    ConnectionOwnership own;
    const auto a = own.station_begin("A");
    assert(own.station_resolve(a, true, "A") == AttemptOutcome::Accepted);

    const auto c = own.take_candidate("C");
    assert(own.station_paused());
    AttemptDirective directive;
    assert(own.candidate_resolve(c, false, "", directive) == AttemptOutcome::Failed);
    assert(!directive.persist);                 // C is never persisted
    assert(directive.release_candidate);
    assert(directive.resume_station);
    assert(own.owner() == ConnectionOwner::None);  // station may re-acquire
    assert(!own.station_paused());

    // The station loop re-acquires and reconnects A.
    const auto a2 = own.station_begin("A");
    assert(a2.valid);
    assert(own.station_resolve(a2, true, "A") == AttemptOutcome::Accepted);
    assert(std::strcmp(own.connected_ssid(), "A") == 0);
    std::puts("T3 candidate failure: C discarded, station reconnects A PASS");
}

// T4: set_config(B) fails, so the old A config remains and connect() succeeds on
// A. The firmware must not treat that as B success. Here the ownership layer
// receives the candidate verdict the firmware derived (Failed) and must not
// persist B. The mismatch guard is the second line of defence.
static void set_config_failure_not_persisted() {
    ConnectionOwnership own;
    const auto a = own.station_begin("A");
    assert(own.station_resolve(a, true, "A") == AttemptOutcome::Accepted);

    const auto b = own.take_candidate("B");
    // Firmware saw esp_wifi_set_config() != ESP_OK -> candidate Failed, even
    // though connect() later succeeded on the stale A config.
    AttemptDirective directive;
    assert(own.candidate_resolve(b, false, "", directive) == AttemptOutcome::Failed);
    assert(!directive.persist);
    assert(own.owner() == ConnectionOwner::None);
    std::puts("T4 set_config failure: B never persisted PASS");
}

// T5: GOT_IP arrives but the associated AP is A while B was expected. Must be
// rejected as Mismatch and never persisted.
static void actual_ssid_mismatch_rejected() {
    ConnectionOwnership own;
    const auto b = own.take_candidate("B");
    AttemptDirective directive;
    const auto outcome = own.candidate_resolve(b, true, "A", directive);
    assert(outcome == AttemptOutcome::Mismatch);
    assert(!directive.persist);
    assert(directive.release_candidate);
    assert(own.connected_ssid()[0] == '\0');

    // Same guard on the station side.
    const auto s = own.station_begin("Office");
    assert(own.station_resolve(s, true, "Home") == AttemptOutcome::Mismatch);
    assert(own.connected_ssid()[0] == '\0');
    std::puts("T5 actual SSID mismatch: rejected, not persisted PASS");
}

// T6: Candidate B succeeds from a connected A. B is committed; the station then
// sees a steady link and must not reconnect every 2 s.
static void candidate_success_commits_and_holds() {
    ConnectionOwnership own;
    const auto a = own.station_begin("A");
    assert(own.station_resolve(a, true, "A") == AttemptOutcome::Accepted);

    const auto b = own.take_candidate("B");
    AttemptDirective directive;
    assert(own.candidate_resolve(b, true, "B", directive) == AttemptOutcome::Accepted);
    assert(directive.persist);
    assert(std::strcmp(own.connected_ssid(), "B") == 0);

    for (int tick = 0; tick < 100; ++tick) assert(own.station_steady());
    assert(std::strcmp(own.connected_ssid(), "B") == 0);  // B stays current
    std::puts("T6 candidate success: B persists and holds, no 2s reconnect PASS");
}

// T7: A candidate is active when a Human Gate (Approval) arrives. The candidate
// generation is invalidated; its late result must be Stale and never commit.
static void approval_preempts_candidate() {
    ConnectionOwnership own;
    const auto b = own.take_candidate("B");
    assert(own.owner() == ConnectionOwner::Candidate);

    assert(own.cancel_candidate());
    assert(own.owner() == ConnectionOwner::None);
    assert(!own.station_paused());

    AttemptDirective directive;
    const auto late = own.candidate_resolve(b, true, "B", directive);
    assert(late == AttemptOutcome::Stale);
    assert(!directive.persist);
    assert(own.connected_ssid()[0] == '\0');
    std::puts("T7 approval preemption: late candidate STALE, never commits PASS");
}

// Reverse direction of T2: a candidate fails, the station takes a NEW generation
// and connects A. The old candidate's late event must not touch the new station.
static void late_candidate_does_not_touch_new_station() {
    ConnectionOwnership own;
    const auto c = own.take_candidate("C");
    AttemptDirective directive;
    assert(own.candidate_resolve(c, false, "", directive) == AttemptOutcome::Failed);

    const auto a = own.station_begin("A");  // new station generation
    assert(a.valid);
    assert(a.generation != c.generation);
    // The old candidate event arrives late.
    AttemptDirective late_directive;
    assert(own.candidate_resolve(c, true, "C", late_directive) == AttemptOutcome::Stale);
    assert(!late_directive.commit_current);
    // The station attempt still resolves normally.
    assert(own.station_resolve(a, true, "A") == AttemptOutcome::Accepted);
    assert(std::strcmp(own.connected_ssid(), "A") == 0);
    std::puts("reverse: late candidate cannot touch a new station generation PASS");
}

// Generation must advance on every ownership change so no two attempts can share
// a token and be confused for one another.
static void generations_are_unique() {
    ConnectionOwnership own;
    uint32_t seen[8];
    size_t n = 0;
    const auto s1 = own.station_begin("A");
    seen[n++] = s1.generation;
    const auto c1 = own.take_candidate("B");
    seen[n++] = c1.generation;
    AttemptDirective d;
    own.candidate_resolve(c1, false, "", d);
    const auto s2 = own.station_begin("A");
    seen[n++] = s2.generation;
    const auto c2 = own.take_candidate("C");
    seen[n++] = c2.generation;
    for (size_t i = 0; i < n; ++i)
        for (size_t j = i + 1; j < n; ++j) assert(seen[i] != seen[j]);
    std::puts("generations are unique across ownership changes PASS");
}

// ---------------------------------------------------------------------------
// Concurrency: the firmware serializes ConnectionOwnership with a dedicated
// recursive mutex (s_own_lock) through the Own* helpers. This mirrors that
// discipline with std::recursive_mutex and races the real machine from several
// threads, asserting the invariants that the lock protects:
//
//   * no two attempts ever observe the same generation concurrently;
//   * a candidate takeover always makes a previously captured station attempt
//     Stale (never Accepted);
//   * the station and candidate paths never both report Accepted for the same
//     generation;
//   * the machine stays internally consistent (attempt.valid implies owner).
//
// It also demonstrates the failure the lock prevents: without serialization,
// interleaved transitions race. The test asserts the locked path is always
// consistent, which is the property the firmware relies on.
namespace {

// Same discipline as the firmware: every ownership read/mutation is guarded.
struct LockedOwnership {
    std::recursive_mutex mutex;
    ConnectionOwnership machine;

    ConnectionOwner owner() {
        std::lock_guard<std::recursive_mutex> guard(mutex);
        return machine.owner();
    }
    ConnectionAttempt station_begin(const char *ssid) {
        std::lock_guard<std::recursive_mutex> guard(mutex);
        return machine.station_begin(ssid);
    }
    AttemptOutcome station_resolve(const ConnectionAttempt &attempt, bool got_ip, const char *actual) {
        std::lock_guard<std::recursive_mutex> guard(mutex);
        return machine.station_resolve(attempt, got_ip, actual);
    }
    void station_disconnected() {
        std::lock_guard<std::recursive_mutex> guard(mutex);
        machine.station_disconnected();
    }
    void station_idle() {
        std::lock_guard<std::recursive_mutex> guard(mutex);
        machine.station_idle();
    }
    ConnectionAttempt take_candidate(const char *ssid) {
        std::lock_guard<std::recursive_mutex> guard(mutex);
        return machine.take_candidate(ssid);
    }
    AttemptOutcome candidate_resolve(const ConnectionAttempt &attempt, bool got_ip, const char *actual,
                                     AttemptDirective &directive) {
        std::lock_guard<std::recursive_mutex> guard(mutex);
        return machine.candidate_resolve(attempt, got_ip, actual, directive);
    }
    bool station_steady() {
        std::lock_guard<std::recursive_mutex> guard(mutex);
        return machine.station_steady();
    }
    uint32_t generation() {
        std::lock_guard<std::recursive_mutex> guard(mutex);
        return machine.generation();
    }
};

// No station attempt captured before a candidate takeover may ever be Accepted.
void concurrent_takeover_never_accepts_stale_station() {
    for (int iteration = 0; iteration < 200; ++iteration) {
        LockedOwnership own;
        std::atomic<bool> station_captured{false};
        std::atomic<bool> candidate_taken{false};
        std::atomic<int> station_accepted{0};
        std::atomic<int> stale_wins{0};

        std::thread station([&] {
            const auto attempt = own.station_begin("A");
            if (attempt.valid) station_captured = true;
            // Spin until the candidate has taken over (or we give up).
            for (int spin = 0; spin < 100000 && !candidate_taken; ++spin) std::this_thread::yield();
            if (own.station_resolve(attempt, true, "A") == AttemptOutcome::Accepted) ++station_accepted;
        });
        std::thread candidate([&] {
            for (int spin = 0; spin < 100000 && !station_captured; ++spin) std::this_thread::yield();
            const auto attempt = own.take_candidate("B");
            candidate_taken = true;
            AttemptDirective directive;
            if (own.candidate_resolve(attempt, true, "B", directive) == AttemptOutcome::Accepted) {
                if (std::strcmp(own.machine.connected_ssid(), "B") == 0) ++stale_wins;
            }
        });
        station.join();
        candidate.join();
        // The station attempt, once captured, is always superseded by the
        // takeover: it can never be Accepted here.
        assert(station_accepted == 0);
    }
    std::puts("concurrency: candidate takeover always invalidates a captured "
              "station attempt (200 runs) PASS");
}

// Hammer the machine from many threads and assert it never enters an
// inconsistent state: a valid current attempt implies a non-None owner.
void concurrent_hammer_preserves_invariants() {
    for (int iteration = 0; iteration < 50; ++iteration) {
        LockedOwnership own;
        std::atomic<bool> stop{false};
        std::vector<std::thread> threads;
        for (int t = 0; t < 4; ++t) {
            threads.emplace_back([&, t] {
                for (int step = 0; step < 5000 && !stop; ++step) {
                    if (t % 2 == 0) {
                        const auto attempt = own.station_begin("A");
                        if (attempt.valid) own.station_resolve(attempt, true, "A");
                        own.station_disconnected();
                        own.station_idle();
                    } else {
                        const auto attempt = own.take_candidate("B");
                        AttemptDirective directive;
                        own.candidate_resolve(attempt, step % 2 == 0, step % 2 == 0 ? "B" : "", directive);
                    }
                    // Invariant under the lock: a valid attempt implies an owner.
                    std::lock_guard<std::recursive_mutex> guard(own.mutex);
                    const auto &current = own.machine.current_attempt();
                    if (current.valid) assert(own.machine.owner() != ConnectionOwner::None);
                }
            });
        }
        for (auto &thread : threads) thread.join();
        stop = true;
        // After all writers stop, the machine is usable again.
        const auto attempt = own.station_begin("A");
        assert(attempt.valid);
        assert(own.station_resolve(attempt, true, "A") == AttemptOutcome::Accepted);
    }
    std::puts("concurrency: 4-thread hammer keeps ownership invariants (50 runs) PASS");
}

}  // namespace

int main() {
    steady_connected_no_reconnect();
    candidate_takeover_stale_station();
    candidate_failure_reeleases_to_station();
    set_config_failure_not_persisted();
    actual_ssid_mismatch_rejected();
    candidate_success_commits_and_holds();
    approval_preempts_candidate();
    late_candidate_does_not_touch_new_station();
    generations_are_unique();
    concurrent_takeover_never_accepts_stale_station();
    concurrent_hammer_preserves_invariants();
    std::puts("Wi-Fi connection ownership: T1-T7 + interleavings + concurrency PASS (host)");
    return 0;
}