package httpapi

import (
	"bytes"
	"context"
	"crypto"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"basicstack.de/proofedmesh/services/signing-service/internal/policy"
	"basicstack.de/proofedmesh/services/signing-service/internal/registry"
	"basicstack.de/proofedmesh/services/signing-service/internal/signer"
	"basicstack.de/proofedmesh/services/signing-service/internal/token"
)

// mockPocketID issues RS256 JWTs and serves OIDC discovery + JWKS over
// httptest so the signing-service's Pocket-ID validator thinks it is
// talking to real Pocket ID. The suite exercises OIDC-Discovery +
// JWK-Cache + kid rotation.
type mockPocketID struct {
	srv  *httptest.Server
	keys map[string]*rsa.PrivateKey // active kid → key
	// activeKIDs returns the order in which keys appear in the JWKS.
	activeKIDs []string
}

func newMockPocketID(t *testing.T) *mockPocketID {
	t.Helper()
	m := &mockPocketID{keys: map[string]*rsa.PrivateKey{}}
	k1, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	m.keys["k1"] = k1
	m.activeKIDs = []string{"k1"}

	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, `{
			"issuer":"%s",
			"authorization_endpoint":"%s/auth",
			"token_endpoint":"%s/token",
			"jwks_uri":"%s/jwks",
			"id_token_signing_alg_values_supported":["RS256"]
		}`, m.issuer(), m.issuer(), m.issuer(), m.issuer())
	})
	mux.HandleFunc("/jwks", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		var keys []map[string]any
		for _, kid := range m.activeKIDs {
			k := m.keys[kid]
			n := base64.RawURLEncoding.EncodeToString(k.N.Bytes())
			e := base64.RawURLEncoding.EncodeToString(big.NewInt(int64(k.E)).Bytes())
			keys = append(keys, map[string]any{
				"kty": "RSA", "use": "sig", "kid": kid, "alg": "RS256", "n": n, "e": e,
			})
		}
		_ = json.NewEncoder(w).Encode(map[string]any{"keys": keys})
	})
	m.srv = httptest.NewServer(mux)
	return m
}

func (m *mockPocketID) issuer() string { return m.srv.URL }

func (m *mockPocketID) close() { m.srv.Close() }

func (m *mockPocketID) rotateAddKey(t *testing.T, kid string) {
	t.Helper()
	k, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	m.keys[kid] = k
	m.activeKIDs = append(m.activeKIDs, kid)
}

func (m *mockPocketID) signToken(t *testing.T, kid, sub, aud string) string {
	t.Helper()
	k, ok := m.keys[kid]
	if !ok {
		t.Fatalf("no key %q", kid)
	}
	header := map[string]any{"alg": "RS256", "typ": "JWT", "kid": kid}
	claims := map[string]any{
		"iss": m.issuer(),
		"sub": sub,
		"aud": aud,
		"iat": time.Now().Unix(),
		"exp": time.Now().Add(10 * time.Minute).Unix(),
	}
	hJ, _ := json.Marshal(header)
	cJ, _ := json.Marshal(claims)
	signingInput := base64.RawURLEncoding.EncodeToString(hJ) + "." + base64.RawURLEncoding.EncodeToString(cJ)
	h := sha256.Sum256([]byte(signingInput))
	sig, err := rsa.SignPKCS1v15(rand.Reader, k, crypto.SHA256, h[:])
	if err != nil {
		t.Fatal(err)
	}
	return signingInput + "." + base64.RawURLEncoding.EncodeToString(sig)
}

// ----- fake TokenReviewer -----

func newFakeTokenReviewer(table map[string]token.K8sIdentity) token.TokenReviewer {
	return token.StaticTokenReviewer(table)
}

// ----- fake registry seeded with Phase-1 Phase 1 agents -----

func newPhase1Registry() registry.Resolver {
	return registry.NewStatic([]registry.Agent{
		{
			AgentID:           "agent-approval-service",
			PocketIDClient:    "agent-approval-service",
			K8sNamespace:      "agent-trust",
			K8sServiceAccount: "approval-service",
			AllowedActions:    []string{"expense.submit"},
		},
		{
			AgentID:           "agent-testagent-transform",
			PocketIDClient:    "agent-testagent-transform",
			K8sNamespace:      "agent-trust",
			K8sServiceAccount: "testagent-transform",
			AllowedActions:    []string{"expense.normalize", "expense.map_to_sink"},
		},
	})
}

