package tokenprobe

import (
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

var httpClient = &http.Client{
	Timeout: 12 * time.Second,
	Transport: &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	},
}

type Result struct {
	OK                   bool   `json:"ok"`
	Detail               string `json:"detail"`
	HTTPStatus           int    `json:"httpStatus"`
	ProbeMessage         string `json:"probeMessage,omitempty"`
	Endpoint             string `json:"endpoint,omitempty"`
	BytesRead            int    `json:"bytesRead,omitempty"`
	RawPreview           string `json:"rawPreview,omitempty"`
	RequestBodyPrefixHex string `json:"requestBodyPrefixHex,omitempty"`
	StreamProtocol       string `json:"streamProtocol,omitempty"`
	StreamNote           string `json:"streamNote,omitempty"`
}

func ProbeOfficial(module, rawToken string) Result {
	tok := normalizeBearerToken(strings.TrimSpace(rawToken))
	if tok == "" {
		return Result{OK: false, Detail: "Token 为空"}
	}
	switch module {
	case "cursor":
		return probeCursor(tok)
	case "windsurf":
		return probeWindsurf(tok)
	case "krio":
		return probeKiro(tok)
	default:
		return Result{OK: false, Detail: "未知模块"}
	}
}

func normalizeBearerToken(s string) string {
	s = strings.TrimSpace(s)
	if i := strings.LastIndex(s, "::"); i >= 0 {
		return strings.TrimSpace(s[i+2:])
	}
	return s
}

func probeCursor(token string) Result {
	url := "https://api2.cursor.sh/auth/full_stripe_profile"
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return Result{OK: false, Detail: "构造请求失败: " + err.Error()}
	}

	req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(token))
	req.Header.Set("X-Cursor-Client-Version", "3.0.16")
	req.Header.Set("X-New-Onboarding-Completed", "false")
	req.Header.Set("X-Ghost-Mode", "true")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Cursor/3.0.16 Chrome/142.0.7444.265 Electron/39.8.1 Safari/537.36")
	req.Header.Set("Accept", "*/*")
	req.Header.Set("Origin", "vscode-file://vscode-app")
	req.Header.Set("Accept-Encoding", "gzip, deflate, br, zstd")
	req.Header.Set("Accept-Language", "zh-CN")
	req.Header.Set("Priority", "u=1, i")

	resp, err := httpClient.Do(req)
	if err != nil {
		return Result{OK: false, Detail: "请求官方账单接口超时/网络失败: " + err.Error(), HTTPStatus: 500}
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(io.LimitReader(resp.Body, 16384))
	jsonStr := string(body)

	res := Result{
		HTTPStatus:     resp.StatusCode,
		Endpoint:       url,
		BytesRead:      len(body),
		RawPreview:     jsonStr,
		StreamProtocol: "HTTP/2 JSON REST",
		StreamNote:     "2026型画像探测",
		ProbeMessage:   "GET full_stripe_profile",
	}

	if resp.StatusCode != http.StatusOK {
		res.OK = false
		res.Detail = fmt.Sprintf("Token已失效或被官方拉黑（HTTP %d）", resp.StatusCode)
		return res
	}

	if strings.Contains(jsonStr, `"noModelsRemaining":true`) ||
		strings.Contains(jsonStr, `"is_usage_limited":true`) ||
		strings.Contains(jsonStr, `"hard_limit_reached"`) ||
		strings.Contains(jsonStr, `"blocked"`) ||
		(strings.Contains(jsonStr, `"membershipType":"free"`) && strings.Contains(jsonStr, `"trialEligible":false`)) {
		res.OK = false
		res.Detail = "Token存活，但属于无额度Free空壳号（上号必弹付费墙）"
		return res
	}

	if len(jsonStr) < 10 {
		res.OK = false
		res.Detail = "官方接口返回异常空数据"
		return res
	}

	res.OK = true
	res.Detail = "检测成功，高速算力/Agent额度健康"
	return res
}

