package webui

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

	"mailgoserver/internal/smime"
)

// TestWebmailSMIMEGenerateAndDownload confirms a mailbox owner can generate a
// self-signed identity (no passphrase — S/MIME keys are stored plain, since S/MIME
// is sign-only here), see it reflected on the page, and download the public
// certificate.
func TestWebmailSMIMEGenerateAndDownload(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	domains, _ := app.DB.ListDomains()
	mailboxID := createTestMailboxWithPassword(t, app, "smime1@example.com", domains[0].ID, "smime-password-1!")
	cookie := webmailLoginSession(t, app, mailboxID)

	req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/smime/identity/generate", nil)
	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.ListSMIMEIdentities(mailboxID)
	if err != nil || len(identities) != 1 {
		t.Fatalf("expected 1 identity stored, got %d (err=%v)", len(identities), err)
	}
	identity := identities[0]
	if !strings.Contains(identity.KeyPEM, "PRIVATE KEY") {
		t.Fatal("expected the stored key to be a usable plain PEM")
	}
	cert, err := smime.ParseCertPEM([]byte(identity.CertPEM))
	if err != nil {
		t.Fatalf("stored cert doesn't parse: %v", err)
	}
	if len(cert.EmailAddresses) != 1 || cert.EmailAddresses[0] != "smime1@example.com" {
		t.Fatalf("unexpected cert EmailAddresses: %v", cert.EmailAddresses)
	}

	pageReq := httptest.NewRequest(http.MethodGet, MailboxPrefix+"/certs", nil)
	pageReq.AddCookie(cookie)
	pageRec := httptest.NewRecorder()
	mux.ServeHTTP(pageRec, pageReq)
	if pageRec.Code != http.StatusOK {
		t.Fatalf("expected the certs page to render, status=%d body=%s", pageRec.Code, pageRec.Body.String())
	}

	dlReq := httptest.NewRequest(http.MethodGet, MailboxPrefix+"/smime/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.CertPEM {
		t.Fatalf("expected downloaded cert to match stored cert, status=%d", dlRec.Code)
	}
}

// TestWebmailSMIMEMultipleIdentities confirms a mailbox can hold more than one
// identity at once.
func TestWebmailSMIMEMultipleIdentities(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	domains, _ := app.DB.ListDomains()
	mailboxID := createTestMailboxWithPassword(t, app, "smime-multi@example.com", domains[0].ID, "smime-password-1!")
	cookie := webmailLoginSession(t, app, mailboxID)

	for i := 0; i < 2; i++ {
		req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/smime/identity/generate", nil)
		req.AddCookie(cookie)
		mux.ServeHTTP(httptest.NewRecorder(), req)
	}

	identities, err := app.DB.ListSMIMEIdentities(mailboxID)
	if err != nil || len(identities) != 2 {
		t.Fatalf("expected 2 identities held simultaneously, got %d (err=%v)", len(identities), err)
	}
}

// TestWebmailSMIMERemoveIdentity confirms removal actually deletes the DB row.
func TestWebmailSMIMERemoveIdentity(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	domains, _ := app.DB.ListDomains()
	mailboxID := createTestMailboxWithPassword(t, app, "smime2@example.com", domains[0].ID, "smime-password-1!")
	cookie := webmailLoginSession(t, app, mailboxID)

	genReq := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/smime/identity/generate", nil)
	genReq.AddCookie(cookie)
	mux.ServeHTTP(httptest.NewRecorder(), genReq)
	identities, _ := app.DB.ListSMIMEIdentities(mailboxID)
	if len(identities) != 1 {
		t.Fatalf("expected 1 identity, got %d", len(identities))
	}

	rmReq := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/smime/identity/"+strconv.FormatInt(identities[0].ID, 10)+"/remove", nil)
	rmReq.AddCookie(cookie)
	rmRec := httptest.NewRecorder()
	mux.ServeHTTP(rmRec, rmReq)
	if rmRec.Code != http.StatusFound {
		t.Fatalf("remove: status=%d", rmRec.Code)
	}

	remaining, err := app.DB.ListSMIMEIdentities(mailboxID)
	if err != nil || len(remaining) != 0 {
		t.Fatalf("expected identity gone, got %d (err=%v)", len(remaining), err)
	}
}

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

	contactCertPEM, _, err := smime.GenerateSelfSigned("contact@other.example", smime.DefaultValidity)
	if err != nil {
		t.Fatal(err)
	}

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

	req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/smime/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.ListSMIMEContacts(mailboxID)
	if err != nil || len(contacts) != 1 || contacts[0].Email != "contact@other.example" {
		t.Fatalf("expected 1 contact, got %+v (err=%v)", contacts, err)
	}

	rmReq := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/smime/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.ListSMIMEContacts(mailboxID)
	if err != nil || len(remaining) != 0 {
		t.Fatalf("expected no contacts left, got %d (err=%v)", len(remaining), err)
	}
}

