schedule_converter/main.go

package main

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

	"schedule_converter/client"
	"schedule_converter/config"
	grpcClient "schedule_converter/grpc"
	"schedule_converter/pkg/logger"
	"schedule_converter/service"

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

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

	if *debug {
		logger.SetLevel(slog.LevelDebug)
		logger.SetTextHandler()
	}

	cfg := config.Load()
	slog.Info("配置加载完成", "mode", *mode, "base_url", cfg.BaseURL)

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

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"
	}

	slog.Info("启动 gRPC 客户端", "runner_id", runnerID, "version", version)

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

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

	go func() {
		sig := <-sigChan
		slog.Info("收到信号，正在关闭...", "signal", sig)
		cancel()
	}()

	reconnectInterval := 5 * time.Second

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

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

		if err := c.Connect(ctx); err != nil {
			slog.Warn("连接失败，稍后重试", "error", err, "retry_after", reconnectInterval)
			time.Sleep(reconnectInterval)
			continue
		}

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

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

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

		if err != nil {
			slog.Warn("连接断开，稍后重连", "error", err, "retry_after", reconnectInterval)
		} else {
			slog.Info("连接已关闭，稍后重连", "retry_after", reconnectInterval)
		}
		time.Sleep(reconnectInterval)
	}
}

func runHTTPServer() {
	client.InitClient()

	r := gin.New()
	r.Use(gin.Recovery())
	r.Use(requestLogger())

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{"message": "pong"})
	})

	r.POST("/recognize", handleRecognize)
	r.POST("/task", handleTask)

	slog.Info("HTTP 服务器启动", "port", 8080)
	if err := r.Run(":8080"); err != nil {
		slog.Error("服务器启动失败", "error", err)
		os.Exit(1)
	}
}

func requestLogger() gin.HandlerFunc {
	return func(c *gin.Context) {
		start := time.Now()
		path := c.Request.URL.Path
		query := c.Request.URL.RawQuery

		c.Next()

		latency := time.Since(start)
		status := c.Writer.Status()

		if query != "" {
			path = path + "?" + query
		}

		slog.Info("HTTP请求",
			"method", c.Request.Method,
			"path", path,
			"status", status,
			"latency", latency.String(),
			"client_ip", c.ClientIP(),
		)
	}
}

func handleRecognize(c *gin.Context) {
	file, header, err := c.Request.FormFile("image")
	if err != nil {
		slog.Warn("图片上传失败", "error", err)
		c.JSON(http.StatusBadRequest, gin.H{"error": "请上传图片文件", "details": err.Error()})
		return
	}
	defer file.Close()

	slog.Info("收到图片上传", "filename", header.Filename, "size", 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 {
		slog.Error("读取图片失败", "error", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "读取图片失败", "details": err.Error()})
		return
	}

	recognizer := service.NewRecognizerService()
	courses, err := recognizer.RecognizeFromBytes(c.Request.Context(), imageData)
	if err != nil {
		slog.Error("识别失败", "error", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "识别失败", "details": err.Error()})
		return
	}

	service.PrintCourses(courses)

	if err := service.SaveToJSON(courses, "schedule_result.json"); err != nil {
		slog.Warn("保存JSON失败", "error", err)
	}

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

func handleTask(c *gin.Context) {
	cfg := config.Get()
	if cfg.Username == "" || cfg.Password == "" {
		slog.Warn("环境变量未设置")
		c.JSON(http.StatusInternalServerError, gin.H{"error": "环境变量未设置"})
		return
	}

	scheduleSvc := service.NewScheduleService()
	_, err := scheduleSvc.FetchSchedule(c.Request.Context(), cfg.Username, cfg.Password, cfg.Semester)
	if err != nil {
		slog.Error("任务执行失败", "error", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

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