package auth

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"math/big"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/crewjam/saml"
	"github.com/crewjam/saml/samlsp"
	"github.com/gofiber/fiber/v2"
	"wucher/internal/config"
)

func TestHandleBrowserLocalLogout_RedirectAndClearSessionCookie(t *testing.T) {
	prev := SAMLMiddleware
	SAMLMiddleware = nil
	t.Cleanup(func() { SAMLMiddleware = prev })

	h := NewHandler(config.Config{
		FrontendURL:  "https://wfm-dev.mybit.co.id",
		CookieDomain: ".mybit.co.id",
		CookieSecure: true,
	}, nil)

	app := fiber.New()
	app.Get("/api/v1/auth/microsoft/logout", h.HandleBrowserLocalLogout)

	req := httptest.NewRequest(http.MethodGet, "/api/v1/auth/microsoft/logout", nil)
	resp, err := app.Test(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	if resp.StatusCode != http.StatusFound {
		t.Fatalf("expected 302, got %d", resp.StatusCode)
	}
	if got := resp.Header.Get("Location"); got != "https://wfm-dev.mybit.co.id/signin?loggedOut=true" {
		t.Fatalf("unexpected redirect: %s", got)
	}

	setCookie := strings.Join(resp.Header.Values("Set-Cookie"), "\n")
	if !strings.Contains(setCookie, "session_token=") {
		t.Fatalf("expected session_token clear cookie, got %q", setCookie)
	}
	if !strings.Contains(strings.ToLower(setCookie), "domain=.mybit.co.id") {
		t.Fatalf("expected cookie domain .mybit.co.id, got %q", setCookie)
	}
}

func TestBuildRedirectSignaturePayload_PreservesRawEncoding(t *testing.T) {
	raw := "SAMLResponse=abc%2B123%3D%3D&RelayState=x%2By%252Fz&SigAlg=http%3A%2F%2Fwww.w3.org%2F2001%2F04%2Fxmldsig-more%23rsa-sha256&Signature=ignored"
	got, err := buildRedirectSignaturePayload(raw, "SAMLResponse")
	if err != nil {
		t.Fatalf("build payload: %v", err)
	}
	want := "SAMLResponse=abc%2B123%3D%3D&RelayState=x%2By%252Fz&SigAlg=http%3A%2F%2Fwww.w3.org%2F2001%2F04%2Fxmldsig-more%23rsa-sha256"
	if string(got) != want {
		t.Fatalf("payload mismatch\nwant: %s\ngot:  %s", want, string(got))
	}
}

func TestHandleSAMLLogoutResponse_GetRedirectSignature_Valid_WithRelayState(t *testing.T) {
	runValidSAMLResponseRedirectTest(t, "relay-state-1", true)
}

func TestHandleSAMLLogoutResponse_GetRedirectSignature_Valid_WithoutRelayState(t *testing.T) {
	runValidSAMLResponseRedirectTest(t, "", false)
}

func TestHandleSAMLLogoutResponse_GetRedirectSignature_Invalid(t *testing.T) {
	priv, cert := testSigningCert(t)
	setTestSAMLMiddlewareWithCert(t, cert, "https://sts.windows.net/tenant/")

	xmlB64 := base64.StdEncoding.EncodeToString([]byte(testLogoutResponseXML("https://sts.windows.net/tenant/")))
	sigAlg := "http://www.w3.org/2001/04/xmldsig-more#rsa-sha256"
	params := []string{
		"SAMLResponse=" + url.QueryEscape(xmlB64),
		"SigAlg=" + url.QueryEscape(sigAlg),
	}
	canonical := strings.Join(params, "&")
	sig := signRedirectPayloadRSA256(t, priv, []byte(canonical))
	// Break signature on purpose
	sig = strings.TrimRight(sig, "=") + "A"

	path := "/api/v1/auth/microsoft/logout?" + canonical + "&Signature=" + url.QueryEscape(sig)
	h := NewHandler(config.Config{FrontendURL: "https://wfm-dev.mybit.co.id"}, nil)
	app := fiber.New()
	app.Get("/api/v1/auth/microsoft/logout", h.HandleSAMLLogoutResponse)

	req := httptest.NewRequest(http.MethodGet, path, nil)
	resp, err := app.Test(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("expected 401, got %d", resp.StatusCode)
	}
}

func TestHandleSAMLLogoutResponse_GetMissingSignature(t *testing.T) {
	_, cert := testSigningCert(t)
	setTestSAMLMiddlewareWithCert(t, cert, "https://sts.windows.net/tenant/")

	xmlB64 := base64.StdEncoding.EncodeToString([]byte(testLogoutResponseXML("https://sts.windows.net/tenant/")))
	path := "/api/v1/auth/microsoft/logout?SAMLResponse=" + url.QueryEscape(xmlB64)

	h := NewHandler(config.Config{FrontendURL: "https://wfm-dev.mybit.co.id"}, nil)
	app := fiber.New()
	app.Get("/api/v1/auth/microsoft/logout", h.HandleSAMLLogoutResponse)

	req := httptest.NewRequest(http.MethodGet, path, nil)
	resp, err := app.Test(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("expected 401, got %d", resp.StatusCode)
	}
}

func runValidSAMLResponseRedirectTest(t *testing.T, relayState string, withRelay bool) {
	t.Helper()
	priv, cert := testSigningCert(t)
	setTestSAMLMiddlewareWithCert(t, cert, "https://sts.windows.net/tenant/")

	xmlB64 := base64.StdEncoding.EncodeToString([]byte(testLogoutResponseXML("https://sts.windows.net/tenant/")))
	sigAlg := "http://www.w3.org/2001/04/xmldsig-more#rsa-sha256"

	parts := []string{"SAMLResponse=" + url.QueryEscape(xmlB64)}
	if withRelay {
		parts = append(parts, "RelayState="+url.QueryEscape(relayState))
	}
	parts = append(parts, "SigAlg="+url.QueryEscape(sigAlg))
	canonical := strings.Join(parts, "&")
	sig := signRedirectPayloadRSA256(t, priv, []byte(canonical))
	rawQuery := canonical + "&Signature=" + url.QueryEscape(sig)

	h := NewHandler(config.Config{FrontendURL: "https://wfm-dev.mybit.co.id"}, nil)
	app := fiber.New()
	app.Get("/api/v1/auth/microsoft/logout", h.HandleSAMLLogoutResponse)

	req := httptest.NewRequest(http.MethodGet, "/api/v1/auth/microsoft/logout?"+rawQuery, nil)
	resp, err := app.Test(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	if resp.StatusCode != http.StatusFound {
		t.Fatalf("expected 302, got %d", resp.StatusCode)
	}
	if got := resp.Header.Get("Location"); got != "https://wfm-dev.mybit.co.id/signin?loggedOut=true" {
		t.Fatalf("unexpected redirect: %s", got)
	}
}

func signRedirectPayloadRSA256(t *testing.T, priv *rsa.PrivateKey, payload []byte) string {
	t.Helper()
	sum := sha256.Sum256(payload)
	sig, err := rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA256, sum[:])
	if err != nil {
		t.Fatalf("sign payload: %v", err)
	}
	return base64.StdEncoding.EncodeToString(sig)
}

func testSigningCert(t *testing.T) (*rsa.PrivateKey, *x509.Certificate) {
	t.Helper()
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("gen key: %v", err)
	}
	serial, _ := rand.Int(rand.Reader, big.NewInt(1<<62))
	tpl := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: "test-idp"},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(24 * time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature,
	}
	der, err := x509.CreateCertificate(rand.Reader, tpl, tpl, &priv.PublicKey, priv)
	if err != nil {
		t.Fatalf("create cert: %v", err)
	}
	cert, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatalf("parse cert: %v", err)
	}
	return priv, cert
}

