schedule_converter/service/recognizer.go

package service

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	"schedule_converter/client"
	"schedule_converter/config"
	"schedule_converter/models"
	apperr "schedule_converter/pkg/errors"
)

var (
	jsonBlockRegex = regexp.MustCompile("(?s)```json\\s*(\\[[\\s\\S]*?\\])\\s*```")
	codeBlockRegex = regexp.MustCompile("(?s)```\\s*(\\[[\\s\\S]*?\\])\\s*```")
	jsonArrayRegex = regexp.MustCompile(`(\[[\s\S]*\])`)
	regexPool      = sync.Pool{
		New: func() any {
			return nil
		},
	}
)

type RecognizerService struct {
	client    *http.Client
	maxRetry  int
	retryWait time.Duration
}

func NewRecognizerService() *RecognizerService {
	return &RecognizerService{
		client:    client.RecognizerClient,
		maxRetry:  3,
		retryWait: 5 * time.Second,
	}
}

func (r *RecognizerService) RecognizeFromFile(ctx context.Context, imagePath string) ([]models.Course, error) {
	log := slog.With("image_path", imagePath)
	log.Info("正在读取图片")

	imageData, err := os.ReadFile(imagePath)
	if err != nil {
		return nil, apperr.New("recognize_file", err, "读取图片失败")
	}

	log.Info("图片读取完成", "size", len(imageData))
	return r.RecognizeFromBytes(ctx, imageData)
}

func (r *RecognizerService) RecognizeFromBytes(ctx context.Context, imageData []byte) ([]models.Course, error) {
	log := slog.With("size", len(imageData))
	log.Info("正在处理图片数据")

	imageBase64 := base64.StdEncoding.EncodeToString(imageData)
	return r.callAPI(ctx, imageBase64)
}

func (r *RecognizerService) RecognizeFromBase64(ctx context.Context, imageBase64 string) ([]models.Course, error) {
	slog.Info("正在处理base64图片数据", "length", len(imageBase64))
	return r.callAPI(ctx, imageBase64)
}

func (r *RecognizerService) callAPI(ctx context.Context, imageBase64 string) ([]models.Course, error) {
	var lastErr error
	for retry := range r.maxRetry {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}

		if retry > 0 {
			slog.Info("重试中", "attempt", retry+1, "max_retry", r.maxRetry)
			time.Sleep(r.retryWait)
		}

		courses, err := r.callAPISingle(ctx, imageBase64)
		if err == nil {
			return courses, nil
		}

		lastErr = err
		slog.Warn("API调用失败", "attempt", retry+1, "error", err)

		if !strings.Contains(err.Error(), "504") {
			return nil, err
		}
	}

	return nil, apperr.New("recognize_api", lastErr, "已达到最大重试次数")
}

func (r *RecognizerService) callAPISingle(ctx context.Context, imageBase64 string) ([]models.Course, error) {
	cfg := config.Get()
	if cfg.Recognizer.APIKey == "" {
		return nil, apperr.ErrRecognitionFailed
	}

	prompt := `请解析以下课表图片，提取所有课程信息并以JSON数组格式返回。

JSON格式要求：
[{"name":"课程名","teacher":"老师名","room":"教室","schedule":[{"day":"星期X","section":[节次数组],"weeks":[周数数组]}]}]

注意事项：
1. day使用"星期一"到"星期日"
2. section节次：上午1-2节为[1,2]，上午3-4节为[3,4]，下午5-6节为[5,6]，下午7-8节为[7,8]，晚上9-10节为[9,10]，晚上11-12节为[11,12]
3. weeks周数：格式如[1,2,3,4,5,6,7,8]或[9,10,11,12,13,14,15,16]
4. 只返回JSON数组，不要任何其他文字。`

	falseVal := false
	zero := 0
	apiReq := models.APIRequest{
		Model:          cfg.Recognizer.Model,
		MaxTokens:      4000,
		EnableThinking: &falseVal,
		ThinkingBudget: &zero,
		Messages: []models.Message{
			{
				Role: "user",
				Content: []models.ContentItem{
					{Type: "text", Text: prompt},
					{
						Type: "image_url",
						ImageURL: &models.ImageURL{
							URL: "data:image/png;base64," + imageBase64,
						},
					},
				},
			},
		},
	}

	jsonData, err := json.Marshal(apiReq)
	if err != nil {
		return nil, apperr.New("marshal_request", err, "")
	}

	req, err := http.NewRequestWithContext(ctx, "POST", cfg.Recognizer.APIURL, bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, apperr.New("create_request", err, "")
	}

	req.Header = http.Header{
		"Content-Type":  []string{"application/json"},
		"Authorization": []string{"Bearer " + cfg.Recognizer.APIKey},
	}

	slog.Info("调用识别API")
	resp, err := r.client.Do(req)
	if err != nil {
		return nil, apperr.New("api_call", err, "")
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		respBody, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("API调用失败: %d, %s", resp.StatusCode, string(respBody))
	}

	var apiResp models.APIResponse
	if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil {
		return nil, apperr.New("decode_response", err, "")
	}

	if len(apiResp.Choices) == 0 {
		return nil, apperr.New("empty_response", apperr.ErrRecognitionFailed, "API返回为空")
	}

	resultText := extractContent(apiResp.Choices[0].Message.Content)
	slog.Debug("API响应接收完成", "response_length", len(resultText))

	jsonMatch := extractJSON(resultText)
	if jsonMatch == "" {
		return nil, apperr.New("extract_json", apperr.ErrRecognitionFailed,
			fmt.Sprintf("未找到JSON数据，原始响应: %s", resultText))
	}

	var courses []models.Course
	if err := json.Unmarshal([]byte(jsonMatch), &courses); err != nil {
		return nil, apperr.New("parse_json", err, fmt.Sprintf("数据: %s", jsonMatch))
	}

	slog.Info("识别完成", "course_count", len(courses))
	return courses, nil
}

func extractContent(content any) string {
	switch c := content.(type) {
	case string:
		return c
	case []any:
		for _, item := range c {
			if itemMap, ok := item.(map[string]any); ok {
				if text, ok := itemMap["text"].(string); ok {
					return text
				}
			}
		}
	}
	return ""
}

func extractJSON(text string) string {
	if match := jsonBlockRegex.FindStringSubmatch(text); len(match) > 1 {
		slog.Debug("从json代码块中提取JSON")
		return match[1]
	}
	if match := codeBlockRegex.FindStringSubmatch(text); len(match) > 1 {
		slog.Debug("从代码块中提取JSON")
		return match[1]
	}
	if match := jsonArrayRegex.FindString(text); match != "" {
		slog.Debug("从文本中提取JSON")
		return match
	}
	return ""
}

func RecognizeScheduleFromFile(imagePath string) ([]models.Course, error) {
	return NewRecognizerService().RecognizeFromFile(context.Background(), imagePath)
}

func RecognizeScheduleFromBytes(imageData []byte) ([]models.Course, error) {
	return NewRecognizerService().RecognizeFromBytes(context.Background(), imageData)
}

func RecognizeScheduleFromBase64(imageBase64 string) ([]models.Course, error) {
	return NewRecognizerService().RecognizeFromBase64(context.Background(), imageBase64)
}