package webui

import (
	"bytes"
	"crypto/x509"
	"encoding/pem"
	"io"
	"mime"
	"net/http"
	"net/mail"
	"strings"

	"mailgoserver/internal/db"
	"mailgoserver/internal/pgp"
	"mailgoserver/internal/smime"
)

// smimeInfo summarizes what unwrapSMIME found, for the message view to show as
// badges. A message can be both signed and encrypted (either nesting order); the
// zero value means "plain, no S/MIME involved."
type smimeInfo struct {
	Signed       bool
	SignatureOK  bool
	SignatureErr string
	SignerEmail  string
	Encrypted    bool
	Decrypted    bool
	DecryptErr   string
}

// pgpInfo summarizes what unwrapCrypto found about PGP encryption — parallel to
// smimeInfo, but PGP has no signing role in this codebase (S/MIME handles that), so
// there's no Signed/SignatureOK equivalent here.
type pgpInfo struct {
	Encrypted   bool
	Decrypted   bool
	DecryptErr  string
	NeedsUnlock bool
	Identities  []db.MailboxPGPIdentity
}

// entityHeaderNames are the only headers that belong to a MIME entity (as opposed to
// the message envelope) — see smime.Entity's doc comment.
var entityHeaderNames = []string{"Content-Type", "Content-Transfer-Encoding", "Content-Disposition"}

func isEntityHeader(name string) bool {
	for _, n := range entityHeaderNames {
		if strings.EqualFold(n, name) {
			return true
		}
	}
	return false
}

// tryDecryptWithIdentities attempts a pkcs7-mime decrypt using every S/MIME
// identity this mailbox holds — deliberately not inspecting the CMS RecipientInfo
// to figure out which identity a message targets; trying each key against
// smime.Decrypt is cheap and simple, and a mailbox realistically holds only a
// handful of identities. No passphrase/unlock step: S/MIME keys are stored plain.
func tryDecryptWithIdentities(entity smime.Entity, identities []db.MailboxSMIMEIdentity) (smime.Entity, bool) {
	for _, id := range identities {
		cert, err := smime.ParseCertPEM([]byte(id.CertPEM))
		if err != nil {
			continue
		}
		key, err := smime.ParseKeyPEM([]byte(id.KeyPEM))
		if err != nil {
			continue
		}
		if inner, err := smime.Decrypt(entity, cert, key); err == nil {
			return inner, true
		}
	}
	return smime.Entity{}, false
}

// tryPGPDecryptWithUnlockedIdentities mirrors tryDecryptWithUnlockedIdentities for
// PGP: tries every identity already unlocked in this session's key cache.
func (a *App) tryPGPDecryptWithUnlockedIdentities(token string, entity pgp.Entity, identities []db.MailboxPGPIdentity) (pgp.Entity, bool) {
	for _, id := range identities {
		unlocked, ok := a.pgpKeys.get(token, id.ID)
		if !ok {
			continue
		}
		if inner, err := pgp.DecryptEntity(entity, unlocked); err == nil {
			return inner, true
		}
	}
	return pgp.Entity{}, false
}

