pj0235-eai_agentplatform/eai_agentplatform/backend-go/internal/api/media.go

package api

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"

	"eai_agentplatform/backend/internal/middleware"
	"eai_agentplatform/backend/internal/model"
	"eai_agentplatform/backend/internal/repository"
	"eai_agentplatform/backend/internal/web"
)

// mediaRepo 素材仓库（便于测试时覆写），包内共享。
var mediaRepo repository.MediaFileRepo

func init() {
	mediaRepo = repository.MediaFileRepo{}
}

var blockedUploadExt = map[string]bool{
	"exe": true, "bin": true, "dll": true, "so": true, "dylib": true,
	"msi": true, "apk": true, "ipa": true, "deb": true, "rpm": true, "pkg": true, "appimage": true,
	"bat": true, "cmd": true, "com": true, "scr": true, "sys": true, "drv": true,
	"ps1": true, "psm1": true, "vbs": true, "vbe": true, "js": true, "jse": true, "wsf": true, "wsh": true,
	"reg": true, "lnk": true, "iso": true, "img": true, "dmg": true,
}

var videoExt = map[string]bool{
	"mp4": true, "mov": true, "avi": true, "mkv": true, "webm": true, "m4v": true, "wmv": true, "flv": true,
}

var validBindType = map[string]bool{"company": true, "product": true, "course": true, "none": true}

const defaultChunkSize = 5 * 1024 * 1024 // 5MB

// randomID 生成随机十六进制标识（文件名/上传会话）
func randomID() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return strconv.FormatInt(time.Now().UnixNano(), 16)
	}
	return hex.EncodeToString(b)
}

func extOf(filename string) string {
	return strings.ToLower(strings.TrimPrefix(filepath.Ext(filename), "."))
}

func isUploadableExt(ext string) bool {
	normalized := strings.ToLower(strings.TrimSpace(ext))
	if normalized == "" {
		return false
	}
	return !blockedUploadExt[normalized]
}

func isVideoExt(ext string) bool {
	return videoExt[strings.ToLower(strings.TrimSpace(ext))]
}

func sizeLimitFor(ext string) int64 {
	if isVideoExt(ext) {
		return Cfg.FileMaxVideo
	}
	return Cfg.FileMaxDoc
}

// 审批前置：未审批文件存 pending/、驳回文件存 rejected/ 子目录（均不对外公开），
// 审批通过后移动到 approved/ 子目录（静态公开，仅 approved 可访问）。
func approvedDir() string { return filepath.Join(Cfg.KBDataDir, "approved") }
func pendingDir() string  { return filepath.Join(Cfg.KBDataDir, "pending") }
func rejectedDir() string { return filepath.Join(Cfg.KBDataDir, "rejected") }

// mediaPathFor 按状态返回素材物理路径
func mediaPathFor(m model.MediaFile) string {
	switch m.Status {
	case "approved":
		return filepath.Join(approvedDir(), m.StoredName)
	case "rejected":
		return filepath.Join(rejectedDir(), m.StoredName)
	default:
		return filepath.Join(pendingDir(), m.StoredName)
	}
}

func parseBind(c *gin.Context) (string, *uint) {
	bt := c.PostForm("bind_type")
	if bt == "" {
		bt = "none"
	}
	if !validBindType[bt] {
		bt = "none"
	}
	var bid *uint
	if s := c.PostForm("bind_id"); s != "" {
		if n, err := strconv.ParseUint(s, 10, 64); err == nil && n > 0 {
			v := uint(n)
			bid = &v
		}
	}
	return bt, bid
}

func parseKnowledgeSpaceKey(c *gin.Context) string {
	key := sanitizeSpaceKey(c.PostForm("knowledge_space_key"))
	if key == "" {
		return "general"
	}
	return ensureKnowledgeSpaceKeyOrDefault(key)
}

func sourceOf(c *gin.Context) string {
	u := middleware.CurrentUser(c)
	if u != nil && u.Role == "admin" {
		return "admin"
	}
	return "employee"
}

// ============ 直传 ============

