lone-services/services/analysis/internal/callback/faceback.go

package callback

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"lone-services/pkg/modelbase"
	pkgoss "lone-services/pkg/oss"
	"lone-services/pkg/utils"
	"lone-services/pkg/wechat/qrcode"
	analysisconfig "lone-services/services/analysis/internal/config"
	"lone-services/services/analysis/internal/dao"
	"lone-services/services/analysis/internal/model"

	"github.com/zeromicro/go-zero/core/logx"
)

type ImageRule struct {
	Scene    string
	Position string
}

type FrontItem struct {
	Name string `json:"name"`
	Path string `json:"path"`
}

type AnalysisSummary struct {
	HealthScore         string `json:"health_score"`
	NormalArea          string `json:"normal_area"`
	MildAbnormality     string `json:"mild_abnormality"`
	ModerateAbnormality string `json:"moderate_abnormality"`
	SevereAbnormality   string `json:"severe_abnormality"`
	PredictedAge        string `json:"predicted_age"`
}

func ProcessPyCallback(subTaskId string) error {
	sub, err := GetSubTask(subTaskId)
	if err != nil {
		return err
	}
	if err = FinishSubTask(subTaskId); err != nil {
		return err
	}

	subCopy := *sub
	go func(s dao.ImageAnalysisSub) {
		defer func() {
			if r := recover(); r != nil {
				logx.Errorf("py callback async panic: %v", r)
			}
		}()
		if ossErr := UploadSubImagesToOSS(&s); ossErr != nil {
			logx.Errorf("upload sub images to oss failed: %v", ossErr)
		}
		if finishErr := TryFinishTask(s.TaskId); finishErr != nil {
			logx.Errorf("try finish task failed: %v", finishErr)
		}
	}(subCopy)

	return nil
}

func GetSubTask(subTaskId string) (*dao.ImageAnalysisSub, error) {
	var data dao.ImageAnalysisSub
	err := model.ImageAnalysisSubModel{}.Init().GetOne(modelbase.Params{
		Eq: map[string]string{"sub_task_id": subTaskId},
	}, &data)
	if err != nil {
		return nil, err
	}
	if data.Id < utils.NumberOne {
		return nil, fmt.Errorf("sub task not found: %s", subTaskId)
	}
	return &data, nil
}

func FinishSubTask(subTaskId string) error {
	_, err := model.ImageAnalysisSubModel{}.Init().Edit(modelbase.Params{
		Eq: map[string]string{"sub_task_id": subTaskId},
	}, map[string]interface{}{
		"status":         dao.AnalysisStatusDone,
		"process_end_ts": time.Now().UnixMilli(),
		"updated_at":     utils.Now(),
	})
	return err
}

func UploadSubImagesToOSS(sub *dao.ImageAnalysisSub) error {
	if sub == nil || sub.SubTaskId == utils.StringEmpty {
		return fmt.Errorf("sub task empty")
	}

	imageDir := sub.ImageDir
	if imageDir == utils.StringEmpty {
		imageDir = subImageDir(sub.FilePath)
	}
	if imageDir == utils.StringEmpty {
		return fmt.Errorf("sub image dir empty")
	}

	scene, originalPaths, err := detectOriginalLighting(imageDir)
	if err != nil {
		return err
	}

	needResult := sceneHasResult(scene)
	originalsReady := sub.LeftPath != utils.StringEmpty && sub.MidPath != utils.StringEmpty && sub.RightPath != utils.StringEmpty
	resultsReady := !needResult || (sub.LeftResultPath != utils.StringEmpty && sub.MidResultPath != utils.StringEmpty && sub.RightResultPath != utils.StringEmpty)
	if originalsReady && resultsReady {
		return nil
	}

	update := map[string]interface{}{
		"scene":      scene,
		"updated_at": utils.Now(),
	}

	if !originalsReady {
		leftURL, midURL, rightURL, upErr := uploadPositionImages(sub.TaskId, originalPaths, true)
		if upErr != nil {
			return fmt.Errorf("upload original images: %w", upErr)
		}
		update["left_path"] = leftURL
		update["mid_path"] = midURL
		update["right_path"] = rightURL
		sub.LeftPath = leftURL
		sub.MidPath = midURL
		sub.RightPath = rightURL
	}

	if needResult && !resultsReady {
		analysis, aErr := GetAnalysis(sub.TaskId)
		if aErr != nil {
			return aErr
		}
		if analysis.FolderPath == utils.StringEmpty {
			return fmt.Errorf("analysis folder_path empty task_id=%s", sub.TaskId)
		}
		resultPaths, pErr := resultImagePaths(analysis.FolderPath, scene)
		if pErr != nil {
			return pErr
		}
		leftURL, midURL, rightURL, upErr := uploadPositionImages(sub.TaskId, resultPaths, false)
		if upErr != nil {
			return fmt.Errorf("upload result images: %w", upErr)
		}
		update["left_result_path"] = leftURL
		update["mid_result_path"] = midURL
		update["right_result_path"] = rightURL
	}

	_, err = model.ImageAnalysisSubModel{}.Init().Edit(modelbase.Params{
		Eq: map[string]string{"sub_task_id": sub.SubTaskId},
	}, update)
	return err
}

