schedule_converter/recognizer/recognizer.go

package recognizer

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

	"jwts/config"
	"jwts/models"
)

// APIRequest API请求结构
type APIRequest struct {
	Model          string    `json:"model"`
	Messages       []Message `json:"messages"`
	MaxTokens      int       `json:"max_tokens"`
	EnableThinking *bool     `json:"enable_thinking,omitempty"`
	ThinkingBudget *int      `json:"thinking_budget,omitempty"`
}

// Message 消息结构
type Message struct {
	Role    string      `json:"role"`
	Content interface{} `json:"content"`
}

// ContentItem 内容项
type ContentItem struct {
	Type     string    `json:"type,omitempty"`
	Text     string    `json:"text,omitempty"`
	ImageURL *ImageURL `json:"image_url,omitempty"`
}

// ImageURL 图片URL
type ImageURL struct {
	URL string `json:"url"`
}

// APIResponse API响应结构
type APIResponse struct {
	Choices []Choice `json:"choices"`
}

// Choice 选择
type Choice struct {
	Message Message `json:"message"`
}

// 全局HTTP客户端
var client *http.Client

// 初始化HTTP客户端
func init() {
	client = &http.Client{
		Timeout: 120 * time.Second,
	}

	if config.RecognizerAPIKey == "" {
		log.Println("警告: RECOGNIZER_API_KEY 未设置，图片识别功能将无法使用")
	}

	log.Println("recognizer模块已初始化")
}

// RecognizeFromFile 从图片文件识别课表
func RecognizeFromFile(imagePath string) ([]models.Course, error) {
	fmt.Printf("[识别] 正在读取图片: %s\n", imagePath)

	// 读取图片文件
	imageData, err := os.ReadFile(imagePath)
	if err != nil {
		return nil, fmt.Errorf("读取图片失败: %v", err)
	}

	// 将图片转为base64
	imageBase64 := base64.StdEncoding.EncodeToString(imageData)
	fmt.Printf("[识别] 图片大小: %d bytes\n", len(imageData))

	return callAPI(imageBase64)
}

// RecognizeFromBytes 从图片数据识别课表
func RecognizeFromBytes(imageData []byte) ([]models.Course, error) {
	fmt.Printf("[识别] 正在处理图片数据, 大小: %d bytes\n", len(imageData))

	// 将图片转为base64
	imageBase64 := base64.StdEncoding.EncodeToString(imageData)

	return callAPI(imageBase64)
}

// RecognizeFromBase64 从base64数据识别课表
func RecognizeFromBase64(imageBase64 string) ([]models.Course, error) {
	fmt.Printf("[识别] 正在处理base64图片数据, 长度: %d\n", len(imageBase64))

	return callAPI(imageBase64)
}

// callAPI 调用API识别课表（带重试机制）
func callAPI(imageBase64 string) ([]models.Course, error) {
	// 重试机制
	maxRetries := 3
	for retry := 0; retry < maxRetries; retry++ {
		if retry > 0 {
			fmt.Printf("[识别] 第 %d 次重试...\n", retry+1)
			time.Sleep(5 * time.Second) // 等待5秒后重试
		}

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

		fmt.Printf("[识别] 尝试 %d/%d 失败: %v\n", retry+1, maxRetries, err)

		// 如果是504超时，继续重试；否则直接返回错误
		if !strings.Contains(err.Error(), "504") {
			return nil, err
		}
	}

	return nil, fmt.Errorf("已达到最大重试次数")
}

