package securechat

import (
	"bytes"
	"crypto/rand"
	"crypto/subtle"
	"encoding/binary"
	"encoding/json"
	"regexp"
	"strconv"
	"strings"

	"golang.org/x/crypto/blake2b"
	"golang.org/x/crypto/chacha20poly1305"
)

// Adjuntos (ADR-033): trozos de 64 KiB con XChaCha20-Poly1305.
//
//	nonce_i = nonceBase[0..16) ‖ u64be(i)
//	AD_i    = SecureChat|v1|att|<attachmentId>|<i>|<1 si es el último, 0 si no>
const (
	AttachmentChunk    = 64 * 1024
	AttachmentMaxBytes = 25 * 1024 * 1024
)

var (
	attachmentIDRe = regexp.MustCompile(`^att_[A-Za-z0-9_-]{8,64}$`)
	mimeRe         = regexp.MustCompile(`^[a-z0-9][a-z0-9!#$&^_.+-]{0,63}/[a-z0-9][a-z0-9!#$&^_.+-]{0,127}$`)
)

// AttachmentMeta viaja DENTRO de un mensaje cifrado (contentType "attachment").
// El orden de los campos es el de envelope.js.
type AttachmentMeta struct {
	V            int     `json:"v"`
	Kind         string  `json:"kind"`
	AttachmentID string  `json:"attachmentId"`
	Name         string  `json:"name"`
	Mime         string  `json:"mime"`
	Size         int     `json:"size"`
	Key          string  `json:"key"`
	Nonce        string  `json:"nonce"`
	Digest       string  `json:"digest"`
	ChunkSize    int     `json:"chunkSize"`
	Caption      *string `json:"caption,omitempty"`
	Width        *int    `json:"width,omitempty"`
	Height       *int    `json:"height,omitempty"`
}

func AttachmentCiphertextLength(plainSize int) int {
	chunks := (plainSize + AttachmentChunk - 1) / AttachmentChunk
	if chunks < 1 {
		chunks = 1
	}
	return plainSize + chunks*16
}

// AttachmentDigest = base64url(BLAKE2b-256 sin clave).
func AttachmentDigest(b []byte) string {
	h := blake2b.Sum256(b)
	return B64uEncode(h[:])
}

func chunkNonce(base []byte, i int) []byte {
	n := make([]byte, 24)
	copy(n, base[:16])
	binary.BigEndian.PutUint64(n[16:], uint64(i))
	return n
}

func chunkAD(attachmentID string, i int, last bool) []byte {
	l := "0"
	if last {
		l = "1"
	}
	return []byte(domain + "|att|" + attachmentID + "|" + strconv.Itoa(i) + "|" + l)
}

// EncryptAttachment cifra un fichero. key/nonceBase == nil → aleatorios (solo los vectores los fijan).
func EncryptAttachment(data []byte, attachmentID, name, mime string, caption *string, key, nonceBase []byte) ([]byte, AttachmentMeta, error) {
	if !attachmentIDRe.MatchString(attachmentID) {
		return nil, AttachmentMeta{}, fail(ErrInvalidInput, "attachmentId inválido")
	}
	if len(data) > AttachmentMaxBytes {
		return nil, AttachmentMeta{}, fail(ErrInvalidInput, "adjunto demasiado grande")
	}
	if key == nil {
		key = make([]byte, 32)
		_, _ = rand.Read(key)
	}
	if nonceBase == nil {
		nonceBase = make([]byte, 24)
		_, _ = rand.Read(nonceBase)
	}
	if err := assertLen(key, 32, "key"); err != nil {
		return nil, AttachmentMeta{}, err
	}
	if err := assertLen(nonceBase, 24, "nonceBase"); err != nil {
		return nil, AttachmentMeta{}, err
	}
	aead, _ := chacha20poly1305.NewX(key)
	chunks := (len(data) + AttachmentChunk - 1) / AttachmentChunk
	if chunks < 1 {
		chunks = 1
	}
	out := make([]byte, 0, AttachmentCiphertextLength(len(data)))
	for i := 0; i < chunks; i++ {
		end := (i + 1) * AttachmentChunk
		if end > len(data) {
			end = len(data)
		}
		out = aead.Seal(out, chunkNonce(nonceBase, i), data[i*AttachmentChunk:end], chunkAD(attachmentID, i, i == chunks-1))
	}
	meta := AttachmentMeta{
		V: 1, Kind: "attachment", AttachmentID: attachmentID, Name: name, Mime: mime, Size: len(data),
		Key: B64uEncode(key), Nonce: B64uEncode(nonceBase), Digest: AttachmentDigest(data), ChunkSize: AttachmentChunk,
		Caption: caption,
	}
	m, err := ValidateAttachmentMeta(meta)
	return out, m, err
}

