package support

import (
	"ai_scheduler/internal/config"
	"ai_scheduler/internal/data/constants"
	errorcode "ai_scheduler/internal/data/error"
	"ai_scheduler/internal/pkg/util"
	"ai_scheduler/internal/pkg/utils_vllm"
	"context"
	"encoding/json"
	"strings"
	"time"

	"github.com/cloudwego/eino/schema"
	"github.com/gofiber/fiber/v2"
)

// HytAddressIngester 货易通地址提取实现
type HytAddressIngester struct {
	cfg *config.Config
}

func NewHytAddressIngester(cfg *config.Config) *HytAddressIngester {
	return &HytAddressIngester{cfg: cfg}
}

// Auth 鉴权逻辑
func (s *HytAddressIngester) Auth(c *fiber.Ctx) error {
	// 读取头
	token := strings.TrimSpace(c.Get("X-Source-Key"))
	ts := strings.TrimSpace(c.Get("X-Timestamp"))

	// 时间窗口校验
	if ts != "" && !util.IsInTimeWindow(ts, 5*time.Minute) {
		return errorcode.AuthNotFound
	}
	// token校验
	if token == "" || token != constants.TokenAddressIngestHyt {
		return errorcode.KeyNotFound
	}
	return nil
}

// Ingest 执行提取逻辑
func (s *HytAddressIngester) Ingest(ctx context.Context, text string) (*AddressIngestResp, error) {
	// 模型调用
	client, cleanup, err := utils_vllm.NewClient(s.cfg)
	if err != nil {
		return nil, err
	}
	defer cleanup()

	res, err := client.Chat(ctx, []*schema.Message{
		{
			Role:    "system",
			Content: constants.SystemPromptAddressIngestHyt,
		},
		{
			Role:    "user",
			Content: text,
		},
	})
	if err != nil {
		return nil, err
	}

	// 解析模型返回结果
	var addr AddressIngestResp
	// 尝试直接解析
	if err := json.Unmarshal([]byte(res.Content), &addr); err != nil {
		// 修复json字符串
		repairedContent, err := util.JSONRepair(res.Content)
		if err != nil {
			return nil, errorcode.ParamErrf("invalid response body: %v", err)
		}
		if err := json.Unmarshal([]byte(repairedContent), &addr); err != nil {
			return nil, errorcode.ParamErrf("invalid response body: %v", err)
		}
	}
	return &addr, nil
}