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
}

// mcpTokenAudience 与 user 模块签发时写入的用途标记一致。
// 必须校验：将来若有别处也用 McpTokenKey 签东西，不校验就会被一并接受。
const mcpTokenAudience = "mcp"

// ResolveUID 决定这次工具调用该用哪个 uid。
//
// argUID 是模型传上来的参数，只用于「与会话 uid 比对」，绝不作为回退值。
func (this *Mcp) ResolveUID(ctx context.Context, argUID string) (string, error) {
	return this.ResolveUIDWithToken(ctx, argUID, "")
}

// ResolveUIDWithToken 在 [ResolveUID] 的基础上多认一条来源：**参数里的签名令牌**。
//
// ## 为什么需要参数这条路
//
// 走 Authorization 头是首选，但 EMAI 这条链路上根本送不到：客户端连百炼只带
// 百炼的 API Key，百炼调 MCP 的 Authorization 头又是控制台里配死的、所有用户
// 共用一个值。唯一能按用户变的通道是 Start 指令的 `biz_params.user_defined_params`，
// 它会作为**工具参数**交到这里。
//
// ⚠️ 「参数里的令牌」和「参数里的 uid」是两回事，别混为一谈：
// 裸 uid 谁都能填，所以永远不采信；令牌带签名，伪造不了，安全性与走 header 等价。
//
// ⚠️ 令牌用的是独立的 McpTokenKey，不是 gateway 的 TokenKey——它经过百炼
// （第三方）之手，万一泄露也只能用于 MCP 这几个只读工具。
//
// ⚠️ `auth_token` **刻意不声明进工具的参数 schema**：声明了大模型就会试着自己
// 填一个值，而它根本无从知道真值，只会编一个出来——既污染参数又让日志里全是
// 「令牌校验失败」。它由百炼从 user_defined_params 注入到调用参数里，
// 不需要模型参与，所以这里只管读、不管声明。
func (this *Mcp) ResolveUIDWithToken(ctx context.Context, argUID, argToken string) (string, error) {
	uid, err := this.resolveUIDFromSources(ctx, argToken)
	if err != nil {
		return "", err
	}
	argUID = strings.TrimSpace(argUID)
	if argUID != "" && argUID != uid {
		return "", fmt.Errorf("uid 与当前会话不符，已拒绝")
	}
	return uid, nil
}

// resolveUIDFromSources 依次尝试：请求头的登录 JWT → 参数里的 MCP 令牌。
// 两条都拿不出身份就拒，**绝不回退到参数 uid**——那等于这层没做。
func (this *Mcp) resolveUIDFromSources(ctx context.Context, argToken string) (string, error) {
	var headerErr error
	if key := this.options.TokenKey; key != "" {
		if raw := strings.TrimSpace(tokenFromContext(ctx)); raw != "" {
			uid, err := parseUID(raw, []byte(key))
			if err == nil {
				return uid, nil
			}
			headerErr = err
		}
	}

	if argToken = strings.TrimSpace(argToken); argToken != "" {
		key := this.options.McpTokenKey
		if key == "" {
			return "", fmt.Errorf("服务端未配置 McpTokenKey，无法校验令牌")
		}
		uid, err := parseUIDWithAudience(argToken, []byte(key), mcpTokenAudience)
		if err != nil {
			return "", fmt.Errorf("令牌校验失败: %v", err)
		}
		return uid, nil
	}

	if headerErr != nil {
		return "", fmt.Errorf("身份校验失败: %v", headerErr)
	}
	if this.options.TokenKey == "" && this.options.McpTokenKey == "" {
		// 两个密钥都没配说明部署时漏了注入。这时**拒绝所有用户类工具**
		// 而不是放行：配置缺失导致鉴权静默失效是最糟的一种失败方式。
		return "", fmt.Errorf("服务端未配置 TokenKey/McpTokenKey，用户数据类工具不可用")
	}
	return "", fmt.Errorf("身份校验失败: 未带 Authorization，也未带 MCP 令牌")
}

// parseUIDWithAudience 解 MCP 令牌：验签之外还要校验用途(aud)。
// jwt/v5 默认**不校验** aud，必须显式声明 WithAudience。
func parseUIDWithAudience(tokenString string, secretKey []byte, audience string) (string, error) {
	tokenString = strings.TrimSpace(tokenString)
	if tokenString == "" {
		return "", fmt.Errorf("令牌为空")
	}
	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
		},
		jwt.WithAudience(audience),
		// 过期是这套令牌的核心保障，必须要求 exp 存在——没有 exp 的令牌等同长期有效
		jwt.WithExpirationRequired(),
	)
	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
}