func loadPolicy(t *testing.T) policy.StructuredDecider {
	t.Helper()
	p := filepath.Join("..", "..", "..", "..", "policies", "sinks", "signing_service.rego")
	opa, err := policy.NewOPAEvaluator(context.Background(), p)
	if err != nil {
		t.Fatalf("load policy: %v", err)
	}
	return opa
}

// ----- harness wiring everything together -----

type prodHarness struct {
	hsrv            *httptest.Server
	mockPID         *mockPocketID
	approvalPub     ed25519.PublicKey
	transformPub    ed25519.PublicKey
	workloadToken   map[string]token.K8sIdentity
	approvalBearer  string
	transformBearer string
}

func newProdHarness(t *testing.T) *prodHarness {
	t.Helper()
	m := newMockPocketID(t)
	t.Cleanup(m.close)

	ctx := context.Background()
	pv, err := token.NewPocketIDValidator(ctx, m.issuer(), "trust-api")
	if err != nil {
		t.Fatal(err)
	}

	signerBackend := signer.NewInMemory()
	_, approvalPriv, _ := ed25519.GenerateKey(rand.Reader)
	signerBackend.AddKey("trust-api/approval-signer-2026-q4", approvalPriv)
	_, transformPriv, _ := ed25519.GenerateKey(rand.Reader)
	signerBackend.AddKey("trust-api/agent-testagent-transform-2026-q4", transformPriv)

	wl := map[string]token.K8sIdentity{
		"wl-approval":  {Namespace: "agent-trust", ServiceAccount: "approval-service"},
		"wl-transform": {Namespace: "agent-trust", ServiceAccount: "testagent-transform"},
	}

	srv := Server{
		Signer:           signerBackend,
		Validator:        pv,
		TokenReviewer:    newFakeTokenReviewer(wl),
		Registry:         newPhase1Registry(),
		StructuredPolicy: loadPolicy(t),
		Policy:           policy.AllowAll{}, // unused in the structured path
		Now:              func() time.Time { return time.Date(2026, 10, 25, 20, 0, 0, 0, time.UTC) },
	}
	mux := http.NewServeMux()
	srv.Mount(mux)
	hsrv := httptest.NewServer(mux)
	t.Cleanup(hsrv.Close)

	return &prodHarness{
		hsrv:            hsrv,
		mockPID:         m,
		approvalPub:     approvalPriv.Public().(ed25519.PublicKey),
		transformPub:    transformPriv.Public().(ed25519.PublicKey),
		workloadToken:   wl,
		approvalBearer:  m.signToken(t, "k1", "agent-approval-service", "trust-api"),
		transformBearer: m.signToken(t, "k1", "agent-testagent-transform", "trust-api"),
	}
}

func sign(h *prodHarness, bearer, workloadToken, body string) (*http.Response, error) {
	req, _ := http.NewRequest("POST", h.hsrv.URL+"/api/v1/sign/dsse", bytes.NewBufferString(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+bearer)
	if workloadToken != "" {
		req.Header.Set(token.WorkloadTokenHeader, workloadToken)
	}
	return http.DefaultClient.Do(req)
}

const approvalStatement = `{
  "key_ref":"trust-api/approval-signer-2026-q4",
  "statement":{
    "_type":"https://in-toto.io/Statement/v1",
    "subject":[{"name":"approval/APR-1","digest":{"sha256":"abcdef"}}],
    "predicateType":"https://basicstack.de/attestations/approval/v1",
    "predicate":{"approval_id":"APR-1"}
  }
}`

const transformStatement = `{
  "key_ref":"trust-api/agent-testagent-transform-2026-q4",
  "statement":{
    "_type":"https://in-toto.io/Statement/v1",
    "subject":[{"name":"transform/TR-1","digest":{"sha256":"abcdef"}}],
    "predicateType":"https://basicstack.de/attestations/transform/v1",
    "predicate":{"transform_id":"TR-1"}
  }
}`

// --- happy paths ---

func TestProd_HappyPath_ApprovalService(t *testing.T) {
	h := newProdHarness(t)
	resp, err := sign(h, h.approvalBearer, "wl-approval", approvalStatement)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := jsonBody(resp)
		t.Fatalf("status %d: %s", resp.StatusCode, b)
	}
	assertSignatureVerifies(t, resp, h.approvalPub, "application/vnd.in-toto+json")
}