func probeWindsurf(apiKey string) Result {
	payload := map[string]interface{}{
		"metadata": map[string]string{
			"apiKey":           apiKey,
			"ideName":          "windsurf",
			"ideVersion":       "0.0.0",
			"extensionName":    "windsurf",
			"extensionVersion": "0.0.0",
			"locale":           "zh",
		},
	}
	raw, err := json.Marshal(payload)
	if err != nil {
		return Result{OK: false, Detail: err.Error()}
	}
	req, err := http.NewRequest(
		http.MethodPost,
		"https://server.codeium.com/exa.seat_management_pb.SeatManagementService/GetUserStatus",
		bytes.NewReader(raw),
	)
	if err != nil {
		return Result{OK: false, Detail: err.Error()}
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Connect-Protocol-Version", "1")

	resp, err := httpClient.Do(req)
	if err != nil {
		return Result{OK: false, Detail: "请求失败: " + err.Error()}
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 8192))

	switch resp.StatusCode {
	case http.StatusOK:
		var wrap map[string]interface{}
		if json.Unmarshal(body, &wrap) == nil {
			if _, ok := wrap["userStatus"]; ok {
				return Result{OK: true, Detail: "Codeium 云端接口响应正常", HTTPStatus: resp.StatusCode}
			}
		}
		if bytes.Contains(body, []byte(`"planStatus"`)) || bytes.Contains(body, []byte(`"userStatus"`)) {
			return Result{OK: true, Detail: "Codeium 云端接口响应正常", HTTPStatus: resp.StatusCode}
		}
		return Result{OK: true, Detail: fmt.Sprintf("HTTP %d，已收到响应", resp.StatusCode), HTTPStatus: resp.StatusCode}
	case http.StatusUnauthorized, http.StatusForbidden:
		return Result{OK: false, Detail: fmt.Sprintf("API Key 无效或已失效（HTTP %d）", resp.StatusCode), HTTPStatus: resp.StatusCode}
	default:
		snip := strings.TrimSpace(string(body))
		if len(snip) > 220 {
			snip = snip[:220] + "…"
		}
		return Result{OK: false, Detail: fmt.Sprintf("HTTP %d %s", resp.StatusCode, snip), HTTPStatus: resp.StatusCode}
	}
}

func probeKiro(accessToken string) Result {
	arn := findProfileArnInJWT(accessToken)
	if arn == "" {
		return Result{
			OK:     false,
			Detail: "无法从 Token 中解析 profileArn，Kiro 暂无法自动探测",
		}
	}

	q := url.Values{}
	q.Set("origin", "AI_EDITOR")
	q.Set("profileArn", arn)
	q.Set("resourceType", "AGENTIC_REQUEST")
	u := "https://q.us-east-1.amazonaws.com/getUsageLimits?" + q.Encode()

	req, err := http.NewRequest(http.MethodGet, u, nil)
	if err != nil {
		return Result{OK: false, Detail: err.Error()}
	}
	req.Header.Set("Authorization", "Bearer "+normalizeBearerToken(accessToken))
	req.Header.Set("Accept", "application/json")

	resp, err := httpClient.Do(req)
	if err != nil {
		return Result{OK: false, Detail: "请求失败: " + err.Error()}
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))

	switch resp.StatusCode {
	case http.StatusOK:
		return Result{OK: true, Detail: "Kiro（AWS Q）用量接口响应正常", HTTPStatus: resp.StatusCode}
	case http.StatusUnauthorized, http.StatusForbidden:
		return Result{OK: false, Detail: fmt.Sprintf("Token 无效或已过期（HTTP %d）", resp.StatusCode), HTTPStatus: resp.StatusCode}
	default:
		snip := strings.TrimSpace(string(body))
		if len(snip) > 220 {
			snip = snip[:220] + "…"
		}
		return Result{OK: false, Detail: fmt.Sprintf("HTTP %d %s", resp.StatusCode, snip), HTTPStatus: resp.StatusCode}
	}
}

func decodeJWTPayloadMap(raw string) (map[string]interface{}, error) {
	tok := normalizeBearerToken(strings.TrimSpace(raw))
	parts := strings.Split(tok, ".")
	if len(parts) < 2 {
		return nil, fmt.Errorf("not a JWT")
	}
	b, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return nil, err
	}
	var m map[string]interface{}
	if err := json.Unmarshal(b, &m); err != nil {
		return nil, err
	}
	return m, nil
}

func findProfileArnInJWT(raw string) string {
	m, err := decodeJWTPayloadMap(raw)
	if err != nil {
		return ""
	}
	return findProfileArnValue(m)
}

func findProfileArnValue(v interface{}) string {
	switch x := v.(type) {
	case map[string]interface{}:
		for k, val := range x {
			lk := strings.ToLower(k)
			if lk == "profilearn" || lk == "profile_arn" {
				if s, ok := val.(string); ok && strings.Contains(s, "arn:") {
					return s
				}
			}
		}
		for _, val := range x {
			if s := findProfileArnValue(val); s != "" {
				return s
			}
		}
	case []interface{}:
		for _, el := range x {
			if s := findProfileArnValue(el); s != "" {
				return s
			}
		}
	case string:
		if strings.Contains(x, "arn:aws:codewhisperer") && strings.Contains(x, ":profile/") {
			return x
		}
	}
	return ""
}