package webui

import (
	"bytes"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"strings"
	"testing"

	"mailgoserver/internal/pgp"
)

// TestWebmailPGPGenerateAndDownload confirms a mailbox owner can generate a
// passphrase-protected PGP key, see it reflected on the Certs page, and download
// the public key.
func TestWebmailPGPGenerateAndDownload(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	domains, _ := app.DB.ListDomains()
	mailboxID := createTestMailboxWithPassword(t, app, "pgp1@example.com", domains[0].ID, "pgp-password-1!")
	cookie := webmailLoginSession(t, app, mailboxID)

	form := url.Values{"label": {"Work key"}, "passphrase": {"correct horse battery staple"}, "passphrase_confirm": {"correct horse battery staple"}}
	req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/pgp/identity/generate", strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.AddCookie(cookie)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusFound {
		t.Fatalf("generate: status=%d body=%s", rec.Code, rec.Body.String())
	}

	identities, err := app.DB.ListPGPIdentities(mailboxID)
	if err != nil || len(identities) != 1 {
		t.Fatalf("expected 1 PGP identity, got %d (err=%v)", len(identities), err)
	}
	identity := identities[0]
	if identity.Label != "Work key" {
		t.Errorf("label = %q, want %q", identity.Label, "Work key")
	}
	if identity.Fingerprint == "" {
		t.Error("expected a non-empty fingerprint")
	}
	if bytes.Contains([]byte(identity.PrivateKeyArmor), []byte("correct horse")) {
		t.Fatal("stored private key armor should not contain the plaintext passphrase")
	}

	pageReq := httptest.NewRequest(http.MethodGet, MailboxPrefix+"/certs", nil)
	pageReq.AddCookie(cookie)
	pageRec := httptest.NewRecorder()
	mux.ServeHTTP(pageRec, pageReq)
	if pageRec.Code != http.StatusOK || !strings.Contains(pageRec.Body.String(), "Work key") {
		t.Fatalf("expected the Certs page to show the PGP identity, status=%d body=%s", pageRec.Code, pageRec.Body.String())
	}
	if !strings.Contains(pageRec.Body.String(), "Locked") {
		t.Fatal("expected the PGP identity to show as locked before any unlock")
	}

	dlReq := httptest.NewRequest(http.MethodGet, MailboxPrefix+"/pgp/identity/"+strconv.FormatInt(identity.ID, 10)+"/download", nil)
	dlReq.AddCookie(cookie)
	dlRec := httptest.NewRecorder()
	mux.ServeHTTP(dlRec, dlReq)
	if dlRec.Code != http.StatusOK || dlRec.Body.String() != identity.PublicKeyArmor {
		t.Fatalf("expected downloaded key to match stored public key, status=%d", dlRec.Code)
	}
}

// TestWebmailPGPUnlockWrongPassphraseRejected confirms the unlock endpoint rejects
// an incorrect passphrase and caches nothing.
func TestWebmailPGPUnlockWrongPassphraseRejected(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	domains, _ := app.DB.ListDomains()
	mailboxID := createTestMailboxWithPassword(t, app, "pgp2@example.com", domains[0].ID, "pgp-password-1!")
	cookie := webmailLoginSession(t, app, mailboxID)

	form := url.Values{"passphrase": {"the real passphrase"}, "passphrase_confirm": {"the real passphrase"}}
	req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/pgp/identity/generate", strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.AddCookie(cookie)
	mux.ServeHTTP(httptest.NewRecorder(), req)
	identities, _ := app.DB.ListPGPIdentities(mailboxID)
	if len(identities) != 1 {
		t.Fatalf("expected 1 identity, got %d", len(identities))
	}

	unlock := func(passphrase string) *httptest.ResponseRecorder {
		f := url.Values{"identity_id": {strconv.FormatInt(identities[0].ID, 10)}, "passphrase": {passphrase}}
		r := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/pgp/unlock", strings.NewReader(f.Encode()))
		r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		r.AddCookie(cookie)
		rec := httptest.NewRecorder()
		mux.ServeHTTP(rec, r)
		return rec
	}

	if rec := unlock("wrong passphrase"); rec.Code != http.StatusFound {
		t.Fatalf("status=%d", rec.Code)
	}
	if _, ok := app.pgpKeys.get(sessionTokenFromCookie(cookie), identities[0].ID); ok {
		t.Fatal("expected no key cached after a wrong-passphrase unlock attempt")
	}

	if rec := unlock("the real passphrase"); rec.Code != http.StatusFound {
		t.Fatalf("status=%d", rec.Code)
	}
	if _, ok := app.pgpKeys.get(sessionTokenFromCookie(cookie), identities[0].ID); !ok {
		t.Fatal("expected the key cached after the correct passphrase")
	}
}

