package agent

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"image"
	"image/png"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
)

func uploadFixture(t *testing.T, s *Service, id, name string, data []byte) Attachment {
	t.Helper()
	a, err := s.UploadAttachment(context.Background(), id, AttachmentInput{Name: name, Data: base64.StdEncoding.EncodeToString(data)})
	if err != nil {
		t.Fatal(err)
	}
	return a
}

func officeFixture(t *testing.T, files map[string]string) []byte {
	t.Helper()
	var buf bytes.Buffer
	z := zip.NewWriter(&buf)
	for name, text := range files {
		w, err := z.Create(name)
		if err != nil {
			t.Fatal(err)
		}
		if _, err = w.Write([]byte(text)); err != nil {
			t.Fatal(err)
		}
	}
	if err := z.Close(); err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}

func pdfFixture() []byte {
	stream := "BT /F1 12 Tf 40 70 Td (Budget 3500 quantity 500) Tj ET"
	objects := []string{"<< /Type /Catalog /Pages 2 0 R >>", "<< /Type /Pages /Kids [3 0 R] /Count 1 >>", "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 300 144] /Resources << /Font << /F1 4 0 R >> >> /Contents 5 0 R >>", "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", fmt.Sprintf("<< /Length %d >>\nstream\n%s\nendstream", len(stream), stream)}
	var buf bytes.Buffer
	buf.WriteString("%PDF-1.4\n")
	var offsets []int
	for i, obj := range objects {
		offsets = append(offsets, buf.Len())
		fmt.Fprintf(&buf, "%d 0 obj\n%s\nendobj\n", i+1, obj)
	}
	xref := buf.Len()
	fmt.Fprintf(&buf, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1)
	for _, offset := range offsets {
		fmt.Fprintf(&buf, "%010d 00000 n \n", offset)
	}
	fmt.Fprintf(&buf, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(objects)+1, xref)
	return buf.Bytes()
}

func TestAttachmentFormatsAndEncoding(t *testing.T) {
	docx := officeFixture(t, map[string]string{"word/document.xml": `<w:document xmlns:w="urn:test"><w:body><w:p><w:r><w:t>预算3500元</w:t></w:r></w:p><w:p><w:r><w:t>数量500个</w:t></w:r></w:p></w:body></w:document>`})
	xlsx := officeFixture(t, map[string]string{"xl/sharedStrings.xml": `<sst><si><t>数量</t></si></sst>`, "xl/worksheets/sheet1.xml": `<worksheet><sheetData><row><c r="A1" t="s"><v>0</v></c><c r="B1"><v>500</v></c><c r="C1" t="inlineStr"><is><t>预算3500</t></is></c></row></sheetData></worksheet>`})
	for _, tc := range []struct {
		name string
		data []byte
		want string
	}{{"资料.txt", []byte("预算3500元"), "预算3500"}, {"utf16.txt", []byte{0xff, 0xfe, '5', 0, '0', 0, '0', 0}, "500"}, {"报价.docx", docx, "数量500"}, {"客户.xlsx", xlsx, "B1=500"}, {"报价.pdf", pdfFixture(), "3500"}} {
		t.Run(tc.name, func(t *testing.T) {
			result, err := extractAttachment(context.Background(), tc.name, tc.data)
			if err != nil || !strings.Contains(result.Text, tc.want) {
				t.Fatalf("extract %s: %q %v", tc.name, result.Text, err)
			}
		})
	}
	var picture bytes.Buffer
	png.Encode(&picture, image.NewRGBA(image.Rect(0, 0, 2, 2)))
	result, err := extractAttachment(context.Background(), "图.png", picture.Bytes())
	if err != nil || result.Kind != "image" || result.MIMEType != "image/png" {
		t.Fatalf("image: %+v %v", result, err)
	}
	for _, tc := range []struct {
		name string
		data []byte
	}{{"a.exe", []byte("hello")}, {"a.txt", []byte{0, 1, 2}}, {"a.png", []byte("not image")}, {"a.pdf", []byte("invalid PDF")}, {"a.docx", []byte("not zip")}} {
		if _, err := extractAttachment(context.Background(), tc.name, tc.data); err == nil {
			t.Fatalf("invalid file accepted: %s", tc.name)
		}
	}
	result, err = extractAttachment(context.Background(), "long.txt", []byte(strings.Repeat("内容", 40000)))
	if err != nil || !result.Truncated || len(result.Text) > maxAttachmentText {
		t.Fatal("long text limit", err)
	}
}

