package service

import (
	"context"
	"sort"
	"strings"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
)

const AstraSamplingTaskType = "modeltrace_sampling"

const (
	AstraSamplingWindow           = 6 * time.Hour
	AstraSamplingDefaultInterval  = 30 * time.Minute
	AstraSamplingUnstableInterval = 30 * time.Minute
	AstraAvailabilityInterval     = 15 * time.Minute
)

type AstraSamplingRound struct {
	ID         string                      `json:"id"`
	ChannelID  int                         `json:"channel_id"`
	StartedAt  int64                       `json:"started_at"`
	FinishedAt int64                       `json:"finished_at,omitempty"`
	Sent       int                         `json:"sent"`
	Version    string                      `json:"version"`
	Result     fingerprint.Result          `json:"result"`
	Reason     string                      `json:"reason,omitempty"`
	Usage      []fingerprint.ProbeResponse `json:"usage,omitempty"`
}
type AstraSamplingState struct {
	Rounds []AstraSamplingRound `json:"rounds"`
}

type ChannelAvailabilityRound struct {
	ChannelID int   `json:"channel_id"`
	At        int64 `json:"at"`
	Success   bool  `json:"success"`
}

type ChannelAvailabilityState struct {
	Rounds []ChannelAvailabilityRound `json:"rounds"`
}

type AstraSamplingHistory struct {
	UserInitiated bool    `json:"user_initiated,omitempty"`
	ID            string  `json:"id"`
	At            int64   `json:"at"`
	Status        string  `json:"status"`
	ValidSamples  int     `json:"valid_samples"`
	Closest       string  `json:"closest,omitempty"`
	Gap           float64 `json:"score_gap,omitempty"`
}
type AstraSamplingSummary struct {
	LatestValid                 *AstraSamplingHistory  `json:"-"`
	Enabled                     bool                   `json:"enabled"`
	Available                   bool                   `json:"available"`
	Channels                    int                    `json:"channels"`
	SampledChannels             int                    `json:"sampled_channels"`
	Matches                     int                    `json:"matches"`
	Mismatches                  int                    `json:"mismatches"`
	Inconclusive                int                    `json:"inconclusive"`
	Errors                      int                    `json:"errors"`
	Running                     int                    `json:"running"`
	Rate                        *float64               `json:"rate"`
	WindowHours                 int                    `json:"window_hours"`
	IntervalMinutes             int                    `json:"interval_minutes"`
	IntervalMaxMinutes          int                    `json:"interval_max_minutes"`
	AvailabilityOnly            bool                   `json:"availability_only,omitempty"`
	AvailabilityIntervalMinutes int                    `json:"availability_interval_minutes,omitempty"`
	AvailabilityChecks          int                    `json:"availability_checks,omitempty"`
	AvailabilitySuccesses       int                    `json:"availability_successes,omitempty"`
	AvailabilityRate            *float64               `json:"availability_rate,omitempty"`
	History                     []AstraSamplingHistory `json:"history"`
}

func AstraSamplingEnabled() bool {
	return common.GetEnvOrDefaultBool("ASTRA_FINGERPRINT_ENABLED", false)
}

func AstraAvailabilityEnabled() bool {
	return common.GetEnvOrDefaultBool("CHANNEL_AVAILABILITY_PROBE_ENABLED", true)
}

func ChannelSupportsAstra(channel *model.Channel) bool {
	if channel == nil {
		return false
	}
	for _, name := range channel.GetModels() {
		if strings.TrimSpace(name) == fingerprint.Model {
			return true
		}
	}
	return false
}

func AstraSamplingIntervalForGroup(group string) time.Duration {
	return AstraSamplingDefaultInterval
}

func AstraSamplingIntervalForChannel(channel model.Channel) time.Duration {
	interval := AstraSamplingDefaultInterval
	for _, group := range channel.GetGroups() {
		interval = min(interval, AstraSamplingIntervalForGroup(group))
	}
	return interval
}

func AstraSamplingSchedulerInterval() time.Duration {
	return 15 * time.Second
}

