schedule_converter/grpc/handler.go

package grpc

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"sync"
	"time"

	pb "schedule_converter/proto/runner"

	"schedule_converter/auth"
	"schedule_converter/captcha"
	"schedule_converter/client"
	"schedule_converter/config"
	"schedule_converter/models"
	"schedule_converter/parser"
	apperr "schedule_converter/pkg/errors"
)

type TaskHandler struct {
	cancelledTasks sync.Map
}

func NewTaskHandler() *TaskHandler {
	return &TaskHandler{}
}

func (h *TaskHandler) Execute(task *pb.Task) *pb.TaskResult {
	log := slog.With("task_id", task.TaskId, "task_type", task.Type.String())
	log.Info("开始执行任务")

	result := &pb.TaskResult{
		TaskId:      task.TaskId,
		CompletedAt: time.Now().UnixMilli(),
	}

	if h.IsCancelled(task.TaskId) {
		result.Status = pb.TaskStatus_TASK_STATUS_CANCELLED
		result.ErrorMessage = "任务已取消"
		log.Info("任务已取消")
		return result
	}

	var err error
	var data []byte

	switch task.Type {
	case pb.TaskType_TASK_TYPE_GET_SCHEDULE:
		data, err = h.getSchedule(context.Background(), task.Payload)
	case pb.TaskType_TASK_TYPE_LOGIN:
		data, err = h.login(context.Background(), task.Payload)
	default:
		result.Status = pb.TaskStatus_TASK_STATUS_FAILED
		result.ErrorMessage = fmt.Sprintf("未知的任务类型: %v", task.Type)
		log.Warn("未知的任务类型")
		return result
	}

	if err != nil {
		result.Status = pb.TaskStatus_TASK_STATUS_FAILED
		result.ErrorMessage = err.Error()
		log.Error("任务执行失败", "error", err)
	} else {
		result.Status = pb.TaskStatus_TASK_STATUS_SUCCESS
		result.Data = data
		log.Info("任务执行成功")
	}

	return result
}

func (h *TaskHandler) CancelTask(taskID string) {
	h.cancelledTasks.Store(taskID, true)
	slog.Info("任务已标记为取消", "task_id", taskID)
}

func (h *TaskHandler) IsCancelled(taskID string) bool {
	v, ok := h.cancelledTasks.Load(taskID)
	return ok && v.(bool)
}

func (h *TaskHandler) getSchedule(ctx context.Context, payload []byte) ([]byte, error) {
	var req pb.GetSchedulePayload
	if err := json.Unmarshal(payload, &req); err != nil {
		return nil, apperr.New("parse_payload", err, "解析任务载荷失败")
	}

	log := slog.With("username", req.Username, "semester", req.Semester)
	log.Info("获取课表任务")

	httpClient, err := createHTTPClient()
	if err != nil {
		return nil, err
	}

	cfg := config.Get()
	setTWFIDCookie(httpClient, cfg.BaseURL, "")

	courses, err := executeLoginAndFetchSchedule(ctx, httpClient, req.Username, req.Password, req.Semester)
	if err != nil {
		return nil, err
	}

	result := &pb.ScheduleResultData{
		StudentId: req.Username,
		Semester:  req.Semester,
		Courses:   convertCourses(courses),
	}

	log.Info("课表获取成功", "course_count", len(courses))
	return json.Marshal(result)
}

func (h *TaskHandler) login(ctx context.Context, payload []byte) ([]byte, error) {
	var req pb.LoginPayload
	if err := json.Unmarshal(payload, &req); err != nil {
		return nil, apperr.New("parse_payload", err, "解析任务载荷失败")
	}

	log := slog.With("username", req.Username)
	log.Info("登录任务")

	httpClient, err := createHTTPClient()
	if err != nil {
		return nil, err
	}

	cfg := config.Get()
	setTWFIDCookie(httpClient, cfg.BaseURL, req.Twfid)

	if err := executeLogin(ctx, httpClient, req.Username, req.Password); err != nil {
		return nil, err
	}

	log.Info("登录成功")
	return json.Marshal(map[string]any{
		"success": true,
		"message": "登录成功",
	})
}

func createHTTPClient() (*http.Client, error) {
	jar, err := cookiejar.New(nil)
	if err != nil {
		return nil, apperr.New("create_jar", err, "创建 cookie jar 失败")
	}
	return &http.Client{Jar: jar}, nil
}

func setTWFIDCookie(httpClient *http.Client, baseURL, twfid string) {
	cfg := config.Get()
	if twfid == "" {
		twfid = cfg.TWFID
	}
	u, _ := url.Parse(baseURL)
	httpClient.Jar.SetCookies(u, []*http.Cookie{
		{Name: "TWFID", Value: twfid},
	})
}

func executeLogin(ctx context.Context, httpClient *http.Client, username, password string) error {
	log := slog.With("username", username)

	if err := client.GetInitialCookiesWithClient(httpClient); err != nil {
		return apperr.New("get_cookies", err, "获取初始 Cookie 失败")
	}

	cfg := config.Get()
	_ = cfg.Recognizer.Timeout

	for attempt := 1; attempt <= 3; attempt++ {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		log.Info("尝试登录", "attempt", attempt, "max_retries", 3)

		captchaCode, err := captcha.GetAndRecognize(httpClient)
		if err != nil {
			log.Warn("验证码获取失败", "error", err, "attempt", attempt)
			time.Sleep(3 * time.Second)
			continue
		}

		success, err := auth.LoginWithClient(httpClient, captchaCode, username, password)
		if err != nil {
			log.Warn("登录请求失败", "error", err, "attempt", attempt)
			time.Sleep(3 * time.Second)
			continue
		}

		if success {
			return nil
		}

		log.Warn("登录验证失败", "attempt", attempt)
		time.Sleep(3 * time.Second)
	}

	return apperr.ErrLoginFailed
}

func executeLoginAndFetchSchedule(ctx context.Context, httpClient *http.Client, username, password, semester string) ([]models.Course, error) {
	if err := executeLogin(ctx, httpClient, username, password); err != nil {
		return nil, err
	}

	cfg := config.Get()
	if semester == "" {
		semester = cfg.Semester
	}

	return parser.FetchScheduleWithClient(httpClient, semester)
}

func convertCourses(courses []models.Course) []*pb.Course {
	result := make([]*pb.Course, len(courses))
	for i, c := range courses {
		schedule := make([]*pb.ScheduleTime, len(c.Schedule))
		for j, s := range c.Schedule {
			schedule[j] = &pb.ScheduleTime{
				Day:      s.Day,
				Sections: convertToInt32(s.Section),
				Weeks:    convertToInt32(s.Weeks),
			}
		}
		result[i] = &pb.Course{
			Name:     c.Name,
			Teacher:  c.Teacher,
			Room:     c.Room,
			Schedule: schedule,
		}
	}
	return result
}

func convertToInt32(arr []int) []int32 {
	result := make([]int32, len(arr))
	for i, v := range arr {
		result[i] = int32(v)
	}
	return result
}