// callAPISingle 单次API调用
func callAPISingle(imageBase64 string) ([]models.Course, error) {
	if config.RecognizerAPIKey == "" {
		return nil, fmt.Errorf("RECOGNIZER_API_KEY 未配置")
	}

	falseVal := false
	zero := 0
	apiReq := APIRequest{
		Model:          config.RecognizerModel,
		MaxTokens:      4000,
		EnableThinking: &falseVal,
		ThinkingBudget: &zero,
		Messages: []Message{
			{
				Role: "user",
				Content: []ContentItem{
					{
						Type: "text",
						Text: `请解析以下课表图片，提取所有课程信息并以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数组，不要任何其他文字。`,
					},
					{
						Type: "image_url",
						ImageURL: &ImageURL{
							URL: "data:image/png;base64," + imageBase64,
						},
					},
				},
			},
		},
	}

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

	httpReq, err := http.NewRequest("POST", config.RecognizerAPIURL, bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, err
	}

	httpReq.Header = http.Header{
		"Content-Type":  []string{"application/json"},
		"Authorization": []string{"Bearer " + config.RecognizerAPIKey},
	}

	fmt.Println("[识别] 正在调用API进行识别...")
	resp, err := client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("API调用失败: %v", 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 APIResponse
	err = json.NewDecoder(resp.Body).Decode(&apiResp)
	if err != nil {
		return nil, fmt.Errorf("API响应解析失败: %v", err)
	}

	if len(apiResp.Choices) == 0 {
		return nil, fmt.Errorf("API返回为空")
	}

	// 处理Content，可能是string或[]ContentItem
	var resultText string
	switch c := apiResp.Choices[0].Message.Content.(type) {
	case string:
		resultText = c
	case []interface{}:
		for _, item := range c {
			if itemMap, ok := item.(map[string]interface{}); ok {
				if text, ok := itemMap["text"].(string); ok {
					resultText = text
					break
				}
			}
		}
	}

	fmt.Println("[识别] API响应已接收，正在解析...")
	fmt.Printf("[识别] 原始响应: %s\n", resultText)

	// 提取JSON部分
	var jsonMatch string

	// 1. 先尝试匹配 ```json ... ``` 代码块
	jsonBlockMatch := regexp.MustCompile("(?s)```json\\s*(\\[[\\s\\S]*?\\])\\s*```").FindStringSubmatch(resultText)
	if len(jsonBlockMatch) > 1 {
		jsonMatch = jsonBlockMatch[1]
		fmt.Println("[识别] 从json代码块中提取JSON")
	} else {
		// 2. 尝试匹配 ``` ... ``` 代码块中的数组
		jsonBlockMatch = regexp.MustCompile("(?s)```\\s*(\\[[\\s\\S]*?\\])\\s*```").FindStringSubmatch(resultText)
		if len(jsonBlockMatch) > 1 {
			jsonMatch = jsonBlockMatch[1]
			fmt.Println("[识别] 从代码块中提取JSON")
		} else {
			// 3. 查找最外层的JSON数组（匹配第一个 [ 到最后一个 ]）
			// 使用非贪婪匹配找到完整的JSON数组
			jsonMatch = regexp.MustCompile(`(\[[\s\S]*\])`).FindString(resultText)
			if jsonMatch == "" {
				return nil, fmt.Errorf("未找到JSON数据，原始响应: %s", resultText)
			}
			fmt.Println("[识别] 从文本中提取JSON")
		}
	}

	// 解析JSON
	var courses []models.Course
	err = json.Unmarshal([]byte(jsonMatch), &courses)
	if err != nil {
		return nil, fmt.Errorf("JSON解析失败: %v, 数据: %s", err, jsonMatch)
	}

	fmt.Printf("[识别] 成功解析 %d 门课程\n", len(courses))
	return courses, nil
}

// SaveToJSON 保存课程数据到JSON文件
func SaveToJSON(courses []models.Course, filename string) error {
	jsonData, err := json.MarshalIndent(courses, "", "  ")
	if err != nil {
		return err
	}

	err = os.WriteFile(filename, jsonData, 0644)
	if err != nil {
		return err
	}

	fmt.Printf("[输出] JSON结果已保存到: %s\n", filename)
	return nil
}

// PrintCourses 打印课程信息
func PrintCourses(courses []models.Course) {
	fmt.Println("\n==================================================")
	fmt.Println("识别结果：")
	fmt.Println("==================================================")
	for i, course := range courses {
		fmt.Printf("\n课程 %d: %s\n", i+1, course.Name)
		fmt.Printf("  教师: %s\n", course.Teacher)
		fmt.Printf("  教室: %s\n", course.Room)
		fmt.Printf("  时间:")
		for _, sch := range course.Schedule {
			fmt.Printf(" %s %v 周%v", sch.Day, sch.Section, sch.Weeks)
		}
		fmt.Println()
	}
	fmt.Println("\n==================================================")
}