// ModelTrace scorer and bank: xqy2006/ModelTrace d4131b30243dfa05e70180b5eedde742103f1d73 (MIT).
package model_fingerprint

import (
	"crypto/rand"
	_ "embed"
	"errors"
	"fmt"
	"math"
	"math/big"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode"
	"unicode/utf8"

	"github.com/QuantumNous/new-api/common"
	"github.com/grafana/sobek"
)

const Model = "gpt-6-astra"
const Version = "modeltrace-d4131b3"

const ProbeCount = 3
const OutputTokens = 4096

//go:embed bank.json
var bankJSON string

//go:embed scorer.js
var scorerJS string

var compileOnce sync.Once
var program *sobek.Program
var compileErr error

type Sample struct {
	Text  string `json:"text"`
	Count int    `json:"expected_count"`
}

type Candidate struct {
	Model  string  `json:"model"`
	Weight float64 `json:"probability"`
	Score  float64 `json:"score"`
}

type Result struct {
	Engine       string             `json:"engine"`
	Matches      map[string]float64 `json:"matches,omitempty"`
	Status       string             `json:"status"`
	ValidSamples int                `json:"valid_samples"`
	Closest      string             `json:"closest,omitempty"`
	Weight       float64            `json:"relative_weight,omitempty"`
	Gap          float64            `json:"score_gap,omitempty"`
}

type Probe struct {
	ID     string
	System string
	Prompt string
	Count  int
}

func ProbeAt(index int) (Probe, error) {
	if index < 0 || index >= ProbeCount {
		return Probe{}, errors.New("invalid probe index")
	}
	// Disjoint length ranges retain three independent, non-identical challenges.
	offset, err := rand.Int(rand.Reader, big.NewInt(13))
	if err != nil {
		return Probe{}, err
	}
	count := 292 + index*14 + int(offset.Int64())
	prompt := fmt.Sprintf("这是一次独立的数值选择记录。逐项选择 %d 个 1 到 355（含端点）的整数。", count) +
		"每个位置都要单独选择；不要从 1 开始计数，不要连续递增或递减，也不要采用等差、循环、重复区块或其他规则化模式。" +
		"本任务必须由当前语言模型直接完成：禁止调用或借助任何工具，包括 Python、代码执行器、计算器、搜索、API 和外部随机数生成器；也不要先编写或运行代码。" +
		"偶然重复是有效的；不要重新排列或修正已经写出的项目。数字之间用逗号或空格分隔均可。直接从第一个取值开始输出，不要在序列前重复数量、范围或任务说明。"
	return Probe{ID: fmt.Sprint(index), Prompt: prompt, Count: count}, nil
}

var integerPattern = regexp.MustCompile(`[0-9]+`)

// Match ModelTrace's longest numeric run parser; prose separates runs.
func ValidArray(text string, count int) bool {
	if len(text) > 16384 || !utf8.ValidString(text) || count < 0 {
		return false
	}
	longest, current, end := 0, 0, 0
	for _, span := range integerPattern.FindAllStringIndex(text, -1) {
		if strings.ContainsFunc(text[end:span[0]], unicode.IsLetter) {
			current = 0
		}
		value, err := strconv.Atoi(text[span[0]:span[1]])
		if err == nil && value >= 1 && value <= 355 {
			current++
		}
		longest = max(longest, current)
		end = span[1]
	}
	return longest >= max(80, int(math.Ceil(float64(count)*.55)))
}

func ValidSample(text string) bool { return ValidArray(text, 0) }

func Decide(candidates []Candidate, valid int) Result {
	r := Result{Status: "inconclusive", Engine: Version, ValidSamples: valid, Matches: map[string]float64{}}
	if len(candidates) < 2 {
		return r
	}
	for _, candidate := range candidates {
		r.Matches[candidate.Model] = candidate.Weight
	}
	top := candidates[0]
	r.Closest = top.Model
	r.Weight = top.Weight
	r.Gap = top.Score - candidates[1].Score
	if valid != 3 || math.IsNaN(r.Gap) || math.IsInf(r.Gap, 0) || top.Weight < .8 || r.Gap < .5 {
		return r
	}
	if top.Model == Model {
		r.Status = "match"
		return r
	}
	for _, candidate := range candidates {
		if candidate.Model == Model && candidate.Weight <= .2 && top.Score-candidate.Score >= .5 {
			r.Status = "mismatch"
		}
	}
	return r
}

// Score runs only the embedded, pinned scorer. No filesystem, network or host
// functions are exposed to the VM; upstream text is data, never JavaScript.
func Score(samples []Sample) (Result, error) {
	if len(samples) > ProbeCount {
		return Result{}, errors.New("too many fingerprint samples")
	}
	valid := make([]Sample, 0, 3)
	for _, sample := range samples {
		if ValidArray(sample.Text, sample.Count) {
			valid = append(valid, sample)
		}
	}
	if len(valid) == 0 {
		return Result{Status: "inconclusive", Engine: Version}, nil
	}
	compileOnce.Do(func() {
		program, compileErr = sobek.Compile("fingerprint", strings.ReplaceAll(scorerJS, "export ", ""), true)
	})
	if compileErr != nil {
		return Result{}, compileErr
	}
	vm := sobek.New()
	timer := time.AfterFunc(3*time.Second, func() { vm.Interrupt("fingerprint scoring timeout") })
	defer timer.Stop()
	if _, err := vm.RunProgram(program); err != nil {
		return Result{}, err
	}
	if err := vm.Set("bankJSON", bankJSON); err != nil {
		return Result{}, err
	}
	encoded, err := common.Marshal(valid)
	if err != nil {
		return Result{}, err
	}
	if err = vm.Set("samplesJSON", string(encoded)); err != nil {
		return Result{}, err
	}
	value, err := vm.RunString("JSON.stringify(analyzeGlobalOutputs(JSON.parse(samplesJSON), JSON.parse(bankJSON)))")
	if err != nil {
		return Result{}, err
	}
	var analysis struct {
		Results []Candidate `json:"results"`
	}
	if err = common.Unmarshal([]byte(value.String()), &analysis); err != nil {
		return Result{}, err
	}
	if len(analysis.Results) < 2 {
		return Result{}, errors.New("missing fingerprint candidates")
	}
	return Decide(analysis.Results, len(valid)), nil
}