package mcp_test

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/lestrrat-go/jwx/v2/jwt"
	"github.com/mathiasbq/hyperguild/ingestion/internal/auth"
	"github.com/mathiasbq/hyperguild/ingestion/internal/mcp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func okHandler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	})
}

func TestBearerAuth_MissingHeader(t *testing.T) {
	handler := mcp.BearerAuth("secret", nil, okHandler())
	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)
	assert.Equal(t, http.StatusUnauthorized, rr.Code)
}

func TestBearerAuth_WrongToken(t *testing.T) {
	handler := mcp.BearerAuth("secret", nil, okHandler())
	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
	req.Header.Set("Authorization", "Bearer wrong")
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)
	assert.Equal(t, http.StatusUnauthorized, rr.Code)
}

func TestBearerAuth_CorrectToken(t *testing.T) {
	called := false
	handler := mcp.BearerAuth("secret", nil, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		called = true
		w.WriteHeader(http.StatusOK)
	}))
	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
	req.Header.Set("Authorization", "Bearer secret")
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)
	assert.Equal(t, http.StatusOK, rr.Code)
	assert.True(t, called)
}

func TestBearerAuth_EmptyConfiguredToken(t *testing.T) {
	handler := mcp.BearerAuth("", nil, okHandler())
	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)
	assert.Equal(t, http.StatusUnauthorized, rr.Code)
}

// JWT auth tests

func buildOIDCServer(t *testing.T) (*httptest.Server, jwk.Key) {
	t.Helper()
	raw, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	priv, err := jwk.FromRaw(raw)
	require.NoError(t, err)
	require.NoError(t, priv.Set(jwk.KeyIDKey, "k1"))
	require.NoError(t, priv.Set(jwk.AlgorithmKey, jwa.RS256))
	pub, err := jwk.PublicKeyOf(priv)
	require.NoError(t, err)

	set := jwk.NewSet()
	require.NoError(t, set.AddKey(pub))
	jwksBytes, err := json.Marshal(set)
	require.NoError(t, err)

	muxSrv := http.NewServeMux()
	var srv *httptest.Server
	muxSrv.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]string{
			"issuer":   srv.URL,
			"jwks_uri": srv.URL + "/jwks",
		})
	})
	muxSrv.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write(jwksBytes)
	})
	srv = httptest.NewServer(muxSrv)
	t.Cleanup(srv.Close)
	return srv, priv
}

func signJWT(t *testing.T, priv jwk.Key, issuer, audience string, exp time.Time) string {
	t.Helper()
	tok, err := jwt.NewBuilder().
		Issuer(issuer).Audience([]string{audience}).
		Subject("s").Expiration(exp).
		Build()
	require.NoError(t, err)
	signed, err := jwt.Sign(tok, jwt.WithKey(jwa.RS256, priv))
	require.NoError(t, err)
	return string(signed)
}

func TestBearerAuth_ValidJWT(t *testing.T) {
	oidcSrv, priv := buildOIDCServer(t)
	v, err := auth.NewValidator(oidcSrv.URL, "brain")
	require.NoError(t, err)

	called := false
	handler := mcp.BearerAuth("static-secret", v, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		called = true
		w.WriteHeader(http.StatusOK)
	}))

	token := signJWT(t, priv, oidcSrv.URL, "brain", time.Now().Add(time.Hour))
	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
	req.Header.Set("Authorization", "Bearer "+token)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)
	assert.Equal(t, http.StatusOK, rr.Code)
	assert.True(t, called)
}

func TestBearerAuth_InvalidJWT_FallsBackToStaticToken(t *testing.T) {
	oidcSrv, _ := buildOIDCServer(t)
	v, err := auth.NewValidator(oidcSrv.URL, "brain")
	require.NoError(t, err)

	handler := mcp.BearerAuth("static-secret", v, okHandler())
	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
	req.Header.Set("Authorization", "Bearer static-secret")
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)
	assert.Equal(t, http.StatusOK, rr.Code)
}

func TestBearerAuth_InvalidJWT_WrongStaticToken(t *testing.T) {
	oidcSrv, priv := buildOIDCServer(t)
	v, err := auth.NewValidator(oidcSrv.URL, "brain")
	require.NoError(t, err)

	handler := mcp.BearerAuth("static-secret", v, okHandler())
	// Expired JWT — JWT fails, static token doesn't match either
	token := signJWT(t, priv, oidcSrv.URL, "brain", time.Now().Add(-time.Hour))
	req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
	req.Header.Set("Authorization", "Bearer "+token)

	_ = context.Background() // satisfies import
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)
	assert.Equal(t, http.StatusUnauthorized, rr.Code)
}