// TestWebmailPGPContactAddAndRemove confirms a contact's public key can be added
// (validated as a real key), listed, and removed again.
func TestWebmailPGPContactAddAndRemove(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	domains, _ := app.DB.ListDomains()
	mailboxID := createTestMailboxWithPassword(t, app, "pgp3@example.com", domains[0].ID, "pgp-password-1!")
	cookie := webmailLoginSession(t, app, mailboxID)

	contactPub, _, err := pgp.GenerateKeyPair("contact@other.example", "contact passphrase")
	if err != nil {
		t.Fatal(err)
	}

	var buf bytes.Buffer
	mw := multipart.NewWriter(&buf)
	mw.WriteField("email", "contact@other.example")
	mw.WriteField("label", "Other")
	fw, err := mw.CreateFormFile("key_file", "contact.asc")
	if err != nil {
		t.Fatal(err)
	}
	fw.Write(contactPub)
	mw.Close()

	req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/pgp/contacts/add", &buf)
	req.Header.Set("Content-Type", mw.FormDataContentType())
	req.AddCookie(cookie)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusFound {
		t.Fatalf("add contact: status=%d body=%s", rec.Code, rec.Body.String())
	}

	contacts, err := app.DB.ListPGPContacts(mailboxID)
	if err != nil || len(contacts) != 1 || contacts[0].Email != "contact@other.example" || contacts[0].Fingerprint == "" {
		t.Fatalf("expected 1 contact with a fingerprint, got %+v (err=%v)", contacts, err)
	}

	rmReq := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/pgp/contacts/"+strconv.FormatInt(contacts[0].ID, 10)+"/remove", nil)
	rmReq.AddCookie(cookie)
	rmRec := httptest.NewRecorder()
	mux.ServeHTTP(rmRec, rmReq)
	if rmRec.Code != http.StatusFound {
		t.Fatalf("remove contact: status=%d", rmRec.Code)
	}
	remaining, err := app.DB.ListPGPContacts(mailboxID)
	if err != nil || len(remaining) != 0 {
		t.Fatalf("expected no contacts left, got %d (err=%v)", len(remaining), err)
	}
}

// TestWebmailPGPAddContactRejectsGarbage confirms an upload that isn't a valid PGP
// public key is rejected rather than silently stored.
func TestWebmailPGPAddContactRejectsGarbage(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	domains, _ := app.DB.ListDomains()
	mailboxID := createTestMailboxWithPassword(t, app, "pgp4@example.com", domains[0].ID, "pgp-password-1!")
	cookie := webmailLoginSession(t, app, mailboxID)

	var buf bytes.Buffer
	mw := multipart.NewWriter(&buf)
	mw.WriteField("email", "nope@example.com")
	fw, err := mw.CreateFormFile("key_file", "notakey.asc")
	if err != nil {
		t.Fatal(err)
	}
	fw.Write([]byte("this is not a pgp key"))
	mw.Close()

	req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/pgp/contacts/add", &buf)
	req.Header.Set("Content-Type", mw.FormDataContentType())
	req.AddCookie(cookie)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusFound {
		t.Fatalf("status=%d", rec.Code)
	}

	contacts, err := app.DB.ListPGPContacts(mailboxID)
	if err != nil || len(contacts) != 0 {
		t.Fatalf("expected the invalid key rejected, got %d contacts (err=%v)", len(contacts), err)
	}
}

// TestWebmailPGPImportAlreadyEncryptedKey confirms importing an existing
// passphrase-protected armored key (as if exported from GnuPG) works end to end
// through the real HTTP import endpoint.
func TestWebmailPGPImportAlreadyEncryptedKey(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	domains, _ := app.DB.ListDomains()
	mailboxID := createTestMailboxWithPassword(t, app, "pgp5@example.com", domains[0].ID, "pgp-password-1!")
	cookie := webmailLoginSession(t, app, mailboxID)

	_, privArmor, err := pgp.GenerateKeyPair("imported@example.com", "existing passphrase")
	if err != nil {
		t.Fatal(err)
	}

	var buf bytes.Buffer
	mw := multipart.NewWriter(&buf)
	mw.WriteField("label", "Imported")
	mw.WriteField("passphrase", "existing passphrase")
	fw, err := mw.CreateFormFile("key_file", "imported.asc")
	if err != nil {
		t.Fatal(err)
	}
	fw.Write(privArmor)
	mw.Close()

	req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/pgp/identity/import", &buf)
	req.Header.Set("Content-Type", mw.FormDataContentType())
	req.AddCookie(cookie)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusFound {
		t.Fatalf("import: status=%d body=%s", rec.Code, rec.Body.String())
	}

	identities, err := app.DB.ListPGPIdentities(mailboxID)
	if err != nil || len(identities) != 1 {
		t.Fatalf("expected 1 identity, got %d (err=%v)", len(identities), err)
	}
	if identities[0].Email != "imported@example.com" {
		t.Errorf("expected the email extracted from the key itself, got %q", identities[0].Email)
	}
}