import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;

/**
 * Taiga MOAD-0003: threading.local session_id — request-scoped identity in thread-local.
 *
 * taiga/events/middleware.py stores the X-Session-ID header in a threading.local.
 * Under thread-pool workers (gunicorn sync workers, celery) the same OS thread can
 * serve multiple requests. If a background task or a signal handler fires after the
 * response is returned but before process_response clears the local, the session_id
 * bleeds into an unrelated request context.
 *
 * Fix: replace threading.local with contextvars.ContextVar — each task/coroutine gets
 * its own copy automatically.
 *
 * This test models the defect with ThreadLocal<String> (Java analog) vs
 * per-invocation context passing and demonstrates cross-thread contamination.
 */
public class TaigaSessionContextVarTest {

    // --- Defective model: ThreadLocal session_id ---
    static final ThreadLocal<String> SESSION_ID_LOCAL = new ThreadLocal<>();

    static String getSessionIdDefect() {
        return SESSION_ID_LOCAL.get();
    }

    static void setSessionIdDefect(String id) {
        SESSION_ID_LOCAL.set(id);
    }

    static void clearSessionIdDefect() {
        SESSION_ID_LOCAL.remove();
    }

    // --- Fixed model: context passed per invocation ---
    // In Python this is ContextVar which propagates along the call chain
    // automatically. In Java we model it as explicit parameter passing.
    static String getSessionIdFixed(Map<String, String> ctx) {
        return ctx.get("session_id");
    }

    // --- Correctness: verify set/get/clear works as expected ---
    static void testCorrectness() {
        setSessionIdDefect("session-abc");
        assert "session-abc".equals(getSessionIdDefect()) : "Expected session-abc";

        clearSessionIdDefect();
        assert getSessionIdDefect() == null : "Expected null after clear";

        Map<String, String> ctx = new HashMap<>();
        ctx.put("session_id", "session-xyz");
        assert "session-xyz".equals(getSessionIdFixed(ctx)) : "Expected session-xyz";
        System.out.println("PASS correctness");
    }

    // --- Contamination test: thread-local bleeds across simulated requests ---
    static void testThreadLocalContamination() throws Exception {
        // Simulate a thread-pool worker handling two requests sequentially.
        // Request 1 sets session_id but "crashes" before clearing.
        // Request 2 should see null, but with ThreadLocal it sees Request 1's value.
        AtomicReference<String> leaked = new AtomicReference<>();

        ExecutorService pool = Executors.newSingleThreadExecutor();
        pool.submit(() -> {
            // Request 1: set and "forget" to clear (simulates crash/early return)
            setSessionIdDefect("leaked-session-111");
            // process_response never called
        }).get();

        pool.submit(() -> {
            // Request 2: no set — should start clean, but thread is reused
            leaked.set(getSessionIdDefect());
        }).get();

        pool.shutdown();

        String leakedVal = leaked.get();
        // This demonstrates the defect: ThreadLocal bleeds "leaked-session-111" into Request 2
        System.out.println("ThreadLocal contamination value: " + leakedVal
            + (leakedVal != null ? " (DEFECT: session leaked)" : " (clean)"));
        // Confirm the defect exists
        assert "leaked-session-111".equals(leakedVal)
            : "Expected leak but got: " + leakedVal;
        System.out.println("PASS contamination demonstration");
    }

    // --- Fixed model has no contamination: context is per-call ---
    static void testContextVarNoContamination() throws Exception {
        AtomicReference<String> request2Session = new AtomicReference<>();

        ExecutorService pool = Executors.newSingleThreadExecutor();
        pool.submit(() -> {
            // Request 1: uses its own local map — doesn't affect other invocations
            Map<String, String> ctx1 = new HashMap<>();
            ctx1.put("session_id", "session-req1");
            getSessionIdFixed(ctx1); // just uses it, never stored globally
        }).get();

        pool.submit(() -> {
            // Request 2: fresh context, no bleed
            Map<String, String> ctx2 = new HashMap<>();
            // session_id not set for request 2
            request2Session.set(getSessionIdFixed(ctx2));
        }).get();

        pool.shutdown();

        assert request2Session.get() == null
            : "Fixed model must not bleed session, got: " + request2Session.get();
        System.out.println("PASS no contamination (fixed model)");
    }

    public static void main(String[] args) throws Exception {
        testCorrectness();
        testThreadLocalContamination();
        testContextVarNoContamination();
        System.out.println("ALL PASS");
    }
}