// Upload POST /api/media/upload —— 文档/小视频直传
func Upload(c *gin.Context) {
	u := middleware.CurrentUser(c)
	file, header, err := c.Request.FormFile("file")
	if err != nil {
		web.Fail(c, web.NewBadRequest("缺少文件字段 file"))
		return
	}
	defer file.Close()

	ext := extOf(header.Filename)
	if !isUploadableExt(ext) {
		web.Fail(c, web.NewBadRequest("不支持的文件类型"))
		return
	}
	if isVideoExt(ext) && header.Size > Cfg.ChunkThreshold {
		web.Fail(c, web.NewBadRequest("视频超过 100MB 请使用分片上传"))
		return
	}
	if header.Size > sizeLimitFor(ext) {
		web.Fail(c, web.NewBadRequest("文件超过大小限制"))
		return
	}

	bindType, bindID := parseBind(c)
	knowledgeSpaceKey := parseKnowledgeSpaceKey(c)
	source := sourceOf(c)
	status := "pending"
	targetDir := pendingDir()
	if source == "admin" {
		status = "approved" // 管理员上传自动通过
		targetDir = approvedDir()
	}

	storedName := randomID() + "." + ext
	dst := filepath.Join(targetDir, storedName)
	if err := os.MkdirAll(targetDir, 0o755); err != nil {
		web.Fail(c, web.NewBadRequest("创建存储目录失败"))
		return
	}
	out, err := os.Create(dst)
	if err != nil {
		web.Fail(c, web.NewBadRequest("保存文件失败"))
		return
	}
	if _, err := io.Copy(out, file); err != nil {
		out.Close()
		web.Fail(c, web.NewBadRequest("写入文件失败"))
		return
	}
	out.Close()

	m := model.MediaFile{
		Filename:          header.Filename,
		StoredName:        storedName,
		StoredPath:        storedName,
		FileExt:           ext,
		FileSize:          header.Size,
		Status:            status,
		Source:            source,
		SubmitterID:       u.ID,
		BindType:          bindType,
		BindID:            bindID,
		KnowledgeSpaceKey: knowledgeSpaceKey,
	}
	if !mediaRepo.Insert(&m) {
		web.Fail(c, web.NewBadRequest("创建素材记录失败"))
		return
	}
	if status == "approved" {
		go runExtractPipeline(m.ID)
	}
	web.OK(c, gin.H{"media_id": m.ID, "status": m.Status, "filename": m.Filename})
}

// ============ 分片上传 ============

type uploadSession struct {
	Filename          string
	FileSize          int64
	Ext               string
	BindType          string
	BindID            *uint
	ChunkSize         int64
	ChunkCount        int
	Chunks            map[int]bool
	KnowledgeSpaceKey string
}

var uploadSessions = struct {
	sync.RWMutex
	m map[string]*uploadSession
}{m: map[string]*uploadSession{}}

// UploadInit POST /api/media/upload-init —— 初始化分片上传
func UploadInit(c *gin.Context) {
	var req struct {
		Filename          string `json:"filename"`
		FileSize          int64  `json:"file_size"`
		BindType          string `json:"bind_type"`
		BindID            *uint  `json:"bind_id"`
		KnowledgeSpaceKey string `json:"knowledge_space_key"`
	}
	if err := c.ShouldBindJSON(&req); err != nil || req.Filename == "" || req.FileSize <= 0 {
		web.Fail(c, web.NewBadRequest("filename/file_size 必填"))
		return
	}
	ext := extOf(req.Filename)
	if !isUploadableExt(ext) {
		web.Fail(c, web.NewBadRequest("不支持的文件类型"))
		return
	}
	if req.FileSize > sizeLimitFor(ext) {
		web.Fail(c, web.NewBadRequest("文件超过大小限制"))
		return
	}
	if req.BindType == "" || !validBindType[req.BindType] {
		req.BindType = "none"
	}

	id := randomID()
	chunkCount := int((req.FileSize + defaultChunkSize - 1) / defaultChunkSize)
	sess := &uploadSession{
		Filename: req.Filename, FileSize: req.FileSize, Ext: ext,
		BindType: req.BindType, BindID: req.BindID,
		ChunkSize: defaultChunkSize, ChunkCount: chunkCount, Chunks: map[int]bool{},
	}
	sess.KnowledgeSpaceKey = ensureKnowledgeSpaceKeyOrDefault(sanitizeSpaceKey(req.KnowledgeSpaceKey))
	uploadSessions.Lock()
	uploadSessions.m[id] = sess
	uploadSessions.Unlock()

	web.OK(c, gin.H{"upload_id": id, "chunk_size": defaultChunkSize, "chunk_count": chunkCount})
}