func sceneHasResult(scene string) bool {
	switch scene {
	case "uv405", "cross", "parallel", "uv365":
		return true
	default:
		return false
	}
}

func detectOriginalLighting(imageDir string) (string, map[string]string, error) {
	entries, err := os.ReadDir(imageDir)
	if err != nil {
		return utils.StringEmpty, nil, err
	}

	scene := utils.StringEmpty
	paths := map[string]string{}
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		rule, ok := parseImageRule(entry.Name())
		if !ok || !strings.HasPrefix(rule.Scene, "original_") {
			continue
		}
		light := strings.TrimPrefix(rule.Scene, "original_")
		if scene == utils.StringEmpty {
			scene = light
		}
		if light != scene {
			continue
		}
		paths[rule.Position] = filepath.Join(imageDir, entry.Name())
	}
	if scene == utils.StringEmpty || paths["left"] == "" || paths["mid"] == "" || paths["right"] == "" {
		return utils.StringEmpty, nil, fmt.Errorf("no complete original lighting in %s", imageDir)
	}
	return scene, paths, nil
}

func resultImagePaths(folderPath, scene string) (map[string]string, error) {
	var dir string
	var names map[string]string
	switch scene {
	case "uv405":
		dir = "output/yolo"
		names = map[string]string{
			"left":  "yolo_uv405_left.jpg",
			"mid":   "yolo_uv405_mid.jpg",
			"right": "yolo_uv405_right.jpg",
		}
	case "cross":
		dir = "output/pigment"
		names = map[string]string{
			"left":  "pigment_cross_left.jpg",
			"mid":   "pigment_cross_mid.jpg",
			"right": "pigment_cross_right.jpg",
		}
	case "parallel":
		dir = "output/sensitivity"
		names = map[string]string{
			"left":  "sensitivity_parallel_left.jpg",
			"mid":   "sensitivity_parallel_mid.jpg",
			"right": "sensitivity_parallel_right.jpg",
		}
	case "uv365":
		dir = "output/redzone"
		names = map[string]string{
			"left":  "red_uv365_left.jpg",
			"mid":   "red_uv365_mid.jpg",
			"right": "red_uv365_right.jpg",
		}
	default:
		return nil, fmt.Errorf("scene %s has no result images", scene)
	}

	paths := map[string]string{}
	for pos, name := range names {
		full := filepath.Join(folderPath, dir, name)
		if _, err := os.Stat(full); err != nil {
			return nil, fmt.Errorf("result image missing: %s", full)
		}
		paths[pos] = full
	}
	return paths, nil
}

