package tokenprobe

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"crypto/tls"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"os"
	"runtime"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/google/uuid"
	"golang.org/x/net/http2"
)

const (
	cursorBackendURL    = "https://api2.cursor.sh"
	cursorAgentPath     = "/aiserver.v1.ChatService/StreamUnifiedChatWithTools"
	cursorClientVersion = "2.6.22"
	cursorHiMaxRead     = 512 * 1024
	// probeHiText 发往官方 Agent 的探测内容（与前端展示 probeMessage 一致）
	probeHiText = "hi"
)

var cursorProbeHTTPClient = newCursorHTTP2Client()

func newCursorHTTP2Client() *http.Client {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{MinVersion: tls.VersionTLS12},
	}
	// 与 Cursor 官方一致走 HTTP/2
	if err := http2.ConfigureTransport(tr); err != nil {
		return &http.Client{Timeout: 40 * time.Second}
	}
	return &http.Client{Transport: tr, Timeout: 40 * time.Second}
}

func cursorClientOS() string {
	switch runtime.GOOS {
	case "windows":
		return "win32"
	case "darwin":
		return "darwin"
	default:
		return "linux"
	}
}

func cursorClientArch() string {
	switch runtime.GOARCH {
	case "amd64":
		return "x64"
	case "arm64":
		return "arm64"
	default:
		return runtime.GOARCH
	}
}

func cursorEnvVersion() string {
	if v := strings.TrimSpace(os.Getenv("CURSOR_CLIENT_VERSION")); v != "" {
		return v
	}
	return cursorClientVersion
}

// --- protobuf wire (与 cursor_api_demo 对齐) ---

func pbVarint(v uint64) []byte {
	var out []byte
	for v >= 0x80 {
		out = append(out, byte(v&0x7f|0x80))
		v >>= 7
	}
	out = append(out, byte(v&0x7f))
	return out
}

func pbField(fieldNum int, wireType int, value interface{}) []byte {
	tag := uint64(fieldNum<<3 | wireType)
	out := pbVarint(tag)
	switch wireType {
	case 0:
		var n uint64
		switch x := value.(type) {
		case int:
			n = uint64(x)
		case int32:
			n = uint64(x)
		case uint32:
			n = uint64(x)
		case uint64:
			n = x
		default:
			n = uint64(0)
		}
		out = append(out, pbVarint(n)...)
	case 2:
		var b []byte
		switch x := value.(type) {
		case string:
			b = []byte(x)
		case []byte:
			b = x
		default:
			b = []byte(fmt.Sprint(x))
		}
		out = append(out, pbVarint(uint64(len(b)))...)
		out = append(out, b...)
	}
	return out
}

func encodeCursorMessage(content string, role int, messageID string, chatModeEnum *int) []byte {
	msg := pbField(1, 2, content)
	msg = append(msg, pbField(2, 0, role)...)
	msg = append(msg, pbField(13, 2, messageID)...)
	if chatModeEnum != nil {
		msg = append(msg, pbField(47, 0, *chatModeEnum)...)
	}
	return msg
}

func encodeCursorModel(modelName string) []byte {
	msg := pbField(1, 2, modelName)
	msg = append(msg, pbField(4, 2, []byte{})...)
	return msg
}

func encodeCursorSetting() []byte {
	inner := pbField(1, 2, []byte{})
	inner = append(inner, pbField(2, 2, []byte{})...)
	msg := pbField(1, 2, `cursor\aisettings`)
	msg = append(msg, pbField(3, 2, []byte{})...)
	msg = append(msg, pbField(6, 2, inner)...)
	msg = append(msg, pbField(8, 0, 1)...)
	msg = append(msg, pbField(9, 0, 1)...)
	return msg
}

