package mcp

import (
	"context"
	"fmt"
	"strings"

	"github.com/golang-jwt/jwt/v5"
)

/*
MCP 的会话身份。

## 这块此前是空的

`authFromRequest` 只把 Authorization 头塞进 context，**全仓没有一处读回来，
也没有任何校验**；三个用户工具（get_user_tasks / allhelp_task / cancel_user_task）
的 uid 都是 `mcp.WithString("uid", Required())` —— 由大模型填进来的普通参数。

而 MCP 是**独立进程、独立端口（7300）**，不经过 gateway 那套 parseToken +
isInWhiteList（那只管 /api/* 和 /web/*），又必须对百炼公网可达。
也就是说这不是「内网接口没做鉴权」，是**公网接口完全没有鉴权**：
知道 uid 就能读改任意人的数据。

记忆中心的条目装的是花销金额、灵感原文、每日行程，比「提醒任务」敏感得多，
所以在挂 memory 工具之前先把这层补上。

## 口径

  - 有 Authorization 且能解出 uid → 用**会话 uid**，忽略参数里的 uid；
  - 参数里的 uid 与会话 uid 不一致 → 直接拒（不是静默改用会话 uid：
    调用方显然误以为自己能指定用户，让它失败比让它以为成功了更安全）；
  - 解不出会话 uid → 拒。**不降级回退到参数 uid** —— 那等于这层没做。

⚠️ TokenKey 必须与 gateway 的 `GATEWAY_TOKEN_KEY` 一致，否则所有 token 都验不过。
两边都从同一个 .env 变量注入。
*/

// authKey is a custom context key for storing the auth token.
type authKey struct{}

// withAuthKey adds an auth key to the context.
func withAuthKey(ctx context.Context, auth string) context.Context {
	return context.WithValue(ctx, authKey{}, auth)
}

// tokenFromContext 取出请求带来的 Authorization 原始串
func tokenFromContext(ctx context.Context) string {
	v, _ := ctx.Value(authKey{}).(string)
	return v
}

// parseUID 从 JWT 里解出 uid。与 gateway/core.go 的 parseToken 同一套口径
// （jwt.RegisteredClaims，uid 放在 ID 字段）。
func parseUID(tokenString string, secretKey []byte) (string, error) {
	tokenString = strings.TrimSpace(tokenString)
	if tokenString == "" {
		return "", fmt.Errorf("缺少 Authorization")
	}
	// 兼容 "Bearer xxx" 写法：gateway 那边收到的就是裸 token，
	// 但 MCP 客户端（百炼）常按标准加 Bearer 前缀。
	if len(tokenString) > 7 && strings.EqualFold(tokenString[:7], "bearer ") {
		tokenString = strings.TrimSpace(tokenString[7:])
	}
	parsed, err := jwt.ParseWithClaims(tokenString, &jwt.RegisteredClaims{},
		func(token *jwt.Token) (interface{}, error) {
			if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
			}
			return secretKey, nil
		})
	if err != nil {
		return "", err
	}
	claims, ok := parsed.Claims.(*jwt.RegisteredClaims)
	if !ok || !parsed.Valid {
		return "", fmt.Errorf("invalid token")
	}
	if claims.ID == "" {
		return "", fmt.Errorf("token 里没有 uid")
	}
	return claims.ID, nil
}

// ResolveUID 决定这次工具调用该用哪个 uid。
//
// argUID 是模型传上来的参数，只用于「与会话 uid 比对」，绝不作为回退值。
func (this *Mcp) ResolveUID(ctx context.Context, argUID string) (string, error) {
	key := this.options.TokenKey
	if key == "" {
		// 没配 TokenKey 说明部署时漏了 GATEWAY_TOKEN_KEY。
		// 这时**拒绝所有用户类工具**而不是放行：配置缺失导致鉴权静默失效
		// 是最糟的一种失败方式。
		return "", fmt.Errorf("服务端未配置 TokenKey，用户数据类工具不可用")
	}
	uid, err := parseUID(tokenFromContext(ctx), []byte(key))
	if err != nil {
		return "", fmt.Errorf("身份校验失败: %v", err)
	}
	argUID = strings.TrimSpace(argUID)
	if argUID != "" && argUID != uid {
		return "", fmt.Errorf("uid 与当前会话不符，已拒绝")
	}
	return uid, nil
}