func uploadPositionImages(taskId string, paths map[string]string, rotate bool) (string, string, string, error) {
	leftURL, err := uploadImageToOSS(taskId, paths["left"], rotate)
	if err != nil {
		return utils.StringEmpty, utils.StringEmpty, utils.StringEmpty, fmt.Errorf("left: %w", err)
	}
	midURL, err := uploadImageToOSS(taskId, paths["mid"], rotate)
	if err != nil {
		return utils.StringEmpty, utils.StringEmpty, utils.StringEmpty, fmt.Errorf("mid: %w", err)
	}
	rightURL, err := uploadImageToOSS(taskId, paths["right"], rotate)
	if err != nil {
		return utils.StringEmpty, utils.StringEmpty, utils.StringEmpty, fmt.Errorf("right: %w", err)
	}
	return leftURL, midURL, rightURL, nil
}

func subImageDir(filePath string) string {
	if filePath == utils.StringEmpty {
		return utils.StringEmpty
	}
	if info, err := os.Stat(filePath); err == nil && info.IsDir() {
		return filePath
	}
	ext := filepath.Ext(filePath)
	if strings.EqualFold(ext, ".zip") {
		return strings.TrimSuffix(filePath, ext)
	}
	return filepath.Dir(filePath)
}

func uploadImageToOSS(taskId, localPath string, rotate bool) (string, error) {
	_ = taskId
	data, err := os.ReadFile(localPath)
	if err != nil {
		return utils.StringEmpty, err
	}
	if rotate {
		data, err = Rotate90Bytes(data, filepath.Ext(localPath))
		if err != nil {
			return utils.StringEmpty, fmt.Errorf("rotate image: %w", err)
		}
	}

	ext := strings.ToLower(filepath.Ext(localPath))
	if ext == utils.StringEmpty {
		ext = ".jpg"
	}
	fileName := fmt.Sprintf("%d%d%s", time.Now().UnixNano(), rand.Intn(10000), ext)
	contentType := "image/jpeg"
	if ext == ".png" {
		contentType = "image/png"
	}

	imageDir := utils.GetConfigString("oss.imageDir")
	res, err := pkgoss.Upload(pkgoss.UploadInput{
		Filename:    fileName,
		Content:     data,
		ContentType: contentType,
		Dir:         imageDir,
	})
	if err != nil {
		return utils.StringEmpty, err
	}
	return res.URL, nil
}

func ListFinishedSubs(taskId string) ([]dao.ImageAnalysisSub, error) {
	var list []dao.ImageAnalysisSub
	err := model.ImageAnalysisSubModel{}.Init().Items(modelbase.Params{
		Eq: map[string]string{
			"task_id": taskId,
			"status":  strconv.Itoa(int(dao.AnalysisStatusDone)),
		},
		Order: "id asc",
	}, &list)
	return list, err
}

func TryFinishTask(taskId string) error {
	count, err := CountFinishedTask(taskId)
	if err != nil || count < 6 {
		return err
	}

	subs, err := ListFinishedSubs(taskId)
	if err != nil {
		return err
	}
	if len(subs) < 6 {
		return nil
	}
	for _, sub := range subs {
		if sub.LeftPath == utils.StringEmpty || sub.MidPath == utils.StringEmpty || sub.RightPath == utils.StringEmpty {
			return nil
		}
		if sceneHasResult(sub.Scene) && (sub.LeftResultPath == utils.StringEmpty || sub.MidResultPath == utils.StringEmpty || sub.RightResultPath == utils.StringEmpty) {
			return nil
		}
	}

	analysis, err := GetAnalysis(taskId)
	if err != nil {
		return err
	}
	if analysis.Status == dao.AnalysisStatusDone {
		return nil
	}

	summary, err := ParseSummary(analysis.FolderPath)
	if err != nil {
		logx.Errorf("read summary failed: %v", err)
	}

	analysis.Cover = UploadCoverToOSS(taskId, analysis.Cover)
	if err = FinishAnalysis(taskId, analysis); err != nil {
		return err
	}
	userImage, err := FinishUserImage(analysis, summary)
	if err != nil {
		return err
	}
	if qrErr := EnsureReportQrcode(taskId, userImage.Id); qrErr != nil {
		logx.Errorf("ensure report qrcode failed: %v", qrErr)
	}
	return nil
}