func encodeCursorMetadata() []byte {
	msg := pbField(1, 2, cursorClientOS())
	msg = append(msg, pbField(2, 2, cursorClientArch())...)
	msg = append(msg, pbField(3, 2, "unknown")...)
	msg = append(msg, pbField(4, 2, "go-platform/tokenprobe")...)
	msg = append(msg, pbField(5, 2, time.Now().Format(time.RFC3339))...)
	return msg
}

func encodeCursorMessageID(messageID string, role int) []byte {
	msg := pbField(1, 2, messageID)
	msg = append(msg, pbField(3, 0, role)...)
	return msg
}

// defaultAgentTools 与 cursor_agent_client.DEFAULT_TOOLS 一致
var defaultAgentTools = []int{5, 6, 3, 15, 7, 8, 42}

func encodeCursorAgentRequest(userContent, modelName string) []byte {
	msgID := uuid.NewString()
	cm := 2 // Agent
	userMsg := encodeCursorMessage(userContent, 1, msgID, &cm)

	var msg []byte
	msg = append(msg, pbField(1, 2, userMsg)...)
	msg = append(msg, pbField(2, 0, 1)...)
	msg = append(msg, pbField(3, 2, []byte{})...)
	msg = append(msg, pbField(4, 0, 1)...)
	msg = append(msg, pbField(5, 2, encodeCursorModel(modelName))...)
	msg = append(msg, pbField(8, 2, "")...)
	msg = append(msg, pbField(13, 0, 1)...)
	msg = append(msg, pbField(15, 2, encodeCursorSetting())...)
	msg = append(msg, pbField(19, 0, 1)...)
	msg = append(msg, pbField(23, 2, uuid.NewString())...)
	msg = append(msg, pbField(26, 2, encodeCursorMetadata())...)
	msg = append(msg, pbField(27, 0, 1)...)
	for _, t := range defaultAgentTools {
		msg = append(msg, pbField(29, 0, t)...)
	}
	msg = append(msg, pbField(30, 2, encodeCursorMessageID(msgID, 1))...)
	msg = append(msg, pbField(35, 0, 0)...)
	msg = append(msg, pbField(38, 0, 0)...)
	msg = append(msg, pbField(46, 0, 2)...)
	msg = append(msg, pbField(47, 2, "")...)
	msg = append(msg, pbField(48, 0, 0)...)
	msg = append(msg, pbField(49, 0, 0)...)
	msg = append(msg, pbField(51, 0, 0)...)
	msg = append(msg, pbField(53, 0, 1)...)
	msg = append(msg, pbField(54, 2, "agent")...)
	return msg
}

func encodeStreamUnifiedChatWithToolsRequest(inner []byte) []byte {
	return pbField(1, 2, inner)
}

func generateCursorAgentFramedBody(userText, model string) []byte {
	inner := encodeCursorAgentRequest(userText, model)
	buf := encodeStreamUnifiedChatWithToolsRequest(inner)
	magic := byte(0x00)
	hexLen := fmt.Sprintf("%08x", len(buf))
	lenBytes, err := hex.DecodeString(hexLen)
	if err != nil || len(lenBytes) != 4 {
		lenB := []byte{byte(len(buf) >> 24), byte(len(buf) >> 16), byte(len(buf) >> 8), byte(len(buf))}
		return append([]byte{magic}, append(lenB, buf...)...)
	}
	return append([]byte{magic}, append(lenBytes, buf...)...)
}

func hashed64Hex(input, salt string) string {
	h := sha256.Sum256([]byte(input + salt))
	return hex.EncodeToString(h[:])
}