func TestAttachmentsPersistAndStayWithinConversation(t *testing.T) {
	s, secrets := testService(t)
	c, _ := s.CreateConversation()
	other, _ := s.CreateConversation()
	a := uploadFixture(t, s, c.ID, "需求.md", []byte("预算3500，数量500，只起草不发送"))
	c.DraftAttachments = []Attachment{a}
	c.Draft = "分析这份资料"
	if _, err := s.SaveConversation(c); err != nil {
		t.Fatal(err)
	}
	restarted := newService(s.configPath, secrets)
	loaded, err := restarted.GetConversation(c.ID)
	if err != nil || len(loaded.DraftAttachments) != 1 || loaded.DraftAttachments[0].Name != a.Name {
		t.Fatalf("restore: %+v %v", loaded, err)
	}
	other.DraftAttachments = []Attachment{a}
	if _, err := s.SaveConversation(other); err == nil {
		t.Fatal("cross-conversation file reference accepted")
	}
	forged := a
	forged.Name = "forged"
	forged.Size = -1
	c.DraftAttachments = []Attachment{forged}
	s.SaveConversation(c)
	loaded, _ = s.GetConversation(c.ID)
	if loaded.DraftAttachments[0].Name != a.Name || loaded.DraftAttachments[0].Size <= 0 {
		t.Fatal("client metadata was trusted")
	}
	if err := s.RemoveAttachment(c.ID, a.ID); err == nil {
		t.Fatal("removed referenced draft")
	}
	c.DraftAttachments = nil
	c.Messages = []ConversationMessage{{ID: "u", Role: "user", Content: "分析附件", Attachments: []Attachment{a}}, {ID: "a", Role: "assistant", Content: "已读取", Status: "complete"}}
	s.SaveConversation(c)
	if err := s.RemoveAttachment(c.ID, a.ID); err == nil {
		t.Fatal("removed sent file")
	}
	message, err := s.messageForModel(c.ID, c.Messages[0])
	if err != nil || !strings.Contains(fmt.Sprint(message.Content), "预算3500") {
		t.Fatalf("missing file context: %+v %v", message, err)
	}
	source, err := s.memorySource(context.Background(), c.ID, c.Messages)
	if err != nil || !strings.Contains(source, "预算3500") {
		t.Fatal("file missing in memory source", err)
	}
	dir, _ := s.attachmentDir(c.ID, "")
	if err := s.DeleteConversation(c.ID); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(dir); !os.IsNotExist(err) {
		t.Fatal("conversation deletion left attachments")
	}
	if _, err := s.UploadAttachment(context.Background(), "../escape", AttachmentInput{Name: "safe.txt", Data: base64.StdEncoding.EncodeToString([]byte("text"))}); err == nil {
		t.Fatal("invalid conversation accepted")
	}
}

func TestChatStreamReceivesAttachmentContentAndImageParts(t *testing.T) {
	s, _ := testService(t)
	c, _ := s.CreateConversation()
	a := uploadFixture(t, s, c.ID, "budget.csv", []byte("budget,quantity\n3500,500"))
	var received []modelMessage
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body struct {
			Messages []modelMessage `json:"messages"`
		}
		json.NewDecoder(r.Body).Decode(&body)
		received = body.Messages
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"预算3500，数量500\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n")
	}))
	defer server.Close()
	s.SaveModelConfig(ModelConfigInput{Config: testConfig(server.URL), APIKey: "key"})
	result, err := s.ChatStreamWithFiles(context.Background(), "files-test", c.ID, "分析预算", []string{a.ID})
	if err != nil || !strings.Contains(result["content"], "3500") || !strings.Contains(fmt.Sprint(received[len(received)-1].Content), "3500,500") {
		t.Fatalf("model file context: %v %v", received, err)
	}
	var picture bytes.Buffer
	png.Encode(&picture, image.NewRGBA(image.Rect(0, 0, 2, 2)))
	img := uploadFixture(t, s, c.ID, "图.png", picture.Bytes())
	message, err := s.messageForModel(c.ID, ConversationMessage{Role: "user", Content: "看图", Attachments: []Attachment{img}})
	if err != nil {
		t.Fatal(err)
	}
	raw, _ := json.Marshal(message)
	if !strings.Contains(string(raw), "data:image/png;base64,") {
		t.Fatal("image not adapted for model")
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := s.UploadAttachment(ctx, c.ID, AttachmentInput{Name: "cancel.txt", Data: base64.StdEncoding.EncodeToString([]byte("cancel"))}); err == nil {
		t.Fatal("cancelled upload continued")
	}
}