func CountFinishedTask(taskId string) (int64, error) {
	return model.ImageAnalysisSubModel{}.Init().CountNumber(modelbase.Params{
		Eq: map[string]string{
			"task_id": taskId,
			"status":  strconv.Itoa(int(dao.AnalysisStatusDone)),
		},
	})
}

func EnsureReportQrcode(taskId string, imageId int64) error {
	if taskId == utils.StringEmpty {
		return fmt.Errorf("task_id empty")
	}
	if imageId <= 0 {
		return fmt.Errorf("image id invalid")
	}

	analysis, err := GetAnalysis(taskId)
	if err != nil {
		return err
	}
	if analysis.Qrcode != utils.StringEmpty {
		return nil
	}

	wxCfg, err := analysisconfig.LoadWechatConfig()
	if err != nil {
		return err
	}
	img, err := wxCfg.NewQrcodeClient().CreateWxaCode(context.Background(), &qrcode.WxaCodeRequest{
		Path: qrcode.BuildPath(wxCfg.ReportPage, map[string]string{
			"type": strconv.Itoa(utils.NumberOne),
			"ids":  strconv.FormatInt(imageId, utils.NumberTen),
		}),
		Width:      wxCfg.Width,
		EnvVersion: wxCfg.EnvVersion,
		AutoColor:  false,
		LineColor:  &qrcode.LineColor{R: 255, G: 255, B: 255},
		IsHyaline:  true,
	})
	if err != nil {
		return fmt.Errorf("fetch mini qrcode: %w", err)
	}

	imageDir := utils.GetConfigString("oss.imageDir")
	fileName := fmt.Sprintf("%d%d.png", time.Now().UnixNano(), rand.Intn(10000))
	res, err := pkgoss.Upload(pkgoss.UploadInput{
		Filename:    fileName,
		Content:     img,
		ContentType: "image/png",
		Dir:         imageDir,
	})
	if err != nil {
		return err
	}

	affected, err := model.ImageAnalysisModel{}.Init().Edit(modelbase.Params{
		Eq: map[string]string{
			"task_id": taskId,
			"qrcode":  "",
		},
	}, map[string]interface{}{
		"qrcode":     res.URL,
		"updated_at": utils.Now(),
	})
	if err != nil {
		return fmt.Errorf("save qrcode: %w", err)
	}
	if affected == 0 {
		logx.Infof("report qrcode already set by concurrent callback task_id=%s", taskId)
		return nil
	}
	logx.Infof("report qrcode saved task_id=%s image_id=%d qrcode=%s", taskId, imageId, res.URL)
	return nil
}

func GetAnalysis(taskId string) (*dao.ImageAnalysis, error) {
	var task dao.ImageAnalysis
	err := model.ImageAnalysisModel{}.Init().GetOne(modelbase.Params{
		Eq: map[string]string{"task_id": taskId},
	}, &task)
	if err != nil {
		return nil, err
	}
	if task.Id < utils.NumberOne {
		return nil, fmt.Errorf("analysis not found: %s", taskId)
	}
	return &task, nil
}

func ParseSummary(folderPath string) (*AnalysisSummary, error) {
	txtPath := path.Join(folderPath, "output/texts", "analysis_report.txt")
	data, err := os.ReadFile(txtPath)
	if err != nil {
		return nil, err
	}

	result := &AnalysisSummary{}
	items := strings.Split(string(data), ";;")
	for _, item := range items {
		item = strings.TrimSpace(item)
		if item == utils.StringEmpty {
			continue
		}
		parts := strings.Split(item, "::")
		if len(parts) != utils.NumberTwo {
			continue
		}
		key := strings.TrimSpace(parts[0])
		value := strings.TrimSpace(parts[1])
		switch key {
		case "Health Score":
			result.HealthScore = value
		case "Normal area":
			result.NormalArea = value
		case "Mild abnormality":
			result.MildAbnormality = value
		case "Moderate abnormality":
			result.ModerateAbnormality = value
		case "Severe abnormality":
			result.SevereAbnormality = value
		case "Predicted Age":
			result.PredictedAge = value
		}
	}
	return result, nil
}