func generateCursorChecksum(authToken string) string {
	machineID := hashed64Hex(authToken, "machineId")
	ts := int(time.Now().UnixMilli() / 1_000_000)
	barr := []byte{
		byte(ts >> 40), byte(ts >> 32), byte(ts >> 24), byte(ts >> 16), byte(ts >> 8), byte(ts),
	}
	t := byte(165)
	for i := range barr {
		barr[i] = ((barr[i] ^ t) + byte(i%256)) & 255
		t = barr[i]
	}
	const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
	var enc strings.Builder
	for i := 0; i < len(barr); i += 3 {
		a := barr[i]
		var b, c byte
		if i+1 < len(barr) {
			b = barr[i+1]
		}
		if i+2 < len(barr) {
			c = barr[i+2]
		}
		enc.WriteByte(alphabet[a>>2])
		enc.WriteByte(alphabet[((a&3)<<4)|(b>>4)])
		if i+1 < len(barr) {
			enc.WriteByte(alphabet[((b&15)<<2)|(c>>6)])
		}
		if i+2 < len(barr) {
			enc.WriteByte(alphabet[c&63])
		}
	}
	return enc.String() + machineID
}

func asciiLowerInPlace(b []byte) {
	for i := range b {
		c := b[i]
		if c >= 'A' && c <= 'Z' {
			b[i] = c + ('a' - 'A')
		}
	}
}

// 社区脚本中的「额度用尽」ASCII 前缀（与明文一致，便于在二进制流中 bytes.Contains，无需整句）
// 对应明文前缀：Get Cursor Pro for more Agent usage
var cursorQuotaExhaustedSigCommunity = []byte{
	0x47, 0x65, 0x74, 0x20, 0x43, 0x75, 0x72, 0x73,
	0x6f, 0x72, 0x20, 0x50, 0x72, 0x6f, 0x20, 0x66,
	0x6f, 0x72, 0x20, 0x6d, 0x6f, 0x72, 0x65, 0x20,
	0x41, 0x67, 0x65, 0x6e, 0x74, 0x20, 0x75, 0x73,
	0x61, 0x67, 0x65,
}

// cursorQuotaTipSig 与常见示例一致：raw 全字节里 bytes.Contains(raw, tipSig) → 额度用尽
var cursorQuotaTipSig = []byte("Get Cursor Pro for more Agent usage, unlimited Tab, and more.")

const cursorLimitTipPrefix = "Get Cursor Pro for more Agent usage, unlimited Tab"

// classifyCursorRawStream 在官方流式二进制/文本中匹配用量与升级提示（ASCII 区不区分大小写 + UTF-8 短语）
func classifyCursorRawStream(raw []byte) (blocked bool, reason string) {
	if len(raw) == 0 {
		return false, ""
	}
	for _, sig := range cursorQuotaExhaustedSigsFromEnv() {
		if bytes.Contains(raw, sig) {
			return true, fmt.Sprintf("流中匹配：CURSOR_QUOTA_EXHAUSTED_SIG_HEX 配置的二进制特征（%d 字节）", len(sig))
		}
	}
	if bytes.Contains(raw, cursorQuotaTipSig) {
		return true, "流中匹配：" + string(cursorQuotaTipSig)
	}
	// 社区脚本：仅到「…Agent usage」的 ASCII 前缀（流里可能只有前半段）
	if bytes.Contains(raw, cursorQuotaExhaustedSigCommunity) {
		return true, "流中匹配：Get Cursor Pro for more Agent usage…（社区 QuotaExhaustedSignature 前缀）"
	}
	if bytes.Contains(raw, []byte(cursorLimitTipPrefix)) {
		return true, "流中匹配：" + cursorLimitTipPrefix + "…"
	}
	low := append([]byte(nil), raw...)
	asciiLowerInPlace(low)
	if bytes.Contains(low, []byte("you've hit your usage limit")) ||
		bytes.Contains(low, []byte("youve hit your usage limit")) ||
		bytes.Contains(low, []byte("hit your usage limit")) {
		return true, "流中匹配：hit your usage limit / you've hit your usage limit"
	}
	if bytes.Contains(low, []byte("get cursor pro for more agent usage")) {
		return true, "流中匹配：get cursor pro for more agent usage"
	}
	if bytes.Contains(low, []byte("upgrade to pro")) {
		return true, "流中匹配：upgrade to pro"
	}
	if bytes.Contains(low, []byte("get cursor pro")) && bytes.Contains(low, []byte("agent")) {
		return true, "流中匹配：get cursor pro + agent"
	}
	if bytes.Contains(low, []byte("usage limit")) {
		return true, "流中匹配：usage limit"
	}
	if bytes.Contains(low, []byte("unlimited tab")) && bytes.Contains(low, []byte("cursor pro")) {
		return true, "流中匹配：unlimited tab + cursor pro"
	}

	flat := strings.ToLower(strings.ToValidUTF8(string(raw), "\uFFFD"))
	flat = strings.ReplaceAll(flat, "\u2019", "'") // 右单引号
	flat = strings.ReplaceAll(flat, "`", "'")
	if strings.Contains(flat, "you've hit your usage limit") {
		return true, "流中匹配：you've hit your usage limit（UTF-8）"
	}
	return false, ""
}