// UploadChunk POST /api/media/upload-chunk —— 上传分片
func UploadChunk(c *gin.Context) {
	id := c.PostForm("upload_id")
	idx, err := strconv.Atoi(c.PostForm("chunk_index"))
	if err != nil || id == "" {
		web.Fail(c, web.NewBadRequest("upload_id/chunk_index 必填"))
		return
	}
	uploadSessions.RLock()
	sess := uploadSessions.m[id]
	uploadSessions.RUnlock()
	if sess == nil {
		web.Fail(c, web.NewNotFoundError("上传会话不存在"))
		return
	}
	if idx < 0 || idx >= sess.ChunkCount {
		web.Fail(c, web.NewBadRequest("chunk_index 越界"))
		return
	}

	file, _, err := c.Request.FormFile("file")
	if err != nil {
		web.Fail(c, web.NewBadRequest("缺少分片文件"))
		return
	}
	defer file.Close()

	tmpDir := filepath.Join(pendingDir(), "tmp", id)
	if err := os.MkdirAll(tmpDir, 0o755); err != nil {
		web.Fail(c, web.NewBadRequest("创建临时目录失败"))
		return
	}
	dst := filepath.Join(tmpDir, fmt.Sprintf("%06d", idx))
	out, err := os.Create(dst)
	if err != nil {
		web.Fail(c, web.NewBadRequest("保存分片失败"))
		return
	}
	if _, err := io.Copy(out, file); err != nil {
		out.Close()
		web.Fail(c, web.NewBadRequest("写入分片失败"))
		return
	}
	out.Close()

	uploadSessions.Lock()
	sess.Chunks[idx] = true
	uploadSessions.Unlock()
	web.OK(c, gin.H{"upload_id": id, "chunk_index": idx, "received": len(sess.Chunks)})
}

// UploadComplete POST /api/media/upload-complete —— 合并分片
func UploadComplete(c *gin.Context) {
	u := middleware.CurrentUser(c)
	var req struct {
		UploadID string `json:"upload_id"`
	}
	if err := c.ShouldBindJSON(&req); err != nil || req.UploadID == "" {
		web.Fail(c, web.NewBadRequest("upload_id 必填"))
		return
	}
	uploadSessions.RLock()
	sess := uploadSessions.m[req.UploadID]
	uploadSessions.RUnlock()
	if sess == nil {
		web.Fail(c, web.NewNotFoundError("上传会话不存在"))
		return
	}
	if len(sess.Chunks) != sess.ChunkCount {
		web.Fail(c, web.NewBadRequest(fmt.Sprintf("分片不完整：%d/%d", len(sess.Chunks), sess.ChunkCount)))
		return
	}

	source := sourceOf(c)
	status := "pending"
	targetDir := pendingDir()
	if source == "admin" {
		status = "approved" // 管理员上传自动通过
		targetDir = approvedDir()
	}

	storedName := randomID() + "." + sess.Ext
	if err := os.MkdirAll(targetDir, 0o755); err != nil {
		web.Fail(c, web.NewBadRequest("创建存储目录失败"))
		return
	}
	final := filepath.Join(targetDir, storedName)
	out, err := os.Create(final)
	if err != nil {
		web.Fail(c, web.NewBadRequest("创建最终文件失败"))
		return
	}
	tmpDir := filepath.Join(pendingDir(), "tmp", req.UploadID)
	for i := 0; i < sess.ChunkCount; i++ {
		part := filepath.Join(tmpDir, fmt.Sprintf("%06d", i))
		f, err := os.Open(part)
		if err != nil {
			out.Close()
			web.Fail(c, web.NewBadRequest("读取分片失败"))
			return
		}
		if _, err := io.Copy(out, f); err != nil {
			f.Close()
			out.Close()
			web.Fail(c, web.NewBadRequest("合并分片失败"))
			return
		}
		f.Close()
	}
	out.Close()
	os.RemoveAll(tmpDir)
	uploadSessions.Lock()
	delete(uploadSessions.m, req.UploadID)
	uploadSessions.Unlock()

	m := model.MediaFile{
		Filename: sess.Filename, StoredName: storedName, StoredPath: storedName,
		FileExt: sess.Ext, FileSize: sess.FileSize, Status: status,
		Source: source, SubmitterID: u.ID, BindType: sess.BindType, BindID: sess.BindID,
		KnowledgeSpaceKey: sess.KnowledgeSpaceKey,
	}
	if !mediaRepo.Insert(&m) {
		web.Fail(c, web.NewBadRequest("创建素材记录失败"))
		return
	}
	if status == "approved" {
		go runExtractPipeline(m.ID)
	}
	web.OK(c, gin.H{"media_id": m.ID, "status": m.Status, "filename": m.Filename})
}