func TestProd_HappyPath_TestAgentTransform(t *testing.T) {
	h := newProdHarness(t)
	resp, err := sign(h, h.transformBearer, "wl-transform", transformStatement)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := jsonBody(resp)
		t.Fatalf("status %d: %s", resp.StatusCode, b)
	}
	assertSignatureVerifies(t, resp, h.transformPub, "application/vnd.in-toto+json")
}

// --- denies ---

func TestProd_DeniesOnRegistryMismatch(t *testing.T) {
	h := newProdHarness(t)
	// Pocket-ID says Approval Service, but the X-Workload-Token is the
	// transform agent's. Registry lookups land on different rows → 403.
	resp, _ := sign(h, h.approvalBearer, "wl-transform", approvalStatement)
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusForbidden {
		t.Fatalf("expected 403 got %d", resp.StatusCode)
	}
	var e errResponse
	_ = json.NewDecoder(resp.Body).Decode(&e)
	if e.Error.Code != "signing_denied" {
		t.Errorf("expected signing_denied code, got %q (%s)", e.Error.Code, e.Error.Message)
	}
}

func TestProd_DeniesOnPolicyWrongKeyRef(t *testing.T) {
	h := newProdHarness(t)
	// Approval service identity on both sides, but asks for the transform
	// agent's key_ref. Policy must deny.
	body := strings.ReplaceAll(approvalStatement,
		"trust-api/approval-signer-2026-q4",
		"trust-api/agent-testagent-transform-2026-q4")
	resp, _ := sign(h, h.approvalBearer, "wl-approval", body)
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusForbidden {
		t.Fatalf("expected 403 got %d", resp.StatusCode)
	}
}

func TestProd_RejectsMissingWorkloadToken(t *testing.T) {
	h := newProdHarness(t)
	resp, _ := sign(h, h.approvalBearer, "", approvalStatement)
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("expected 401 got %d", resp.StatusCode)
	}
}

func TestProd_RejectsBadWorkloadToken(t *testing.T) {
	h := newProdHarness(t)
	resp, _ := sign(h, h.approvalBearer, "wl-does-not-exist", approvalStatement)
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("expected 401 got %d", resp.StatusCode)
	}
}

func TestProd_RejectsBadPocketIDToken(t *testing.T) {
	h := newProdHarness(t)
	resp, _ := sign(h, "not.a.jwt", "wl-approval", approvalStatement)
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("expected 401 got %d", resp.StatusCode)
	}
}

// --- kid rotation ---

func TestProd_PocketID_KidRotation(t *testing.T) {
	h := newProdHarness(t)
	// Baseline works with k1.
	resp, _ := sign(h, h.approvalBearer, "wl-approval", approvalStatement)
	resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("baseline: expected 200 got %d", resp.StatusCode)
	}
	// Rotate: new kid k2 appears in JWKS. A token signed with k2 must verify
	// after the validator refreshes the JWKS (go-oidc does this on kid miss).
	h.mockPID.rotateAddKey(t, "k2")
	bearer := h.mockPID.signToken(t, "k2", "agent-approval-service", "trust-api")
	resp2, _ := sign(h, bearer, "wl-approval", approvalStatement)
	defer resp2.Body.Close()
	if resp2.StatusCode != http.StatusOK {
		b, _ := jsonBody(resp2)
		t.Fatalf("after rotation: expected 200 got %d: %s", resp2.StatusCode, b)
	}
}

// --- helpers ---

func jsonBody(resp *http.Response) (string, error) {
	var buf bytes.Buffer
	_, err := buf.ReadFrom(resp.Body)
	return buf.String(), err
}

func assertSignatureVerifies(t *testing.T, resp *http.Response, pub ed25519.PublicKey, payloadType string) {
	t.Helper()
	var out signResponse
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		t.Fatalf("decode: %v", err)
	}
	payload, err := base64.RawURLEncoding.DecodeString(out.Envelope.Payload)
	if err != nil {
		t.Fatalf("decode payload: %v", err)
	}
	pae := []byte(fmt.Sprintf("DSSEv1 %d %s %d ", len(payloadType), payloadType, len(payload)))
	pae = append(pae, payload...)
	sig, err := base64.RawURLEncoding.DecodeString(out.Envelope.Signatures[0].Sig)
	if err != nil {
		t.Fatalf("decode sig: %v", err)
	}
	if !ed25519.Verify(pub, pae, sig) {
		t.Fatal("returned signature does not verify")
	}
}