package identityprovider

import (
	"context"
	"crypto"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"vpn-platform/backend/internal/model"
	"vpn-platform/backend/internal/store"
)

type transportFunc func(*http.Request) (*http.Response, error)

func (f transportFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) }

func TestFeishuExchangeUsesJSONAndRejectsUnexpectedIdentity(t *testing.T) {
	previous := upstreamClient
	defer func() { upstreamClient = previous }()
	upstreamClient = &http.Client{Transport: transportFunc(func(req *http.Request) (*http.Response, error) {
		if req.URL.String() == feishuToken {
			if req.Header.Get("Content-Type") != "application/json" {
				t.Fatal("Feishu token exchange must use JSON")
			}
			var body map[string]string
			if json.NewDecoder(req.Body).Decode(&body) != nil || body["client_id"] != "app-id" || body["redirect_uri"] != "https://vpn.example.com/id/feishu/callback" {
				t.Fatal("wrong Feishu exchange request")
			}
			return &http.Response{StatusCode: 200, Body: io.NopCloser(strings.NewReader(`{"code":0,"access_token":"feishu-token"}`)), Header: make(http.Header)}, nil
		}
		if req.URL.String() != feishuUserInfo || req.Header.Get("Authorization") != "Bearer feishu-token" {
			t.Fatal("unexpected Feishu user request")
		}
		return &http.Response{StatusCode: 200, Body: io.NopCloser(strings.NewReader(`{"code":0,"data":{"open_id":"ou_1","tenant_key":"tenant_1","name":"Alice"}}`)), Header: make(http.Header)}, nil
	})}
	token, err := postFeishuToken(context.Background(), "app-id", "app-secret", "feishu-code", "https://vpn.example.com/id/feishu/callback")
	if err != nil || token != "feishu-token" {
		t.Fatalf("exchange: %q %v", token, err)
	}
	user, tenant, err := fetchFeishuUser(context.Background(), token)
	if err != nil || tenant != "tenant_1" || user.Subject != "tenant_1:ou_1" {
		t.Fatalf("identity: %+v %s %v", user, tenant, err)
	}
}

func TestFeishuExchangeRejectsMalformedSuccess(t *testing.T) {
	previous := upstreamClient
	defer func() { upstreamClient = previous }()
	upstreamClient = &http.Client{Transport: transportFunc(func(req *http.Request) (*http.Response, error) {
		return &http.Response{StatusCode: 200, Body: io.NopCloser(strings.NewReader(`{"access_token":"fake","data":{"open_id":"ou_1","tenant_key":"tenant_1"}}`)), Header: make(http.Header)}, nil
	})}
	if _, err := postFeishuToken(context.Background(), "id", "secret", "code", "https://vpn.example.com/id/feishu/callback"); err == nil {
		t.Fatal("missing Feishu success code accepted")
	}
	if _, _, err := fetchFeishuUser(context.Background(), "fake"); err == nil {
		t.Fatal("missing Feishu user response code accepted")
	}
}

func TestGrantSingleUseAndSignedToken(t *testing.T) {
	dir := t.TempDir()
	t.Setenv("VPN_DATA_DIR", dir)
	if err := store.Init(dir); err != nil {
		t.Fatal(err)
	}
	code, err := randomToken()
	if err != nil {
		t.Fatal(err)
	}
	grant := codeData{requestState: requestState{ClientID: "opsx-vpn", Redirect: "https://vpn.example.com/oauth2/callback", Nonce: "nonce-1"}, Subject: "tenant:ou_1"}
	if err := saveGrant("code", code, grant, time.Minute); err != nil {
		t.Fatal(err)
	}
	var used codeData
	if err := consumeGrant("code", code, &used); err != nil {
		t.Fatal(err)
	}
	if err := consumeGrant("code", code, &used); err == nil {
		t.Fatal("grant replay accepted")
	}
	key, kid, err := signingKey()
	if err != nil {
		t.Fatal(err)
	}
	jwt, err := signIDToken(key, kid, "https://vpn.example.com/id", model.IdentityProvider{ClientID: "opsx-vpn"}, used)
	if err != nil {
		t.Fatal(err)
	}
	parts := strings.Split(jwt, ".")
	if len(parts) != 3 {
		t.Fatal("invalid JWT")
	}
	signature, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		t.Fatal(err)
	}
	hash := sha256.Sum256([]byte(parts[0] + "." + parts[1]))
	if err := rsa.VerifyPKCS1v15(&key.PublicKey, crypto.SHA256, hash[:], signature); err != nil {
		t.Fatal(err)
	}
	claimsRaw, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		t.Fatal(err)
	}
	var claims map[string]any
	if json.Unmarshal(claimsRaw, &claims) != nil || claims["sub"] != "tenant:ou_1" || claims["nonce"] != "nonce-1" || claims["aud"] != "opsx-vpn" {
		t.Fatalf("unexpected claims: %v", claims)
	}
	if info, err := filepath.Glob(filepath.Join(dir, "identity", "*.pem")); err != nil || len(info) != 1 {
		t.Fatal("missing signing key")
	}
}

func TestRedirectIsExactAndNoOpenRedirect(t *testing.T) {
	if !validateRedirect("https://vpn.example.com/oauth2/callback", "vpn.example.com") {
		t.Fatal("valid callback rejected")
	}
	for _, bad := range []string{"http://vpn.example.com/oauth2/callback", "https://vpn.example.com.evil/oauth2/callback", "https://vpn.example.com/oauth2/callback/extra", "https://vpn.example.com/oauth2/callback?next=evil"} {
		if validateRedirect(bad, "vpn.example.com") {
			t.Errorf("accepted callback %q", bad)
		}
	}
	if err := consumeGrant("state", "short", new(requestState)); err == nil {
		t.Fatal("accepted short state")
	}
}