func truncateUTF8Preview(raw []byte, maxBytes int) string {
	s := strings.ToValidUTF8(string(raw), "\uFFFD")
	if maxBytes <= 0 || len(s) <= maxBytes {
		return s
	}
	// 按字节截断并保证合法 UTF-8
	s = s[:maxBytes]
	for len(s) > 0 && !utf8.ValidString(s) {
		s = s[:len(s)-1]
	}
	return s + "…(已截断)"
}

func prefixHexBody(b []byte, max int) string {
	if len(b) > max {
		b = b[:max]
	}
	return hex.EncodeToString(b)
}

func looksLikeGzip(raw []byte) bool {
	return len(raw) >= 3 && raw[0] == 0x1f && raw[1] == 0x8b && raw[2] == 0x08
}

func gunzipBytes(raw []byte) ([]byte, error) {
	zr, err := gzip.NewReader(bytes.NewReader(raw))
	if err != nil {
		return nil, err
	}
	defer zr.Close()
	return io.ReadAll(io.LimitReader(zr, cursorHiMaxRead))
}

func decodeConnectFramedBody(raw []byte) ([]byte, string, bool) {
	if len(raw) < 5 {
		return nil, "", false
	}

	var out bytes.Buffer
	offset := 0
	frameCount := 0
	compressedFrames := 0

	for offset+5 <= len(raw) {
		flags := raw[offset]
		n := int(raw[offset+1])<<24 | int(raw[offset+2])<<16 | int(raw[offset+3])<<8 | int(raw[offset+4])
		offset += 5
		if n < 0 || offset+n > len(raw) {
			return nil, "", false
		}
		payload := raw[offset : offset+n]
		offset += n
		frameCount++

		isCompressed := flags&0x01 == 0x01
		if isCompressed || looksLikeGzip(payload) {
			decoded, err := gunzipBytes(payload)
			if err != nil {
				out.Write(payload)
			} else {
				out.Write(decoded)
				compressedFrames++
			}
		} else {
			out.Write(payload)
		}
	}

	if frameCount == 0 || offset != len(raw) {
		return nil, "", false
	}

	note := fmt.Sprintf("响应体已按 Connect 分帧解析（%d 帧", frameCount)
	if compressedFrames > 0 {
		note += fmt.Sprintf("，其中 %d 帧已做 gzip 解压", compressedFrames)
	}
	note += "）后分析"
	return out.Bytes(), note, true
}