// unwrapCrypto strips any S/MIME signing and/or PGP encryption layers off raw (in
// whichever order they were applied — sign-then-encrypt or encrypt-then-sign, and
// either protocol can be outermost), returning a flat RFC822 message with the
// original envelope headers and the innermost plaintext MIME entity, ready for
// mailview.Parse. A verified S/MIME signature also gets the signer's certificate
// auto-captured into the mailbox's S/MIME contacts, the same way any mail client
// would on receiving a good signature (PGP has no signing role here, so no
// equivalent auto-capture on that side). Never returns an error for "this isn't
// S/MIME or PGP" — that's the common case, and raw is returned unchanged.
func (a *App) unwrapCrypto(r *http.Request, mailboxID int64, raw []byte) ([]byte, smimeInfo, pgpInfo) {
	msg, err := mail.ReadMessage(bytes.NewReader(raw))
	if err != nil {
		return raw, smimeInfo{}, pgpInfo{}
	}
	body, err := io.ReadAll(msg.Body)
	if err != nil {
		return raw, smimeInfo{}, pgpInfo{}
	}

	entity := smime.Entity{Body: body}
	for _, name := range entityHeaderNames {
		if v := msg.Header.Get(name); v != "" {
			entity.Headers = append(entity.Headers, name+": "+v)
		}
	}

	var sInfo smimeInfo
	var pInfo pgpInfo
	var envelopeHeaders []string
	for key, vals := range msg.Header {
		if isEntityHeader(key) {
			continue
		}
		for _, v := range vals {
			envelopeHeaders = append(envelopeHeaders, key+": "+v)
		}
	}

unwrap:
	for i := 0; i < 5; i++ { // cap against pathological nesting; real S/MIME/PGP never nests this deep
		ct := smime.HeaderValue(entity.Headers, "Content-Type")
		mediaType := strings.ToLower(strings.TrimSpace(strings.SplitN(ct, ";", 2)[0]))
		switch mediaType {
		case "multipart/signed":
			sInfo.Signed = true
			inner, signer, verr := smime.VerifySigned(entity)
			if verr != nil {
				sInfo.SignatureErr = verr.Error()
			} else {
				sInfo.SignatureOK = true
				if len(signer.EmailAddresses) > 0 {
					sInfo.SignerEmail = signer.EmailAddresses[0]
					a.captureSignerContact(mailboxID, sInfo.SignerEmail, signer)
				}
			}
			entity = inner
		case "application/pkcs7-mime":
			sInfo.Encrypted = true
			identities, ierr := a.DB.ListSMIMEIdentities(mailboxID)
			if ierr != nil {
				sInfo.DecryptErr = ierr.Error()
				break unwrap
			}
			if inner, ok := tryDecryptWithIdentities(entity, identities); ok {
				sInfo.Decrypted = true
				entity = inner
			} else {
				sInfo.DecryptErr = "no S/MIME certificate on file could decrypt this message"
				break unwrap
			}
		case "multipart/encrypted":
			_, params, perr := mime.ParseMediaType(ct)
			if perr != nil || !strings.EqualFold(params["protocol"], "application/pgp-encrypted") {
				break unwrap
			}
			pInfo.Encrypted = true
			pgpIdentities, ierr := a.DB.ListPGPIdentities(mailboxID)
			if ierr != nil {
				pInfo.DecryptErr = ierr.Error()
				break unwrap
			}
			if inner, ok := a.tryPGPDecryptWithUnlockedIdentities(sessionToken(r), pgp.Entity(entity), pgpIdentities); ok {
				pInfo.Decrypted = true
				entity = smime.Entity(inner)
			} else if len(pgpIdentities) == 0 {
				pInfo.DecryptErr = "no PGP key to decrypt with — set one up on the Certs page"
				break unwrap
			} else {
				pInfo.NeedsUnlock = true
				pInfo.Identities = pgpIdentities
				break unwrap
			}
		default:
			break unwrap
		}
	}

	var buf bytes.Buffer
	for _, h := range envelopeHeaders {
		buf.WriteString(h)
		buf.WriteString("\r\n")
	}
	for _, h := range entity.Headers {
		buf.WriteString(h)
		buf.WriteString("\r\n")
	}
	buf.WriteString("\r\n")
	buf.Write(entity.Body)
	return buf.Bytes(), sInfo, pInfo
}

func (a *App) captureSignerContact(mailboxID int64, email string, cert *x509.Certificate) {
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: cert.Raw})
	if err := a.DB.UpsertSMIMEContact(mailboxID, email, string(certPEM)); err != nil {
		a.Logger.Error("auto-capture signer cert for mailbox %d: %v", mailboxID, err)
	}
}