func setTestSAMLMiddlewareWithCert(t *testing.T, cert *x509.Certificate, entityID string) {
	t.Helper()
	prev := SAMLMiddleware
	SAMLMiddleware = &samlsp.Middleware{}
	SAMLMiddleware.ServiceProvider.IDPMetadata = &saml.EntityDescriptor{
		EntityID: entityID,
		IDPSSODescriptors: []saml.IDPSSODescriptor{
			{
				SSODescriptor: saml.SSODescriptor{
					RoleDescriptor: saml.RoleDescriptor{
						KeyDescriptors: []saml.KeyDescriptor{
							{
								KeyInfo: dsigKeyInfoFromCert(cert),
							},
						},
					},
				},
			},
		},
	}
	t.Cleanup(func() { SAMLMiddleware = prev })
}

func dsigKeyInfoFromCert(cert *x509.Certificate) saml.KeyInfo {
	if cert == nil {
		return saml.KeyInfo{}
	}
	return saml.KeyInfo{
		X509Data: saml.X509Data{
			X509Certificates: []saml.X509Certificate{
				{Data: base64.StdEncoding.EncodeToString(cert.Raw)},
			},
		},
	}
}

func testLogoutResponseXML(issuer string) string {
	return fmt.Sprintf(
		`<LogoutResponse xmlns="urn:oasis:names:tc:SAML:2.0:protocol" ID="_id123" Version="2.0" IssueInstant="2026-05-28T00:00:00Z"><Issuer xmlns="urn:oasis:names:tc:SAML:2.0:assertion">%s</Issuer><Status><StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/></Status></LogoutResponse>`,
		issuer,
	)
}

func TestCertificateThumbprintSHA1(t *testing.T) {
	_, cert := testSigningCert(t)
	got := certificateThumbprintSHA1(cert)
	if got == "" {
		t.Fatal("thumbprint should not be empty")
	}
	sum := sha1.Sum(cert.Raw)
	want := strings.ToUpper(hex.EncodeToString(sum[:]))
	if got != want {
		t.Fatalf("thumbprint mismatch: got=%s want=%s", got, want)
	}
}