func decodeCursorResponseBody(raw []byte, contentEncoding string) ([]byte, string) {
	if decoded, note, ok := decodeConnectFramedBody(raw); ok {
		return decoded, note
	}

	enc := strings.ToLower(strings.TrimSpace(contentEncoding))
	if strings.Contains(enc, "gzip") || looksLikeGzip(raw) {
		decoded, err := gunzipBytes(raw)
		if err != nil {
			if strings.Contains(enc, "gzip") {
				return raw, "响应头声明 gzip，但解压失败，已回退为原始字节预览"
			}
			return raw, "检测到 gzip 魔数，但解压失败，已回退为原始字节预览"
		}
		if strings.Contains(enc, "gzip") {
			return decoded, "响应体已按 gzip 解压后分析"
		}
		return decoded, "响应体虽未显式声明 Content-Encoding，但按 gzip 魔数解压后分析"
	}
	if enc != "" {
		return raw, "响应头 Content-Encoding=" + enc + "，当前未额外解码，按原始字节分析"
	}
	return raw, "响应体未压缩或未声明压缩，且未识别为 Connect 分帧，按原始字节分析"
}

// cursorStreamProtocol 与官方客户端一致：Connect-RPC + protobuf 体，HTTP/2 流式
const cursorStreamProtocol = "Connect-Protocol-Version:1 + application/connect+proto，HTTP/2 二进制流（gRPC 兼容形态，非 JSON REST）"

// cursorStreamNote 说明 rawPreview / ok 的含义边界（与「仅通 200」结论一致）
const cursorStreamNote = `【协议】本 URL 为 Cursor 官方 Agent 流式接口，请求体为 protobuf（requestBodyPrefixHex 可见非表单/JSON）。` +
	`【响应】正文为分包二进制流，rawPreview 是按 UTF-8 有损解码的片段，绝大多数情况下会像乱码，属正常现象，不能当普通 UTF-8 接口正文解析。` +
	`【HTTP 200】仅表示 TLS/代理/网络到 api2.cursor.sh 通畅，不代表 protobuf 业务层、鉴权、设备指纹、风控配额或「能持续对话」已全部通过。` +
	`【ok 字段】当前仅在解码片段上做英文关键词启发式匹配；未命中不代表账户可用，命中也不覆盖「须在 IDE 内完整走流式协议」的场景。` +
	`【若要等价客户端】需完整实现 Connect 帧解析、会话与校验头、可能的 gzip/分包及双向流，本探测只做粗连通与可观测性辅助。` +
	`【二进制特征】若提示语被包在 protobuf 字段内、ASCII 子串匹配不到，可在运行环境设置 CURSOR_QUOTA_EXHAUSTED_SIG_HEX=hex1,hex2（逗号分隔十六进制，可选 0x 前缀），在原始响应字节上做 bytes.Contains，无需解析整条 proto；特征需自行对比「额度正常」与「用尽」两次抓包提取。`

// cursorQuotaExhaustedSigsFromEnv 从环境变量解析额度用尽时的二进制特征（不转 UTF-8）
func cursorQuotaExhaustedSigsFromEnv() [][]byte {
	s := strings.TrimSpace(os.Getenv("CURSOR_QUOTA_EXHAUSTED_SIG_HEX"))
	if s == "" {
		return nil
	}
	var out [][]byte
	for _, part := range strings.Split(s, ",") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		part = strings.TrimPrefix(strings.TrimPrefix(part, "0x"), "0X")
		b, err := hex.DecodeString(part)
		if err != nil || len(b) == 0 {
			continue
		}
		out = append(out, b)
	}
	return out
}

func cursorProbeResult(ok bool, detail string, httpStatus int, reqBody, raw, preview []byte) Result {
	if preview == nil {
		preview = raw
	}
	return Result{
		OK:                   ok,
		Detail:               detail,
		HTTPStatus:           httpStatus,
		ProbeMessage:         probeHiText,
		Endpoint:             cursorBackendURL + cursorAgentPath,
		BytesRead:            len(raw),
		RawPreview:           truncateUTF8Preview(preview, 24000),
		RequestBodyPrefixHex: prefixHexBody(reqBody, 128),
		StreamProtocol:       cursorStreamProtocol,
		StreamNote:           cursorStreamNote,
	}
}

