From 0aaafdbaa3ff78d171f2b4b59e627ed6c2246986 Mon Sep 17 00:00:00 2001
From: yoge <yogecoder@gmail.com>
Date: Fri, 26 Dec 2025 15:48:14 +0800
Subject: [PATCH] feat: add file export

---
 api/v1/task/handler.go         | 28 +++++++++++
 internal/model/task/request.go |  5 ++
 internal/service/task.go       | 86 ++++++++++++++++++++++++++++++++++
 3 files changed, 119 insertions(+)

diff --git a/api/v1/task/handler.go b/api/v1/task/handler.go
index 74ee88d..62aaaa4 100644
--- a/api/v1/task/handler.go
+++ b/api/v1/task/handler.go
@@ -61,3 +61,31 @@ func (h *TaskEndpoint) GetTaskList(c *gin.Context) {
 
 	c.JSON(http.StatusOK, common.SuccessResponse(c, resp))
 }
+
+func (h *TaskEndpoint) ExportTask(c *gin.Context) {
+	var req task.ExportTaskRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		log.Error(c, "func", "ExportTask", "msg", "Invalid parameters", "error", err)
+		c.JSON(http.StatusOK, common.ErrorResponse(c, common.CodeParamError, "Invalid parameters"))
+		return
+	}
+	fileData, contentType, err := h.taskService.ExportTask(c, &req)
+	if err != nil {
+		c.JSON(http.StatusOK, common.ErrorResponse(c, common.CodeSystemError, "导出任务失败"))
+		return
+	}
+
+	// set filename based on export type
+	var filename string
+	switch req.Type {
+	case "pdf":
+		filename = "texpixel_export.pdf"
+	case "docx":
+		filename = "texpixel_export.docx"
+	default:
+		filename = "texpixel_export"
+	}
+
+	c.Header("Content-Disposition", "attachment; filename="+filename)
+	c.Data(http.StatusOK, contentType, fileData)
+}
diff --git a/internal/model/task/request.go b/internal/model/task/request.go
index 883384f..9842d9c 100644
--- a/internal/model/task/request.go
+++ b/internal/model/task/request.go
@@ -34,3 +34,8 @@ type TaskListResponse struct {
 	TaskList []*TaskListDTO `json:"task_list"`
 	Total    int64          `json:"total"`
 }
+
+type ExportTaskRequest struct {
+	TaskID int64  `json:"task_id" binding:"required"`
+	Type   string `json:"type" binding:"required,oneof=pdf docx"`
+}
diff --git a/internal/service/task.go b/internal/service/task.go
index bc5f979..64a52bf 100644
--- a/internal/service/task.go
+++ b/internal/service/task.go
@@ -1,8 +1,13 @@
 package service
 
 import (
+	"bytes"
 	"context"
 	"errors"
+	"fmt"
+	"io"
+	"mime/multipart"
+	"net/http"
 	"strings"
 
 	"gitea.com/bitwsd/document_ai/internal/model/task"
@@ -111,3 +116,84 @@ func (svc *TaskService) GetTaskList(ctx context.Context, req *task.TaskListReque
 	}
 	return resp, nil
 }
+
+func (svc *TaskService) ExportTask(ctx context.Context, req *task.ExportTaskRequest) ([]byte, string, error) {
+	recognitionTask, err := svc.recognitionTaskDao.GetTaskByID(dao.DB.WithContext(ctx), req.TaskID)
+	if err != nil {
+		log.Error(ctx, "func", "ExportTask", "msg", "get task by task id failed", "error", err)
+		return nil, "", err
+	}
+
+	if recognitionTask == nil {
+		log.Error(ctx, "func", "ExportTask", "msg", "task not found")
+		return nil, "", errors.New("task not found")
+	}
+
+	if recognitionTask.Status != dao.TaskStatusCompleted {
+		log.Error(ctx, "func", "ExportTask", "msg", "task not finished")
+		return nil, "", errors.New("task not finished")
+	}
+
+	recognitionResult, err := svc.recognitionResultDao.GetByTaskID(dao.DB.WithContext(ctx), req.TaskID)
+	if err != nil {
+		log.Error(ctx, "func", "ExportTask", "msg", "get recognition result by task id failed", "error", err)
+		return nil, "", err
+	}
+
+	if recognitionResult == nil {
+		log.Error(ctx, "func", "ExportTask", "msg", "recognition result not found")
+		return nil, "", errors.New("recognition result not found")
+	}
+
+	markdown := recognitionResult.Markdown
+	if markdown == "" {
+		log.Error(ctx, "func", "ExportTask", "msg", "markdown not found")
+		return nil, "", errors.New("markdown not found")
+	}
+
+	// call http://localhost:8055/export
+	body := &bytes.Buffer{}
+	writer := multipart.NewWriter(body)
+	_ = writer.WriteField("markdown", markdown)
+	_ = writer.WriteField("type", req.Type)
+	writer.Close()
+
+	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://localhost:8055/export", body)
+	if err != nil {
+		log.Error(ctx, "func", "ExportTask", "msg", "create http request failed", "error", err)
+		return nil, "", err
+	}
+	httpReq.Header.Set("Content-Type", writer.FormDataContentType())
+
+	client := &http.Client{}
+	resp, err := client.Do(httpReq)
+	if err != nil {
+		log.Error(ctx, "func", "ExportTask", "msg", "http request failed", "error", err)
+		return nil, "", err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		log.Error(ctx, "func", "ExportTask", "msg", "http request failed", "status", resp.StatusCode)
+		return nil, "", fmt.Errorf("export service returned status: %d", resp.StatusCode)
+	}
+
+	fileData, err := io.ReadAll(resp.Body)
+	if err != nil {
+		log.Error(ctx, "func", "ExportTask", "msg", "read response body failed", "error", err)
+		return nil, "", err
+	}
+
+	// determine content type based on export type
+	var contentType string
+	switch req.Type {
+	case "pdf":
+		contentType = "application/pdf"
+	case "docx":
+		contentType = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+	default:
+		contentType = "application/octet-stream"
+	}
+
+	return fileData, contentType, nil
+}