schedule_converter/main.go

package main

import (
	"context"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"jwts/auth"
	"jwts/captcha"
	"jwts/client"
	"jwts/config"
	grpcClient "jwts/grpc"
	"jwts/models"
	"jwts/parser"

	"github.com/gin-gonic/gin"
)

func main() {
	mode := flag.String("mode", "http", "运行模式: http 或 grpc")
	grpcAddr := flag.String("grpc-addr", "localhost:50051", "gRPC 服务器地址")
	flag.Parse()

	if *mode == "grpc" {
		runGRPCClient(*grpcAddr)
	} else {
		runHTTPServer()
	}
}

// runGRPCClient 运行 gRPC 客户端模式
func runGRPCClient(serverAddr string) {
	runnerID := os.Getenv("RUNNER_ID")
	if runnerID == "" {
		hostname, _ := os.Hostname()
		runnerID = "runner-" + hostname
	}

	version := os.Getenv("RUNNER_VERSION")
	if version == "" {
		version = "1.0.0"
	}

	log.Printf("启动 gRPC 客户端, Runner ID: %s, 版本: %s", runnerID, version)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-sigChan
		log.Println("正在关闭...")
		cancel()
	}()

	reconnectInterval := 5 * time.Second

	for {
		select {
		case <-ctx.Done():
			log.Println("客户端已停止")
			return
		default:
		}

		client := grpcClient.NewClient(serverAddr, runnerID, version)
		handler := grpcClient.NewTaskHandler()
		client.SetHandler(handler)

		if err := client.Connect(ctx); err != nil {
			log.Printf("连接失败: %v, %v 后重试...", err, reconnectInterval)
			time.Sleep(reconnectInterval)
			continue
		}

		heartbeatCtx, heartbeatCancel := context.WithCancel(ctx)
		go client.HeartbeatLoop(heartbeatCtx, 30*time.Second)

		err := client.Run(ctx)
		heartbeatCancel()
		client.Close()

		if ctx.Err() != nil {
			return
		}

		if err != nil {
			log.Printf("连接断开: %v, %v 后重连...", err, reconnectInterval)
		} else {
			log.Printf("连接已关闭, %v 后重连...", reconnectInterval)
		}
		time.Sleep(reconnectInterval)
	}
}

// runHTTPServer 运行 HTTP 服务器模式
func runHTTPServer() {
	// 初始化客户端
	client.InitClient()

	r := gin.Default()

	// 健康检查接口
	r.GET("/ping", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"message": "pong",
		})
	})

	// 课表图片识别接口
	r.POST("/recognize", func(c *gin.Context) {
		file, header, err := c.Request.FormFile("image")
		if err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "请上传图片文件", "details": err.Error()})
			return
		}
		defer file.Close()

		log.Printf("收到图片上传: %s, 大小: %d bytes", header.Filename, header.Size)

		if header.Size > 10*1024*1024 {
			c.JSON(http.StatusBadRequest, gin.H{"error": "图片文件过大", "max_size": "10MB"})
			return
		}

		imageData, err := io.ReadAll(file)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "读取图片失败", "details": err.Error()})
			return
		}

		courses, err := captcha.RecognizeScheduleFromBytes(imageData)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "识别失败", "details": err.Error()})
			return
		}

		models.PrintCourses(courses)

		err = models.SaveToJSON(courses, "schedule_result.json")
		if err != nil {
			log.Printf("保存JSON失败: %v", err)
		}

		c.JSON(http.StatusOK, gin.H{
			"message": "识别成功",
			"courses": courses,
			"count":   len(courses),
		})
	})

	// 任务处理接口
	r.POST("/task", func(c *gin.Context) {
		// 这里可以根据请求体参数决定执行什么任务
		// 目前默认执行原有的登录和获取课表逻辑
		if config.Username == "" || config.Password == "" {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "环境变量未设置"})
			return
		}

		// 执行任务
		err := runTask()
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}

		c.JSON(http.StatusOK, gin.H{"message": "任务执行成功"})
	})

	// 启动服务
	log.Println("Server starting on :8080")
	if err := r.Run(":8080"); err != nil {
		log.Fatalf("Failed to run server: %v", err)
	}
}

func runTask() error {
	// 步骤1: 获取初始Cookie
	err := client.GetInitialCookies()
	if err != nil {
		log.Printf("获取Cookie失败: %v", err)
		return err
	}

	// 重试登录
	maxRetries := 3
	for attempt := 1; attempt <= maxRetries; attempt++ {
		log.Printf("尝试登录 (第%d次)...", attempt)

		// 获取并识别验证码
		captchaCode, err := captcha.GetAndRecognize(client.Client)
		if err != nil {
			log.Printf("    验证码获取失败: %v", err)
			continue
		}

		// 登录
		success, err := auth.Login(captchaCode)
		if err != nil {
			log.Printf("    登录失败: %v", err)
			continue
		}

		if success {
			// 获取课表
			_, err = parser.FetchSchedule(client.Client)
			if err != nil {
				log.Printf("获取课表失败: %v", err)
			}
			return nil
		}

		log.Println("    登录失败，3秒后重试...")
	}

	return nil
}