func FinishAnalysis(taskID string, analysis *dao.ImageAnalysis) error {
	now := utils.Now()
	_, err := model.ImageAnalysisModel{}.Init().Edit(modelbase.Params{
		Eq: map[string]string{"task_id": taskID},
	}, map[string]interface{}{
		"status":     dao.AnalysisStatusDone,
		"cover":      analysis.Cover,
		"updated_at": now,
	})
	if err != nil {
		return err
	}

	wLog := modelbase.Params{
		Eq: map[string]string{
			"user_id": strconv.FormatInt(analysis.UserId, utils.NumberTen),
		},
	}
	var logRow dao.UserImageLog
	err = model.UserImageLogModel{}.Init().GetOne(wLog, &logRow)
	if err != nil {
		return err
	}

	if logRow.Id < utils.NumberOne {
		logModel := model.UserImageLogModel{}.Init()
		return logModel.Create(&dao.UserImageLog{
			UserId:          analysis.UserId,
			DetectionNumber: utils.NumberOne,
			Cover:           analysis.Cover,
			DetectionTime:   now,
			Status:          uint8(utils.StatusOk),
		})
	}

	newNumber := logRow.DetectionNumber + utils.NumberOne
	updateData := map[string]interface{}{
		"cover":            analysis.Cover,
		"detection_time":   now,
		"detection_number": newNumber,
	}
	if newNumber == utils.NumberTwo {
		updateData["status"] = utils.NumberTwo
	}
	_, err = model.UserImageLogModel{}.Init().Edit(wLog, updateData)
	return err
}

func FinishUserImage(analysis *dao.ImageAnalysis, summaryStruct *AnalysisSummary) (*dao.Image, error) {
	subs, err := ListFinishedSubs(analysis.TaskId)
	if err != nil {
		return nil, err
	}
	if len(subs) == 0 {
		return nil, errors.New("no finished sub tasks")
	}

	userImage, err := BuildUserImageFromSubs(analysis, summaryStruct, subs)
	if err != nil {
		return nil, err
	}
	imgModel := model.ImageModel{}.Init()
	if err = imgModel.Create(userImage); err != nil {
		return nil, err
	}
	return userImage, nil
}

func BuildUserImageFromSubs(analysis *dao.ImageAnalysis, summaryStruct *AnalysisSummary, subs []dao.ImageAnalysisSub) (*dao.Image, error) {
	summaryBytes, err := json.Marshal(summaryStruct)
	if err != nil {
		return nil, err
	}

	images := make(map[string]map[string]string)
	var front []FrontItem

	for _, sub := range subs {
		if sub.Scene == utils.StringEmpty {
			continue
		}
		if sub.LeftPath == utils.StringEmpty || sub.MidPath == utils.StringEmpty || sub.RightPath == utils.StringEmpty {
			continue
		}
		originalScene := "original_" + sub.Scene
		images[originalScene] = map[string]string{
			"left":  sub.LeftPath,
			"mid":   sub.MidPath,
			"right": sub.RightPath,
		}
		front = append(front, FrontItem{
			Name: originalScene,
			Path: sub.MidPath,
		})

		if sub.LeftResultPath == utils.StringEmpty || sub.MidResultPath == utils.StringEmpty || sub.RightResultPath == utils.StringEmpty {
			continue
		}
		images[sub.Scene] = map[string]string{
			"left":  sub.LeftResultPath,
			"mid":   sub.MidResultPath,
			"right": sub.RightResultPath,
		}
		front = append(front, FrontItem{
			Name: sub.Scene,
			Path: sub.MidResultPath,
		})
	}

	if len(images) == 0 {
		return nil, errors.New("no oss images on finished sub tasks")
	}

	result := map[string]interface{}{
		"images": images,
		"front":  front,
	}
	imagesBytes, err := json.Marshal(result)
	if err != nil {
		return nil, err
	}

	return &dao.Image{
		UserId:     analysis.UserId,
		TaskId:     analysis.TaskId,
		Images:     string(imagesBytes),
		Contents:   string(summaryBytes),
		Cover:      analysis.Cover,
		Status:     int8(utils.StatusOk),
		CreateTime: utils.Now(),
	}, nil
}