func probeCursorHiAgent(authToken string) Result {
	if strings.Contains(authToken, "::") {
		if i := strings.LastIndex(authToken, "::"); i >= 0 {
			authToken = strings.TrimSpace(authToken[i+2:])
		}
	}
	if authToken == "" {
		return Result{OK: false, Detail: "Token 为空"}
	}

	sessionID := uuid.NewSHA1(uuid.NameSpaceDNS, []byte(authToken)).String()
	clientKey := hashed64Hex(authToken, "")
	checksum := generateCursorChecksum(authToken)
	conversationID := uuid.NewString()
	reqID := uuid.NewString()

	body := generateCursorAgentFramedBody(probeHiText, "default")
	fullURL := cursorBackendURL + cursorAgentPath
	req, err := http.NewRequest(http.MethodPost, fullURL, bytes.NewReader(body))
	if err != nil {
		r := cursorProbeResult(false, err.Error(), 0, body, nil, nil)
		return r
	}
	req.Header.Set("Authorization", "Bearer "+authToken)
	req.Header.Set("Connect-Accept-Encoding", "gzip")
	req.Header.Set("Connect-Protocol-Version", "1")
	req.Header.Set("Content-Type", "application/connect+proto")
	req.Header.Set("User-Agent", "connect-es/1.6.1")
	req.Header.Set("X-Amzn-Trace-Id", "Root="+reqID)
	req.Header.Set("X-Client-Key", clientKey)
	req.Header.Set("X-Cursor-Checksum", checksum)
	req.Header.Set("X-Cursor-Client-Version", cursorEnvVersion())
	req.Header.Set("X-Cursor-Client-Type", "ide")
	req.Header.Set("X-Cursor-Client-Os", cursorClientOS())
	req.Header.Set("X-Cursor-Client-Arch", cursorClientArch())
	req.Header.Set("X-Cursor-Client-Os-Version", "unknown")
	req.Header.Set("X-Cursor-Client-Device-Type", "desktop")
	req.Header.Set("X-Cursor-Config-Version", uuid.NewString())
	req.Header.Set("X-Cursor-Timezone", "UTC")
	req.Header.Set("X-Ghost-Mode", "false")
	req.Header.Set("X-New-Onboarding-Completed", "true")
	req.Header.Set("X-Request-Id", reqID)
	req.Header.Set("X-Session-Id", sessionID)
	req.Header.Set("X-Conversation-Id", conversationID)
	req.Host = "api2.cursor.sh"

	resp, err := cursorProbeHTTPClient.Do(req)
	if err != nil {
		return cursorProbeResult(false, "请求 Cursor Agent 失败: "+err.Error(), 0, body, nil, nil)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(io.LimitReader(resp.Body, cursorHiMaxRead))
		decoded, decodeNote := decodeCursorResponseBody(raw, resp.Header.Get("Content-Encoding"))
		blocked, reason := classifyCursorRawStream(decoded)
		if blocked {
			return cursorProbeResult(false, reason+"；"+decodeNote, resp.StatusCode, body, raw, decoded)
		}
		detail := fmt.Sprintf("HTTP %d（非 200）；%s；说明与协议边界见 streamNote", resp.StatusCode, decodeNote)
		return cursorProbeResult(false, detail, resp.StatusCode, body, raw, decoded)
	}

	var buf bytes.Buffer
	_, _ = io.Copy(&buf, io.LimitReader(resp.Body, cursorHiMaxRead))
	raw := buf.Bytes()
	decoded, decodeNote := decodeCursorResponseBody(raw, resp.Header.Get("Content-Encoding"))
	blocked, reason := classifyCursorRawStream(decoded)
	if blocked {
		return cursorProbeResult(false, reason+"；"+decodeNote, resp.StatusCode, body, raw, decoded)
	}
	detail := "HTTP 200；未命中内置英文关键词；" + decodeNote + "；二进制流含义与 ok 边界见 streamNote"
	return cursorProbeResult(true, detail, resp.StatusCode, body, raw, decoded)
}