package ai.sentryagent.idp;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.net.http.HttpClient;
import java.time.Duration;
import java.util.concurrent.atomic.AtomicInteger;

import static org.junit.jupiter.api.Assertions.*;

class TokenManagerTest {

    private MockServer srv;
    private HttpClient httpClient;

    private static final String TOKEN_BODY = """
        {"access_token":"eyJ.abc.def","token_type":"Bearer","expires_in":3600,"scope":"agents:read"}
        """;

    @BeforeEach
    void setUp() throws IOException {
        srv = new MockServer();
        httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(5)).build();
    }

    @AfterEach
    void tearDown() { srv.stop(); }

    @Test
    void getToken_issuesToken() {
        srv.addHandler("/api/v1/token", 200, TOKEN_BODY);
        TokenManager tm = new TokenManager(srv.baseUrl(), "cid", "secret", "agents:read", httpClient);
        assertEquals("eyJ.abc.def", tm.getToken());
    }

    @Test
    void getToken_cachesToken() {
        AtomicInteger calls = new AtomicInteger(0);
        srv.addHandler("/api/v1/token", exchange -> {
            calls.incrementAndGet();
            byte[] body = TOKEN_BODY.getBytes();
            exchange.getResponseHeaders().set("Content-Type", "application/json");
            exchange.sendResponseHeaders(200, body.length);
            exchange.getResponseBody().write(body);
            exchange.getResponseBody().close();
        });
        TokenManager tm = new TokenManager(srv.baseUrl(), "cid", "secret", "agents:read", httpClient);
        tm.getToken();
        tm.getToken();
        assertEquals(1, calls.get(), "Should only call the token endpoint once");
    }

    @Test
    void getToken_authFailure_throwsAgentIdPException() {
        srv.addHandler("/api/v1/token", 401,
            "{\"error\":\"invalid_client\",\"error_description\":\"Bad credentials.\"}");
        TokenManager tm = new TokenManager(srv.baseUrl(), "cid", "bad-secret", "agents:read", httpClient);
        AgentIdPException ex = assertThrows(AgentIdPException.class, tm::getToken);
        assertEquals("invalid_client", ex.getCode());
        assertEquals(401, ex.getHttpStatus());
    }

    @Test
    void clearCache_forcesRefetch() {
        AtomicInteger calls = new AtomicInteger(0);
        srv.addHandler("/api/v1/token", exchange -> {
            calls.incrementAndGet();
            byte[] body = TOKEN_BODY.getBytes();
            exchange.getResponseHeaders().set("Content-Type", "application/json");
            exchange.sendResponseHeaders(200, body.length);
            exchange.getResponseBody().write(body);
            exchange.getResponseBody().close();
        });
        TokenManager tm = new TokenManager(srv.baseUrl(), "cid", "secret", "agents:read", httpClient);
        tm.getToken();
        tm.clearCache();
        tm.getToken();
        assertEquals(2, calls.get(), "Should call token endpoint again after clearCache");
    }

    @Test
    void getToken_threadSafe() throws InterruptedException {
        srv.addHandler("/api/v1/token", 200, TOKEN_BODY);
        TokenManager tm = new TokenManager(srv.baseUrl(), "cid", "secret", "agents:read", httpClient);

        Thread[] threads = new Thread[10];
        String[] results = new String[10];
        for (int i = 0; i < threads.length; i++) {
            int idx = i;
            threads[idx] = new Thread(() -> results[idx] = tm.getToken());
        }
        for (Thread t : threads) t.start();
        for (Thread t : threads) t.join();

        for (String result : results) {
            assertEquals("eyJ.abc.def", result);
        }
    }
}