// DecryptAttachment descifra y comprueba tamaño y huella. DECRYPT_FAILED ante cualquier manipulación.
func DecryptAttachment(ciphertext []byte, meta AttachmentMeta) ([]byte, error) {
	m, err := ValidateAttachmentMeta(meta)
	if err != nil {
		return nil, err
	}
	key, _ := B64uDecode(m.Key)
	nb, _ := B64uDecode(m.Nonce)
	if len(ciphertext) != AttachmentCiphertextLength(m.Size) {
		return nil, fail(ErrDecryptFailed, "longitud del adjunto inesperada")
	}
	aead, _ := chacha20poly1305.NewX(key)
	chunks := (m.Size + AttachmentChunk - 1) / AttachmentChunk
	if chunks < 1 {
		chunks = 1
	}
	out := make([]byte, 0, m.Size)
	c := 0
	for i := 0; i < chunks; i++ {
		plainLen := m.Size - i*AttachmentChunk
		if plainLen > AttachmentChunk {
			plainLen = AttachmentChunk
		}
		part := ciphertext[c : c+plainLen+16]
		c += plainLen + 16
		var err error
		out, err = aead.Open(out, chunkNonce(nb, i), part, chunkAD(m.AttachmentID, i, i == chunks-1))
		if err != nil {
			return nil, fail(ErrDecryptFailed, "trozo %d del adjunto no autentica", i)
		}
	}
	if subtle.ConstantTimeCompare([]byte(AttachmentDigest(out)), []byte(m.Digest)) != 1 {
		return nil, fail(ErrDecryptFailed, "la huella del adjunto no coincide")
	}
	return out, nil
}

// ValidateAttachmentMeta aplica las reglas de validateAttachmentMeta de envelope.js.
func ValidateAttachmentMeta(m AttachmentMeta) (AttachmentMeta, error) {
	if m.V != 1 || m.Kind != "attachment" {
		return m, fail(ErrUnsupportedVersion, "meta de adjunto de versión desconocida")
	}
	if !attachmentIDRe.MatchString(m.AttachmentID) {
		return m, fail(ErrInvalidInput, "attachmentId inválido")
	}
	if !validUTF8([]byte(m.Name)) || isJSBlank(m.Name) || utf16Len(m.Name) > 255 || strings.ContainsAny(m.Name, "/\\") || hasControl(m.Name) {
		return m, fail(ErrInvalidInput, "nombre de adjunto inválido")
	}
	if !mimeRe.MatchString(m.Mime) {
		return m, fail(ErrInvalidInput, "tipo de adjunto inválido")
	}
	if m.Size < 0 || m.Size > AttachmentMaxBytes {
		return m, fail(ErrInvalidInput, "tamaño de adjunto inválido")
	}
	if m.ChunkSize != AttachmentChunk {
		return m, fail(ErrUnsupportedVersion, "tamaño de trozo no soportado")
	}
	if _, err := decodeLen(m.Key, 32, "key"); err != nil {
		return m, err
	}
	if _, err := decodeLen(m.Nonce, 24, "nonce"); err != nil {
		return m, err
	}
	if _, err := decodeLen(m.Digest, 32, "digest"); err != nil {
		return m, err
	}
	if m.Caption != nil && (!validUTF8([]byte(*m.Caption)) || utf16Len(*m.Caption) > 4000) {
		return m, fail(ErrInvalidInput, "pie de adjunto inválido")
	}
	if m.Width != nil || m.Height != nil {
		if m.Width == nil || m.Height == nil || *m.Width < 1 || *m.Height < 1 || *m.Width > 20000 || *m.Height > 20000 {
			return m, fail(ErrInvalidInput, "dimensiones de adjunto inválidas")
		}
	}
	return m, nil
}

func hasControl(s string) bool {
	for i := 0; i < len(s); i++ {
		if s[i] < 0x20 {
			return true
		}
	}
	return false
}

// AttachmentMessageText: el JSON de la meta, que se cifra con EncryptReply.
func AttachmentMessageText(m AttachmentMeta) (string, error) {
	v, err := ValidateAttachmentMeta(m)
	if err != nil {
		return "", err
	}
	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	enc.SetEscapeHTML(false)
	if err := enc.Encode(v); err != nil {
		return "", err
	}
	return strings.TrimSuffix(buf.String(), "\n"), nil
}

// ParseAttachmentMessage: inversa de AttachmentMessageText. INVALID_INPUT si no es una meta válida.
func ParseAttachmentMessage(text string) (AttachmentMeta, error) {
	var m AttachmentMeta
	if err := json.Unmarshal([]byte(text), &m); err != nil {
		return m, fail(ErrInvalidInput, "el mensaje de adjunto no es JSON válido")
	}
	return ValidateAttachmentMeta(m)
}