// ============ 预览 / 状态 ============

// Preview GET /api/media/preview/{media_id} —— 仅 approved 可预览
func Preview(c *gin.Context) {
	id, ok := parseID(c, "media_id")
	if !ok {
		return
	}
	m, found := mediaRepo.GetByID(id)
	if !found {
		web.Fail(c, web.NewNotFoundError("素材不存在"))
		return
	}
	if m.Status != "approved" {
		web.Fail(c, web.NewForbiddenError("素材未通过审批，不可预览"))
		return
	}
	web.OK(c, gin.H{
		"preview_url": "/media/" + m.StoredName,
		"file_ext":    m.FileExt,
		"can_preview": true,
	})
}

// MediaStatus GET /api/media/status/{media_id}
func MediaStatus(c *gin.Context) {
	id, ok := parseID(c, "media_id")
	if !ok {
		return
	}
	m, found := mediaRepo.GetByID(id)
	if !found {
		web.Fail(c, web.NewNotFoundError("素材不存在"))
		return
	}
	web.OK(c, gin.H{
		"status":              m.Status,
		"extracted":           m.Extracted,
		"chunk_count":         chunkRepo.CountByMediaFile(m.ID),
		"knowledge_space_key": m.KnowledgeSpaceKey,
	})
}

// ============ 审批 ============

// AuditList GET /api/media/audit-list?status=&page=&size=
func AuditList(c *gin.Context) {
	page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
	size, _ := strconv.Atoi(c.DefaultQuery("size", "20"))
	if page < 1 {
		page = 1
	}
	if size < 1 || size > 100 {
		size = 20
	}
	total, items := mediaRepo.ListForAudit(
		c.Query("status"), sanitizeSpaceKey(c.Query("knowledge_space_key")), page, size)
	web.OK(c, gin.H{"total": total, "items": items})
}