func AstraAvailabilityIntervalMinutes() int {
	minutes := common.GetEnvOrDefault("CHANNEL_AVAILABILITY_INTERVAL_MINUTES", int(AstraAvailabilityInterval/time.Minute))
	if minutes <= 0 {
		return int(AstraAvailabilityInterval / time.Minute)
	}
	return minutes
}

// Public statistics classify completed rounds by their closest valid fingerprint,
// including partial rounds. Raw probe results retain the scorer's confidence.
func AstraSamplingResultStatus(result fingerprint.Result) string {
	if result.Engine != "" {
		return result.Status
	}
	closest := strings.TrimSpace(result.Closest)
	if result.Status != "running" && result.ValidSamples > 0 && closest != "" {
		if strings.EqualFold(closest, fingerprint.Model) {
			return "match"
		}
		return "mismatch"
	}
	if result.Status == "match" || result.Status == "mismatch" {
		return "inconclusive"
	}
	return result.Status
}

// No credentials are fetched by the public status query. Each eligible channel
// is sampled once even when it belongs to multiple groups.
func AstraSamplingChannels(ctx context.Context) ([]model.Channel, error) {
	return samplingChannels(ctx, true)
}

func samplingChannels(ctx context.Context, requireAstra bool) ([]model.Channel, error) {
	var rows []model.Channel
	err := model.DB.WithContext(ctx).Select("id", "type", "group", "models", "status").Where("status = ?", common.ChannelStatusEnabled).Find(&rows).Error
	if err != nil {
		return nil, err
	}
	channels := make([]model.Channel, 0, len(rows))
	for _, channel := range rows {
		if !requireAstra {
			channels = append(channels, channel)
			continue
		}
		if ChannelSupportsAstra(&channel) {
			channels = append(channels, channel)
		}
	}
	return channels, nil
}

func AstraSamplingChannelsDue(ctx context.Context, now time.Time) ([]model.Channel, error) {
	due, _, err := AutomaticDetectionChannelsDue(ctx, "modeltrace", now)
	return due, err
}

func AstraSamplingSummaries(ctx context.Context, now time.Time, settings map[string]model.StatusPageGroup) (map[string]*AstraSamplingSummary, error) {
	channels, err := samplingChannels(ctx, false)
	if err != nil {
		return nil, err
	}
	var tasks []*model.SystemTask
	err = model.DB.WithContext(ctx).Select("status", "state").Where("type = ? AND created_at >= ?", AstraSamplingTaskType, now.Add(-AstraSamplingWindow).Unix()).Order("id asc").Find(&tasks).Error
	if err != nil {
		return nil, err
	}
	var availabilityTasks []*model.SystemTask
	err = model.DB.WithContext(ctx).Select("status", "state").Where("type = ? AND created_at >= ?", model.SystemTaskTypeChannelAvailability, now.Add(-AstraSamplingWindow).Unix()).Order("id asc").Find(&availabilityTasks).Error
	if err != nil {
		return nil, err
	}
	summaries := aggregateAstraSampling(channels, tasks, now, settings)
	AggregateChannelAvailability(summaries, channels, availabilityTasks, now)
	ranges, err := DetectionScheduleRanges(ctx, channels, "modeltrace")
	if err != nil {
		return nil, err
	}
	for group, summary := range summaries {
		if r, ok := ranges[group]; ok {
			summary.IntervalMinutes, summary.IntervalMaxMinutes = r[0], r[1]
		} else {
			summary.IntervalMaxMinutes = summary.IntervalMinutes
		}
	}
	return summaries, nil
}

func AggregateAstraSampling(channels []model.Channel, tasks []*model.SystemTask, now time.Time) map[string]*AstraSamplingSummary {
	return aggregateAstraSampling(channels, tasks, now, nil)
}

