package pgp

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/textproto"
	"strings"

	"github.com/ProtonMail/go-crypto/openpgp"
	"github.com/ProtonMail/go-crypto/openpgp/armor"
)

// Entity is a MIME entity: its own part-level headers plus its body — deliberately
// the same shape as smime.Entity, so webui's compose/read handlers can pass the same
// value between either package's Encrypt/Decrypt without conversion glue.
type Entity struct {
	Headers []string
	Body    []byte
}

func (e Entity) bytes() []byte {
	var buf bytes.Buffer
	for _, h := range e.Headers {
		buf.WriteString(h)
		buf.WriteString("\r\n")
	}
	buf.WriteString("\r\n")
	buf.Write(e.Body)
	return buf.Bytes()
}

// parseEntity splits raw bytes (headers, a blank line, then body) back into an
// Entity — used to recover the inner MIME entity after Decrypt.
func parseEntity(raw []byte) Entity {
	idx := bytes.Index(raw, []byte("\r\n\r\n"))
	sep := 4
	if idx < 0 {
		idx = bytes.Index(raw, []byte("\n\n"))
		sep = 2
	}
	if idx < 0 {
		return Entity{Body: raw}
	}
	var headers []string
	for _, line := range strings.Split(string(raw[:idx]), "\n") {
		line = strings.TrimRight(line, "\r")
		if line == "" {
			continue
		}
		headers = append(headers, line)
	}
	return Entity{Headers: headers, Body: raw[idx+sep:]}
}

func headerValue(headers []string, name string) string {
	for _, h := range headers {
		if i := strings.Index(h, ":"); i >= 0 && strings.EqualFold(strings.TrimSpace(h[:i]), name) {
			return strings.TrimSpace(h[i+1:])
		}
	}
	return ""
}

// EncryptEntity wraps entity's bytes as an RFC 3156 PGP/MIME multipart/encrypted
// structure, encrypted to recipients. Pass every recipient's public key, including
// the sender's own, so a copy kept in Sent stays readable — mirrors
// smime.Encrypt's same convention.
func EncryptEntity(entity Entity, recipients []*openpgp.Entity) (Entity, error) {
	if len(recipients) == 0 {
		return Entity{}, errors.New("pgp: no recipient keys provided")
	}

	var armored bytes.Buffer
	aw, err := armor.Encode(&armored, "PGP MESSAGE", nil)
	if err != nil {
		return Entity{}, fmt.Errorf("pgp: encrypt: %w", err)
	}
	pt, err := openpgp.Encrypt(aw, recipients, nil, nil, defaultConfig())
	if err != nil {
		return Entity{}, fmt.Errorf("pgp: encrypt: %w", err)
	}
	if _, err := pt.Write(entity.bytes()); err != nil {
		return Entity{}, fmt.Errorf("pgp: encrypt: %w", err)
	}
	if err := pt.Close(); err != nil {
		return Entity{}, fmt.Errorf("pgp: encrypt: %w", err)
	}
	if err := aw.Close(); err != nil {
		return Entity{}, fmt.Errorf("pgp: encrypt: %w", err)
	}

	var body bytes.Buffer
	mw := multipart.NewWriter(&body)
	ctrlPart, err := mw.CreatePart(textproto.MIMEHeader{
		"Content-Type":              {"application/pgp-encrypted"},
		"Content-Transfer-Encoding": {"7bit"},
	})
	if err != nil {
		return Entity{}, err
	}
	if _, err := ctrlPart.Write([]byte("Version: 1\r\n")); err != nil {
		return Entity{}, err
	}
	dataPart, err := mw.CreatePart(textproto.MIMEHeader{
		"Content-Type":              {`application/octet-stream; name="encrypted.asc"`},
		"Content-Disposition":       {`inline; filename="encrypted.asc"`},
		"Content-Transfer-Encoding": {"7bit"},
	})
	if err != nil {
		return Entity{}, err
	}
	if _, err := dataPart.Write(armored.Bytes()); err != nil {
		return Entity{}, err
	}
	if err := mw.Close(); err != nil {
		return Entity{}, err
	}

	return Entity{
		Headers: []string{
			fmt.Sprintf(`Content-Type: multipart/encrypted; protocol="application/pgp-encrypted"; boundary="%s"`, mw.Boundary()),
		},
		Body: body.Bytes(),
	}, nil
}

// DecryptEntity reverses EncryptEntity, decrypting with unlockedIdentity (its
// private key material must already be unlocked via UnlockPrivateKey — this
// function never takes a passphrase itself).
func DecryptEntity(entity Entity, unlockedIdentity *openpgp.Entity) (Entity, error) {
	ct := headerValue(entity.Headers, "Content-Type")
	mediaType, params, err := mime.ParseMediaType(ct)
	if err != nil || mediaType != "multipart/encrypted" {
		return Entity{}, errors.New("pgp: not a multipart/encrypted message")
	}
	if !strings.EqualFold(params["protocol"], "application/pgp-encrypted") {
		return Entity{}, fmt.Errorf("pgp: unsupported multipart/encrypted protocol %q", params["protocol"])
	}
	boundary := params["boundary"]
	if boundary == "" {
		return Entity{}, errors.New("pgp: missing multipart boundary")
	}

	mr := multipart.NewReader(bytes.NewReader(entity.Body), boundary)
	// First part is the application/pgp-encrypted control part ("Version: 1") — not
	// needed, the actual ciphertext is the second part.
	if _, err := mr.NextPart(); err != nil {
		return Entity{}, fmt.Errorf("pgp: read control part: %w", err)
	}
	dataPart, err := mr.NextPart()
	if err != nil {
		return Entity{}, fmt.Errorf("pgp: read data part: %w", err)
	}
	armoredCiphertext, err := io.ReadAll(dataPart)
	if err != nil {
		return Entity{}, err
	}

	block, err := armor.Decode(bytes.NewReader(armoredCiphertext))
	if err != nil {
		return Entity{}, fmt.Errorf("pgp: invalid armored ciphertext: %w", err)
	}
	md, err := openpgp.ReadMessage(block.Body, openpgp.EntityList{unlockedIdentity}, nil, defaultConfig())
	if err != nil {
		return Entity{}, fmt.Errorf("pgp: decrypt: %w", err)
	}
	plaintext, err := io.ReadAll(md.UnverifiedBody)
	if err != nil {
		return Entity{}, fmt.Errorf("pgp: decrypt: %w", err)
	}
	return parseEntity(plaintext), nil
}