// AuditMedia POST /api/media/audit/{media_id} —— 审批（通过→触发异步提取）
func AuditMedia(c *gin.Context) {
	u := middleware.CurrentUser(c)
	id, ok := parseID(c, "media_id")
	if !ok {
		return
	}
	m, found := mediaRepo.GetByID(id)
	if !found {
		web.Fail(c, web.NewNotFoundError("素材不存在"))
		return
	}
	if m.Status != "pending" {
		web.Fail(c, web.NewConflictError("该素材已审批，不可重复操作"))
		return
	}
	var req struct {
		Action       string `json:"action"`
		RejectReason string `json:"reject_reason"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		web.Fail(c, web.NewBadRequest("请求参数错误"))
		return
	}
	now := time.Now()
	auditBy := u.ID
	switch req.Action {
	case "approve":
		// 审批前置：通过后才把文件从 pending/ 移动到 approved/（公开）
		src := filepath.Join(pendingDir(), m.StoredName)
		dst := filepath.Join(approvedDir(), m.StoredName)
		if err := os.MkdirAll(approvedDir(), 0o755); err != nil {
			web.Fail(c, web.NewBadRequest("创建公开目录失败"))
			return
		}
		if err := os.Rename(src, dst); err != nil {
			web.Fail(c, web.NewBadRequest("移动文件失败："+err.Error()))
			return
		}
		m.Status = "approved"
		m.AuditBy = &auditBy
		m.AuditAt = &now
		m.RejectReason = ""
	case "reject":
		if strings.TrimSpace(req.RejectReason) == "" {
			web.Fail(c, web.NewBadRequest("驳回理由必填"))
			return
		}
		// 驳回：把文件从 pending/ 移到 rejected/，落实「审批前置」物理隔离
		src := filepath.Join(pendingDir(), m.StoredName)
		dst := filepath.Join(rejectedDir(), m.StoredName)
		if err := os.MkdirAll(rejectedDir(), 0o755); err != nil {
			web.Fail(c, web.NewBadRequest("创建驳回目录失败"))
			return
		}
		if err := os.Rename(src, dst); err != nil {
			web.Fail(c, web.NewBadRequest("移动文件失败："+err.Error()))
			return
		}
		m.Status = "rejected"
		m.RejectReason = req.RejectReason
		m.AuditBy = &auditBy
		m.AuditAt = &now
	default:
		web.Fail(c, web.NewBadRequest("action 必须为 approve 或 reject"))
		return
	}
	if !mediaRepo.Update(&m) {
		web.Fail(c, web.NewBadRequest("审批失败"))
		return
	}
	if m.Status == "approved" {
		go runExtractPipeline(m.ID)
	}
	web.OK(c, gin.H{"id": m.ID, "status": m.Status})
}

// ============ 异步文档提取管线 ============

// runExtractPipeline 审批通过后：文档转 PDF → pdftotext → 切片入库
func runExtractPipeline(mediaID uint) {
	m, found := mediaRepo.GetByID(mediaID)
	if !found {
		return
	}
	docExts := map[string]bool{"ppt": true, "pptx": true, "doc": true, "docx": true}
	var text string
	var err error
	src := mediaPathFor(m)

	switch {
	case docExts[m.FileExt]:
		text, err = libreofficeExtract(src)
	case m.FileExt == "pdf":
		text, err = pdftotextExtract(src)
	default:
		// 视频/图片：仅标记，不提取
		mediaRepo.MarkExtracted(m.ID)
		return
	}
	if err != nil {
		log.Printf("[提取失败] media_id=%d: %v", mediaID, err)
		return
	}

	chunks := splitIntoChunks(text, 1000)
	spaceKey := resolveMediaKnowledgeSpaceKey(m)
	sourceID := strconv.FormatUint(uint64(m.ID), 10)
	items := make([]model.KnowledgeChunk, 0, len(chunks))
	for i, chunk := range chunks {
		if strings.TrimSpace(chunk) == "" {
			continue
		}
		items = append(items, model.KnowledgeChunk{
			MediaFileID:       &m.ID,
			SourceType:        m.FileExt,
			SourceID:          sourceID,
			KnowledgeSpaceKey: spaceKey,
			ChunkIndex:        i,
			Content:           chunk,
		})
	}
	chunkRepo.BulkInsert(items)
	mediaRepo.MarkExtracted(m.ID)
	triggerKnowledgeIndexRebuild()
	log.Printf("[提取完成] media_id=%d chunks=%d", mediaID, len(chunks))
}

func libreofficeExtract(src string) (string, error) {
	tmp, err := os.MkdirTemp("", "lo-convert")
	if err != nil {
		return "", err
	}
	defer os.RemoveAll(tmp)
	// 独立 UserInstallation：systemd 加固（ProtectHome=true）下 HOME 不可写，
	// 避免 LibreOffice profile 初始化失败。
	profileURI := "file://" + filepath.Join(tmp, "profile")
	cmd := exec.Command(Cfg.LibreOfficeBin, "-env:UserInstallation="+profileURI,
		"--headless", "--convert-to", "pdf", "--outdir", tmp, src)
	if out, err := cmd.CombinedOutput(); err != nil {
		return "", fmt.Errorf("libreoffice 转换失败: %v %s", err, out)
	}
	base := strings.TrimSuffix(filepath.Base(src), filepath.Ext(src)) + ".pdf"
	return pdftotextExtract(filepath.Join(tmp, base))
}

func pdftotextExtract(pdf string) (string, error) {
	cmd := exec.Command(Cfg.PdftotextBin, "-enc", "UTF-8", pdf, "-")
	out, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("pdftotext 提取失败: %v", err)
	}
	return string(out), nil
}

func splitIntoChunks(text string, maxChars int) []string {
	paragraphs := strings.Split(text, "\n\n")
	var chunks []string
	current := ""
	for _, p := range paragraphs {
		p = strings.TrimSpace(p)
		if p == "" {
			continue
		}
		if len(current)+len(p) > maxChars {
			if current != "" {
				chunks = append(chunks, current)
			}
			current = p
		} else if current == "" {
			current = p
		} else {
			current += "\n\n" + p
		}
	}
	if current != "" {
		chunks = append(chunks, current)
	}
	return chunks
}