func UploadCoverToOSS(taskId, localPath string) string {
	if localPath == utils.StringEmpty {
		return utils.StringEmpty
	}
	if strings.HasPrefix(localPath, "http://") || strings.HasPrefix(localPath, "https://") {
		return localPath
	}

	candidates := []string{localPath}
	if !filepath.IsAbs(localPath) {
		candidates = append(candidates, filepath.Join(utils.GetConfigString("file.upload_path"), localPath))
	}

	for _, fullPath := range candidates {
		if _, err := os.Stat(fullPath); err != nil {
			continue
		}
		url, err := uploadImageToOSS(taskId, fullPath, true)
		if err != nil {
			logx.Errorf("upload cover to oss failed: %v", err)
			return localPath
		}
		return url
	}

	logx.Errorf("upload cover skipped, file not found: %s", localPath)
	return localPath
}

func parseImageRule(fileName string) (ImageRule, bool) {
	switch fileName {
	case "yolo_uv405_left.jpg":
		return ImageRule{"uv405", "left"}, true
	case "yolo_uv405_mid.jpg":
		return ImageRule{"uv405", "mid"}, true
	case "yolo_uv405_right.jpg":
		return ImageRule{"uv405", "right"}, true
	case "pigment_cross_left.jpg":
		return ImageRule{"cross", "left"}, true
	case "pigment_cross_mid.jpg":
		return ImageRule{"cross", "mid"}, true
	case "pigment_cross_right.jpg":
		return ImageRule{"cross", "right"}, true
	case "sensitivity_cross_left.jpg":
		return ImageRule{"parallel", "left"}, true
	case "sensitivity_cross_mid.jpg":
		return ImageRule{"parallel", "mid"}, true
	case "sensitivity_cross_right.jpg":
		return ImageRule{"parallel", "right"}, true
	case "red_uv365_left.jpg":
		return ImageRule{"uv365", "left"}, true
	case "red_uv365_mid.jpg":
		return ImageRule{"uv365", "mid"}, true
	case "red_uv365_right.jpg":
		return ImageRule{"uv365", "right"}, true
	}

	if strings.HasPrefix(fileName, "parallel_") {
		return buildOriginalRule(fileName, "original_parallel")
	}
	if strings.HasPrefix(fileName, "cross_") {
		return buildOriginalRule(fileName, "original_cross")
	}
	if strings.HasPrefix(fileName, "visible_") {
		return buildOriginalRule(fileName, "original_visible")
	}
	if strings.HasPrefix(fileName, "uv365_") {
		return buildOriginalRule(fileName, "original_uv365")
	}
	if strings.HasPrefix(fileName, "uv385_") {
		return buildOriginalRule(fileName, "original_uv385")
	}
	if strings.HasPrefix(fileName, "uv405_") {
		return buildOriginalRule(fileName, "original_uv405")
	}
	return ImageRule{}, false
}

func buildOriginalRule(fileName string, scene string) (ImageRule, bool) {
	switch {
	case strings.HasSuffix(fileName, "_left.jpg"):
		return ImageRule{Scene: scene, Position: "left"}, true
	case strings.HasSuffix(fileName, "_mid.jpg"):
		return ImageRule{Scene: scene, Position: "mid"}, true
	case strings.HasSuffix(fileName, "_right.jpg"):
		return ImageRule{Scene: scene, Position: "right"}, true
	}
	return ImageRule{}, false
}