func aggregateAstraSampling(channels []model.Channel, tasks []*model.SystemTask, now time.Time, settings map[string]model.StatusPageGroup) map[string]*AstraSamplingSummary {
	groups := map[string]*AstraSamplingSummary{}
	channelGroups := map[int][]string{}
	seen := map[string]map[int]bool{}
	for _, channel := range channels {
		channelGroups[channel.Id] = channel.GetGroups()
		astra := ChannelSupportsAstra(&channel)
		for _, group := range channel.GetGroups() {
			if groups[group] == nil {
				groups[group] = &AstraSamplingSummary{Enabled: AstraSamplingEnabled(), Available: true, WindowHours: int(AstraSamplingWindow / time.Hour), IntervalMinutes: int(AstraSamplingIntervalForGroup(group) / time.Minute), History: []AstraSamplingHistory{}}
				seen[group] = map[int]bool{}
			}
			if astra {
				groups[group].Channels++
			}
		}
	}
	for _, task := range tasks {
		var state AstraSamplingState
		if task.DecodeState(&state) != nil {
			continue
		}
		for _, round := range state.Rounds {
			if round.Result.Engine != fingerprint.Version && round.Result.Status != "running" {
				continue
			}
			if round.StartedAt < now.Add(-AstraSamplingWindow).Unix() {
				continue
			}
			if !containsAstraChannel(channels, round.ChannelID) {
				continue
			}
			status := AstraSamplingResultStatus(round.Result)
			if status == "running" && (task.Status == model.SystemTaskStatusFailed || task.Status == model.SystemTaskStatusSucceeded) {
				status = "error"
			}
			for _, group := range channelGroups[round.ChannelID] {
				if cutoff := settings[group].ResetBefore; cutoff > 0 && round.StartedAt <= cutoff {
					continue
				}
				s := groups[group]
				switch status {
				case "match":
					s.Matches++
				case "mismatch":
					s.Mismatches++
				case "inconclusive":
					s.Inconclusive++
				case "error":
					s.Errors++
				case "running":
					s.Running++
				default:
					continue
				}
				if round.FinishedAt > 0 {
					seen[group][round.ChannelID] = true
				}
				s.History = append(s.History, AstraSamplingHistory{ID: round.ID, At: round.StartedAt, Status: status, ValidSamples: round.Result.ValidSamples, Closest: round.Result.Closest, Gap: round.Result.Gap})
			}
		}
	}
	for group, s := range groups {
		if s.Channels == 0 {
			s.AvailabilityOnly = true
			s.IntervalMinutes = AstraAvailabilityIntervalMinutes()
			s.AvailabilityIntervalMinutes = AstraAvailabilityIntervalMinutes()
		}
		s.SampledChannels = len(seen[group])
		if n := s.Matches + s.Mismatches; n > 0 {
			rate := float64(s.Mismatches) / float64(n)
			s.Rate = &rate
		}
		sort.SliceStable(s.History, func(i, j int) bool { return s.History[i].At > s.History[j].At })
		for i := range s.History {
			round := &s.History[i]
			if round.ValidSamples > 0 && round.Closest != "" && (round.Status == "match" || round.Status == "mismatch") {
				s.LatestValid = round
				break
			}
		}
		if len(s.History) > 12 {
			s.History = s.History[:12]
		}
	}
	return groups
}

func containsAstraChannel(channels []model.Channel, id int) bool {
	for i := range channels {
		if channels[i].Id == id {
			return ChannelSupportsAstra(&channels[i])
		}
	}
	return false
}

func AggregateChannelAvailability(summaries map[string]*AstraSamplingSummary, channels []model.Channel, tasks []*model.SystemTask, now time.Time) {
	channelGroups := map[int][]string{}
	for i := range channels {
		if !ChannelSupportsAstra(&channels[i]) {
			channelGroups[channels[i].Id] = channels[i].GetGroups()
		}
	}
	for _, task := range tasks {
		var state ChannelAvailabilityState
		if task.DecodeState(&state) != nil {
			continue
		}
		for _, round := range state.Rounds {
			if round.At < now.Add(-AstraSamplingWindow).Unix() {
				continue
			}
			for _, group := range channelGroups[round.ChannelID] {
				s := summaries[group]
				s.AvailabilityChecks++
				if round.Success {
					s.AvailabilitySuccesses++
				}
			}
		}
	}
	for _, s := range summaries {
		if s.AvailabilityChecks > 0 {
			rate := float64(s.AvailabilitySuccesses) / float64(s.AvailabilityChecks)
			s.AvailabilityRate = &rate
		}
	}
}