go_gameshield/pkg/excel/excel.go

package excel

import (
	"fmt"
	"github.com/xuri/excelize/v2"
	"io"
	"math"
	"net/http"
	"strconv"
	"time"
)

// ExcelGenerator 通用Excel生成器
type ExcelGenerator struct {
	file       *excelize.File
	sheetName  string
	headers    []string
	headerMap  map[string]string // 表头映射（英文字段名->中文显示名）
	currentRow int               // 当前行号
}

// NewExcelGenerator 创建新的Excel生成器
func NewExcelGenerator(sheetName string, headers []string, headerMap map[string]string) *ExcelGenerator {
	f := excelize.NewFile()
	// 默认sheet名称是Sheet1，如果传入的是其他名称，则创建新sheet并删除默认sheet
	defaultSheetName := "Sheet1"
	if sheetName != defaultSheetName {
		f.NewSheet(sheetName)
		f.DeleteSheet(defaultSheetName)
	}
	
	return &ExcelGenerator{
		file:       f,
		sheetName:  sheetName,
		headers:    headers,
		headerMap:  headerMap,
		currentRow: 1, // 从第1行开始（通常第1行是表头）
	}
}

// WriteHeaders 写入表头
func (g *ExcelGenerator) WriteHeaders() error {
	// 写入表头
	for i, header := range g.headers {
		// 获取对应的显示名，如果没有映射则使用原字段名
		displayName, exists := g.headerMap[header]
		if !exists {
			displayName = header
		}
		
		cell := fmt.Sprintf("%s%d", columnName(i), g.currentRow)
		if err := g.file.SetCellValue(g.sheetName, cell, displayName); err != nil {
			return err
		}
		
		// 设置表头样式（加粗、居中等）
		style, err := g.file.NewStyle(&excelize.Style{
			Font: &excelize.Font{
				Bold: true,
			},
			Alignment: &excelize.Alignment{
				Horizontal: "center",
				Vertical:   "center",
			},
		})
		if err != nil {
			return err
		}
		if err := g.file.SetCellStyle(g.sheetName, cell, cell, style); err != nil {
			return err
		}
	}
	
	g.currentRow++ // 表头写完后，行号+1
	return nil
}

// WriteRows 写入多行数据
func (g *ExcelGenerator) WriteRows(data []map[string]interface{}) error {
	for _, row := range data {
		if err := g.WriteRow(row); err != nil {
			return err
		}
	}
	return nil
}

// WriteRow 写入单行数据
func (g *ExcelGenerator) WriteRow(rowData map[string]interface{}) error {
	for i, field := range g.headers {
		value, exists := rowData[field]
		if !exists {
			value = "" // 如果数据中不存在该字段，则写入空值
		}
		
		cell := fmt.Sprintf("%s%d", columnName(i), g.currentRow)
		
		// 根据不同的数据类型处理
		switch v := value.(type) {
		case time.Time:
			// 时间格式化为 YYYY-MM-DD HH:MM:SS
			if err := g.file.SetCellValue(g.sheetName, cell, v.Format("2006-01-02 15:04:05")); err != nil {
				return err
			}
		default:
			if err := g.file.SetCellValue(g.sheetName, cell, v); err != nil {
				return err
			}
		}
	}
	
	g.currentRow++ // 一行写完后，行号+1
	return nil
}

// SaveToWriter 保存到io.Writer接口
func (g *ExcelGenerator) SaveToWriter(w io.Writer) error {
	return g.file.Write(w)
}

// SaveToBuffer 保存到内存
func (g *ExcelGenerator) SaveToBuffer() ([]byte, error) {
	buffer, err := g.file.WriteToBuffer()
	if err != nil {
		return nil, err
	}
	return buffer.Bytes(), nil
}

// columnName 将列索引转换为Excel列名（A, B, C, ... Z, AA, AB, ...）
func columnName(colIndex int) string {
	if colIndex < 26 {
		return string('A' + colIndex)
	}
	
	// 超过26列时需要用两个或更多字母表示
	result := ""
	for colIndex >= 0 {
		remainder := colIndex % 26
		result = string('A'+remainder) + result
		colIndex = colIndex/26 - 1
		if colIndex < 0 {
			break
		}
	}
	
	return result
}

// TransferOption 传输选项
type TransferOption struct {
	FileName    string
	ContentType string
}

// ExportType 导出类型枚举
type ExportType int

const (
	// ExportTypeNormal 普通导出（小文件，加载到内存后直接传输）
	ExportTypeNormal ExportType = iota
	
	// ExportTypeStream 流式导出（大文件，流式传输避免占用过多内存）
	ExportTypeStream
	
	// ExportTypeChunk 分块导出（超大文件，分块处理并传输）
	ExportTypeChunk
)

// SmartExport 智能选择导出方式
// dataCount: 数据条数
// rowSize: 每行数据的估计大小（字节数）
func SmartExport(dataCount int, rowSize int) ExportType {
	// 估算导出文件大小（表头+数据）
	estimatedSize := (dataCount + 1) * rowSize
	
	// 根据估算大小选择不同的导出方式
	switch {
	case estimatedSize <= 5*1024*1024: // 5MB以下用普通导出
		return ExportTypeNormal
		
	case estimatedSize <= 50*1024*1024: // 5MB-50MB用流式导出
		return ExportTypeStream
		
	default: // 超过50MB用分块导出
		return ExportTypeChunk
	}
}

// NormalExport 普通导出（小文件，一次性加载到内存）
func NormalExport(g *ExcelGenerator, w http.ResponseWriter, option TransferOption) error {
	// 设置响应头
	w.Header().Set("Content-Type", option.ContentType)
	w.Header().Set("Content-Disposition", "attachment; filename="+option.FileName)
	
	// 直接写入响应
	return g.SaveToWriter(w)
}

// StreamExport 流式导出（大文件，流式传输）
func StreamExport(g *ExcelGenerator, w http.ResponseWriter, option TransferOption) error {
	// 设置响应头
	w.Header().Set("Content-Type", option.ContentType)
	w.Header().Set("Content-Disposition", "attachment; filename="+option.FileName)
	w.Header().Set("Transfer-Encoding", "chunked")
	
	// 流式写入
	return g.SaveToWriter(w)
}

// ChunkExport 分块导出（超大文件）
// 这个方法需要配合前端实现，例如通过分页API多次获取数据并合并
func ChunkExport(w http.ResponseWriter, option TransferOption, totalRecords int, pageSize int) {
	totalPages := int(math.Ceil(float64(totalRecords) / float64(pageSize)))
	
	// 设置响应头
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("X-Total-Pages", strconv.Itoa(totalPages))
	w.Header().Set("X-Total-Records", strconv.Itoa(totalRecords))
	w.Header().Set("X-Page-Size", strconv.Itoa(pageSize))
	
	// 返回分块导出信息
	w.Write([]byte(fmt.Sprintf(`{
		"message": "File is too large for direct download. Please use paginated export.",
		"total_records": %d,
		"total_pages": %d,
		"page_size": %d
	}`, totalRecords, totalPages, pageSize)))
}