package service

import (
	"context"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type anchoredSamplingTestHandler struct{ stubScheduledHandler }

func (*anchoredSamplingTestHandler) ScheduleFromCreation() bool { return true }

func TestAstraSamplingThirtyMinuteSchedulePersistsAcrossRestart(t *testing.T) {
	truncate(t)
	handler := &anchoredSamplingTestHandler{stubScheduledHandler{taskType: AstraSamplingTaskType, enabled: true, interval: 30 * time.Minute}}
	withSystemTaskRegistry(t, handler)
	now := common.GetTimestamp()
	state, _ := common.Marshal(AstraSamplingState{Rounds: []AstraSamplingRound{{ChannelID: 42, StartedAt: now - 1900, FinishedAt: now - 1800, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", ValidSamples: 3, Closest: fingerprint.Model}}}})
	row := model.SystemTask{TaskID: "sampling-old", Type: AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 1900, UpdatedAt: now - 10, State: string(state)}
	if err := model.DB.Create(&row).Error; err != nil {
		t.Fatal(err)
	}
	runSystemTaskScheduler()
	if countSystemTasks(t, AstraSamplingTaskType) != 2 {
		t.Fatal("schedule drifted from completion instead of creation")
	}
	runSystemTaskScheduler()
	if countSystemTasks(t, AstraSamplingTaskType) != 2 {
		t.Fatal("restart duplicated pending job")
	}
	channels, err := AstraSamplingChannels(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	_ = channels
	loaded, err := model.GetSystemTaskByTaskID(row.TaskID)
	if err != nil {
		t.Fatal(err)
	}
	summary := AggregateAstraSampling([]model.Channel{{Id: 42, Group: "group", Models: fingerprint.Model}}, []*model.SystemTask{loaded}, time.Now())
	if summary["group"].Matches != 1 {
		t.Fatal("completed results did not survive reload")
	}
}

func TestAstraSamplingAggregationPrivacyAndFailures(t *testing.T) {
	now := time.Now()
	channels := []model.Channel{{Id: 1, Group: "first,shared", Models: fingerprint.Model}, {Id: 2, Group: "first", Models: fingerprint.Model}, {Id: 3, Group: "empty", Models: fingerprint.Model}}
	rounds := []AstraSamplingRound{}
	for _, status := range []string{"match", "mismatch", "inconclusive", "error", "running"} {
		closest := ""
		if status == "match" {
			closest = fingerprint.Model
		}
		if status == "mismatch" {
			closest = "gpt-5.6-luna"
		}
		rounds = append(rounds, AstraSamplingRound{ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: status, ValidSamples: 3, Closest: closest}})
	}
	rounds = append(rounds, AstraSamplingRound{ChannelID: 99, StartedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch"}})
	state, _ := common.Marshal(AstraSamplingState{Rounds: rounds})
	summaries := AggregateAstraSampling(channels, []*model.SystemTask{{Status: model.SystemTaskStatusFailed, State: string(state)}}, now)
	first := summaries["first"]
	if first.Rate == nil || *first.Rate != .5 || first.Errors != 2 || first.Inconclusive != 1 || first.Running != 0 || first.SampledChannels != 1 || first.Channels != 2 {
		t.Fatalf("%+v", first)
	}
	if *summaries["shared"].Rate != .5 || summaries["empty"].Rate != nil {
		t.Fatal("shared results / empty denominator")
	}
	public, _ := common.Marshal(summaries)
	var raw map[string]any
	if common.Unmarshal(public, &raw) != nil {
		t.Fatal("invalid public JSON")
	}
	for _, h := range first.History {
		if h.Status == "running" {
			t.Fatal("restart must fail unfinished sample")
		}
	}
	if len(first.History) != 5 {
		t.Fatal("deleted/unrelated channel leaked into current group")
	}
}

func TestAstraSamplingExcludesInconclusiveFromRateWithinSixHours(t *testing.T) {
	now := time.Now()
	channel := model.Channel{Id: 1, Group: "gpt-s-unstable", Models: fingerprint.Model}
	rounds := []AstraSamplingRound{
		{ID: "recent-match", ChannelID: 1, StartedAt: now.Add(-time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", ValidSamples: 3, Closest: fingerprint.Model}},
		{ID: "recent-luna", ChannelID: 1, StartedAt: now.Add(-2 * time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-luna", Gap: .9}},
		{ID: "old-luna", ChannelID: 1, StartedAt: now.Add(-7 * time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-luna", Gap: .9}},
	}
	state, _ := common.Marshal(AstraSamplingState{Rounds: rounds})
	summary := AggregateAstraSampling([]model.Channel{channel}, []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(state)}}, now)["gpt-s-unstable"]
	if summary.Rate == nil || *summary.Rate != 0 || summary.Matches != 1 || summary.Mismatches != 0 || summary.Inconclusive != 1 {
		t.Fatalf("%+v", summary)
	}
	if summary.WindowHours != 6 || summary.IntervalMinutes != 30 || len(summary.History) != 2 {
		t.Fatalf("unexpected window, interval or history: %+v", summary)
	}
}

func TestAstraSamplingAggregationMarksNonAstraGroupsAvailabilityOnly(t *testing.T) {
	t.Setenv("CHANNEL_AVAILABILITY_INTERVAL_MINUTES", "30")
	now := time.Now()
	channels := []model.Channel{{Id: 1, Group: "opencode", Models: "gpt-5.6-sol"}}
	summary := AggregateAstraSampling(channels, nil, now)["opencode"]
	if !summary.AvailabilityOnly || summary.Channels != 0 || summary.IntervalMinutes != 30 || summary.Rate != nil {
		t.Fatalf("%+v", summary)
	}
	state, _ := common.Marshal(ChannelAvailabilityState{Rounds: []ChannelAvailabilityRound{
		{ChannelID: 1, At: now.Add(-time.Hour).Unix(), Success: true},
		{ChannelID: 1, At: now.Add(-2 * time.Hour).Unix(), Success: false},
		{ChannelID: 1, At: now.Add(-7 * time.Hour).Unix(), Success: false},
	}})
	AggregateChannelAvailability(map[string]*AstraSamplingSummary{"opencode": summary}, channels, []*model.SystemTask{{State: string(state)}}, now)
	if summary.AvailabilityRate == nil || *summary.AvailabilityRate != .5 || summary.AvailabilityChecks != 2 || summary.AvailabilitySuccesses != 1 {
		t.Fatalf("%+v", summary)
	}
}

func TestAstraSamplingExcludesOldEngineAndUncertainCandidates(t *testing.T) {
	now := time.Now()
	results := []fingerprint.Result{
		{Engine: "meow-4.5.4", Status: "mismatch", ValidSamples: 32, Closest: "other"},
		{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "other"},
		{Engine: fingerprint.Version, Status: "match", ValidSamples: 3, Closest: fingerprint.Model},
	}
	var rounds []AstraSamplingRound
	for _, result := range results {
		rounds = append(rounds, AstraSamplingRound{ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: result})
	}
	state, err := common.Marshal(AstraSamplingState{Rounds: rounds})
	require.NoError(t, err)
	summary := AggregateAstraSampling([]model.Channel{{Id: 1, Group: "test", Models: fingerprint.Model}}, []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(state)}}, now)["test"]
	require.NotNil(t, summary.Rate)
	assert.Equal(t, 0.0, *summary.Rate)
	assert.Equal(t, 1, summary.Matches)
	assert.Equal(t, 0, summary.Mismatches)
	assert.Equal(t, 1, summary.Inconclusive)
	assert.Len(t, summary.History, 2)
}

func TestModelTraceVerdictDoesNotPromoteClosestCandidate(t *testing.T) {
	for _, status := range []string{"match", "mismatch", "inconclusive", "error"} {
		assert.Equal(t, status, AstraSamplingResultStatus(fingerprint.Result{Engine: fingerprint.Version, Status: status, ValidSamples: 32, Closest: fingerprint.Model, Weight: .63436}))
	}
}