package imap

import (
	"encoding/base64"
	"strings"
	"testing"
)

func TestExtractAttachments_plainAttachment(t *testing.T) {
	pdfData := []byte("%PDF-1.4\n")
	raw := buildMultipartMessage(t, "mixed", []mimePart{
		{
			contentType: "text/plain",
			body:        []byte("Hello world"),
		},
		{
			contentType: "application/pdf; name=\"doc.pdf\"",
			disposition: "attachment; filename=\"doc.pdf\"",
			body:        pdfData,
			transferEnc: "base64",
		},
	})

	attachments, err := ExtractAttachments(raw)
	if err != nil {
		t.Fatalf("ExtractAttachments() error = %v", err)
	}
	if len(attachments) != 1 {
		t.Fatalf("len(attachments) = %d, want 1", len(attachments))
	}

	att := attachments[0]
	if att.Filename != "doc.pdf" {
		t.Fatalf("Filename = %q, want doc.pdf", att.Filename)
	}
	if att.ContentType != "application/pdf" {
		t.Fatalf("ContentType = %q, want application/pdf", att.ContentType)
	}
	if att.IsInline {
		t.Fatal("IsInline = true, want false")
	}
	if att.ContentID != "" {
		t.Fatalf("ContentID = %q, want empty", att.ContentID)
	}
	if string(att.Data) != string(pdfData) {
		t.Fatalf("Data = %q, want %q", att.Data, pdfData)
	}
}

func TestExtractAttachments_inlineWithCID(t *testing.T) {
	pngData := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}
	raw := buildMultipartMessage(t, "related", []mimePart{
		{
			contentType: "text/html",
			body:        []byte(`<html><body><img src="cid:logo@cid"></body></html>`),
		},
		{
			contentType: "image/png; name=\"logo.png\"",
			disposition: "inline; filename=\"logo.png\"",
			contentID:   "<logo@cid>",
			body:        pngData,
			transferEnc: "base64",
		},
	})

	attachments, err := ExtractAttachments(raw)
	if err != nil {
		t.Fatalf("ExtractAttachments() error = %v", err)
	}
	if len(attachments) != 1 {
		t.Fatalf("len(attachments) = %d, want 1", len(attachments))
	}

	att := attachments[0]
	if att.Filename != "logo.png" {
		t.Fatalf("Filename = %q, want logo.png", att.Filename)
	}
	if att.ContentType != "image/png" {
		t.Fatalf("ContentType = %q, want image/png", att.ContentType)
	}
	if !att.IsInline {
		t.Fatal("IsInline = false, want true")
	}
	if att.ContentID != "logo@cid" {
		t.Fatalf("ContentID = %q, want logo@cid", att.ContentID)
	}
	if string(att.Data) != string(pngData) {
		t.Fatalf("Data = %q, want %q", att.Data, pngData)
	}
}

func TestExtractAttachments_inlineCIDWithoutFilename(t *testing.T) {
	pngData := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}
	raw := buildMultipartMessage(t, "related", []mimePart{
		{
			contentType: "text/html",
			body:        []byte(`<html><body><img src="cid:part1"></body></html>`),
		},
		{
			contentType: "image/png",
			disposition: "inline",
			contentID:   "<part1>",
			body:        pngData,
			transferEnc: "base64",
		},
	})

	attachments, err := ExtractAttachments(raw)
	if err != nil {
		t.Fatalf("ExtractAttachments() error = %v", err)
	}
	if len(attachments) != 1 {
		t.Fatalf("len(attachments) = %d, want 1", len(attachments))
	}

	att := attachments[0]
	if att.ContentID != "part1" {
		t.Fatalf("ContentID = %q, want part1", att.ContentID)
	}
	if !att.IsInline {
		t.Fatal("IsInline = false, want true")
	}
	if att.Filename == "" {
		t.Fatal("Filename empty, want generated name")
	}
	if string(att.Data) != string(pngData) {
		t.Fatalf("Data mismatch")
	}
}

