package dkim

import (
	"os"
	"strings"
	"testing"

	msgdkim "github.com/emersion/go-msgauth/dkim"
	"mailgoserver/internal/db"
)

func TestGenerateSignVerify(t *testing.T) {
	f, err := os.CreateTemp("", "dkim-test-*.db")
	if err != nil {
		t.Fatal(err)
	}
	f.Close()
	defer os.Remove(f.Name())

	database, err := db.Open(f.Name())
	if err != nil {
		t.Fatal(err)
	}
	defer database.Close()

	if _, err := database.Exec(`INSERT INTO esrv_domains (domain_name, is_active) VALUES ('example.com', 1)`); err != nil {
		t.Fatal(err)
	}

	mgr := New(database, 1024) // small key for test speed
	ok, err := mgr.GenerateDKIMKeypair("example.com", "sel1", false)
	if err != nil || !ok {
		t.Fatalf("GenerateDKIMKeypair: ok=%v err=%v", ok, err)
	}

	rec, err := mgr.GetDKIMPublicKeyRecord("example.com")
	if err != nil || rec == nil {
		t.Fatalf("GetDKIMPublicKeyRecord: %v %v", rec, err)
	}
	if rec.Name != "sel1._domainkey.example.com" || rec.Type != "TXT" || !strings.HasPrefix(rec.Value, `"v=DKIM1; k=rsa; p=`) {
		t.Fatalf("unexpected DNS record: %+v", rec)
	}

	msg := "From: sender@example.com\r\nTo: rcpt@example.org\r\nSubject: hi\r\nDate: Mon, 01 Jan 2024 00:00:00 +0000\r\nMessage-ID: <abc@example.com>\r\nMIME-Version: 1.0\r\nContent-Type: text/plain; charset=UTF-8\r\nContent-Transfer-Encoding: 7bit\r\n\r\nhello world\r\n"
	signed := mgr.Sign(msg, "example.com")
	if signed == msg {
		t.Fatal("Sign did not add a signature")
	}
	if !strings.HasPrefix(signed, "DKIM-Signature:") {
		t.Fatalf("expected DKIM-Signature as first header, got: %s", signed[:60])
	}

	// Verify against the key we just generated instead of a live DNS lookup
	// (example.com has no real TXT record for our test selector).
	verifications, err := msgdkim.VerifyWithOptions(strings.NewReader(signed), &msgdkim.VerifyOptions{
		LookupTXT: func(domain string) ([]string, error) {
			if domain != rec.Name {
				t.Fatalf("unexpected TXT lookup domain: %s (want %s)", domain, rec.Name)
			}
			return []string{strings.Trim(rec.Value, `"`)}, nil
		},
	})
	if err != nil {
		t.Fatalf("Verify error: %v", err)
	}
	if len(verifications) != 1 {
		t.Fatalf("expected 1 verification, got %d", len(verifications))
	}
	if verifications[0].Err != nil {
		t.Fatalf("verification failed: %v", verifications[0].Err)
	}

	// Re-signing must strip the old signature, not stack two.
	resigned := mgr.Sign(signed, "example.com")
	if strings.Count(resigned, "DKIM-Signature:") != 1 {
		t.Fatalf("expected exactly one DKIM-Signature after re-sign, got: %s", resigned)
	}
}