// TestWebmailSMIMEAddContactRejectsGarbage confirms an upload that isn't a valid
// certificate is rejected rather than silently stored.
func TestWebmailSMIMEAddContactRejectsGarbage(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	domains, _ := app.DB.ListDomains()
	mailboxID := createTestMailboxWithPassword(t, app, "smime4@example.com", domains[0].ID, "smime-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("cert_file", "notacert.pem")
	if err != nil {
		t.Fatal(err)
	}
	fw.Write([]byte("this is not a certificate"))
	mw.Close()

	req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/smime/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.ListSMIMEContacts(mailboxID)
	if err != nil || len(contacts) != 0 {
		t.Fatalf("expected the invalid cert rejected, got %d contacts (err=%v)", len(contacts), err)
	}
}

// TestWebmailSMIMEScopedToOwnMailbox confirms one mailbox owner can't remove
// another's contact by guessing its ID, and can't reach another's identity.
func TestWebmailSMIMEScopedToOwnMailbox(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()
	domains, _ := app.DB.ListDomains()
	victimID := createTestMailboxWithPassword(t, app, "victim3@example.com", domains[0].ID, "victim-password-1!")
	attackerID := createTestMailboxWithPassword(t, app, "attacker3@example.com", domains[0].ID, "attacker-password-1!")

	if err := app.DB.UpsertSMIMEContact(victimID, "someone@example.com", "irrelevant-pem-for-this-test"); err != nil {
		t.Fatal(err)
	}
	contacts, _ := app.DB.ListSMIMEContacts(victimID)

	attackerCookie := webmailLoginSession(t, app, attackerID)
	req := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/smime/contacts/"+strconv.FormatInt(contacts[0].ID, 10)+"/remove", nil)
	req.AddCookie(attackerCookie)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusFound {
		t.Fatalf("status=%d", rec.Code)
	}

	stillThere, err := app.DB.ListSMIMEContacts(victimID)
	if err != nil || len(stillThere) != 1 {
		t.Fatalf("expected the victim's contact untouched, got %d (err=%v)", len(stillThere), err)
	}

	// Attacker downloading the victim's identity by guessing its ID should 404, not
	// leak the victim's cert.
	victimCookie := webmailLoginSession(t, app, victimID)
	genReq := httptest.NewRequest(http.MethodPost, MailboxPrefix+"/smime/identity/generate", nil)
	genReq.AddCookie(victimCookie)
	mux.ServeHTTP(httptest.NewRecorder(), genReq)
	victimIdentities, _ := app.DB.ListSMIMEIdentities(victimID)
	if len(victimIdentities) != 1 {
		t.Fatalf("expected 1 victim identity, got %d", len(victimIdentities))
	}

	dlReq := httptest.NewRequest(http.MethodGet, MailboxPrefix+"/smime/identity/"+strconv.FormatInt(victimIdentities[0].ID, 10)+"/download", nil)
	dlReq.AddCookie(attackerCookie)
	dlRec := httptest.NewRecorder()
	mux.ServeHTTP(dlRec, dlReq)
	if dlRec.Code != http.StatusNotFound {
		t.Fatalf("expected 404 for an attacker guessing another mailbox's identity ID, got %d", dlRec.Code)
	}
}

func sessionTokenFromCookie(c *http.Cookie) string { return c.Value }