func TestExtractAttachments_nestedMixedRelatedInline(t *testing.T) {
	pngData := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}
	pdfData := []byte("%PDF-1.4\n")
	raw := []byte(
		"From: ikea@example.com\r\n" +
			"To: user@example.com\r\n" +
			"Subject: ticket\r\n" +
			"MIME-Version: 1.0\r\n" +
			"Content-Type: multipart/mixed; boundary=\"mix\"\r\n" +
			"\r\n" +
			"--mix\r\n" +
			"Content-Type: multipart/related; boundary=\"rel\"\r\n" +
			"\r\n" +
			"--rel\r\n" +
			"Content-Type: multipart/alternative; boundary=\"alt\"\r\n" +
			"\r\n" +
			"--alt\r\n" +
			"Content-Type: text/plain\r\n" +
			"\r\n" +
			"Merci pour votre achat\r\n" +
			"--alt\r\n" +
			"Content-Type: text/html\r\n" +
			"\r\n" +
			"<img src=\"cid:Color.png\" alt=\"Logo\">\r\n" +
			"--alt--\r\n" +
			"--rel\r\n" +
			"Content-Type: image/png; name=\"Color.png\"\r\n" +
			"Content-Transfer-Encoding: base64\r\n" +
			"Content-ID: <Color.png>\r\n" +
			"Content-Disposition: inline; filename=\"Color.png\"\r\n" +
			"\r\n" +
			base64.StdEncoding.EncodeToString(pngData) + "\r\n" +
			"--rel--\r\n" +
			"--mix\r\n" +
			"Content-Type: application/pdf; name=\"ticket.pdf\"\r\n" +
			"Content-Transfer-Encoding: base64\r\n" +
			"Content-Disposition: attachment; filename=\"ticket.pdf\"\r\n" +
			"\r\n" +
			base64.StdEncoding.EncodeToString(pdfData) + "\r\n" +
			"--mix--\r\n",
	)

	attachments, err := ExtractAttachments(raw)
	if err != nil {
		t.Fatalf("ExtractAttachments() error = %v", err)
	}
	if len(attachments) != 2 {
		t.Fatalf("len(attachments) = %d, want 2 (inline png + pdf)", len(attachments))
	}

	var inlineFound, pdfFound bool
	for _, att := range attachments {
		switch {
		case att.IsInline && att.ContentID == "Color.png":
			inlineFound = true
		case !att.IsInline && att.Filename == "ticket.pdf":
			pdfFound = true
		}
	}
	if !inlineFound {
		t.Fatal("inline Color.png attachment missing")
	}
	if !pdfFound {
		t.Fatal("pdf attachment missing")
	}
}

func TestExtractAttachments_skipsBodyParts(t *testing.T) {
	raw := buildMultipartMessage(t, "alternative", []mimePart{
		{
			contentType: "text/plain",
			body:        []byte("plain body"),
		},
		{
			contentType: "text/html",
			body:        []byte("<p>html body</p>"),
		},
	})

	attachments, err := ExtractAttachments(raw)
	if err != nil {
		t.Fatalf("ExtractAttachments() error = %v", err)
	}
	if len(attachments) != 0 {
		t.Fatalf("len(attachments) = %d, want 0", len(attachments))
	}
}

type mimePart struct {
	contentType string
	disposition string
	contentID   string
	transferEnc string
	body        []byte
}

func buildMultipartMessage(t *testing.T, subtype string, parts []mimePart) []byte {
	t.Helper()

	const boundary = "test-boundary"
	var b strings.Builder
	b.WriteString("From: sender@example.com\r\n")
	b.WriteString("To: recipient@example.com\r\n")
	b.WriteString("Subject: attachment test\r\n")
	b.WriteString("MIME-Version: 1.0\r\n")
	b.WriteString("Content-Type: multipart/" + subtype + "; boundary=\"" + boundary + "\"\r\n")
	b.WriteString("\r\n")

	for _, part := range parts {
		b.WriteString("--" + boundary + "\r\n")
		b.WriteString("Content-Type: " + part.contentType + "\r\n")
		if part.disposition != "" {
			b.WriteString("Content-Disposition: " + part.disposition + "\r\n")
		}
		if part.contentID != "" {
			b.WriteString("Content-ID: " + part.contentID + "\r\n")
		}
		if part.transferEnc != "" {
			b.WriteString("Content-Transfer-Encoding: " + part.transferEnc + "\r\n")
		}
		b.WriteString("\r\n")

		if part.transferEnc == "base64" {
			b.WriteString(base64.StdEncoding.EncodeToString(part.body))
		} else {
			b.Write(part.body)
		}
		b.WriteString("\r\n")
	}
	b.WriteString("--" + boundary + "--\r\n")

	return []byte(b.String())
}
