package controller

import (
	"bytes"
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/middleware"
	"github.com/QuantumNous/new-api/model"
	benchmark "github.com/QuantumNous/new-api/pkg/model_benchmark"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
	"github.com/QuantumNous/new-api/relay"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/dto"
	"github.com/QuantumNous/new-api/relaykit/types"
	"github.com/QuantumNous/new-api/service"
	"github.com/gin-gonic/gin"
)

type astraSamplingHandler struct{}

type astraSamplingTaskPayload struct {
	Force bool `json:"force,omitempty"`
}

func (astraSamplingHandler) Type() string               { return service.AstraSamplingTaskType }
func (astraSamplingHandler) Enabled() bool              { return service.AstraSamplingEnabled() }
func (astraSamplingHandler) Interval() time.Duration    { return service.AstraSamplingSchedulerInterval() }
func (astraSamplingHandler) ScheduleFromCreation() bool { return true }
func (astraSamplingHandler) NewPayload() any            { return nil }

func (astraSamplingHandler) ScheduleDue(ctx context.Context, now time.Time) (bool, error) {
	channels, err := service.AstraSamplingChannelsDue(ctx, now)
	return len(channels) > 0, err
}

func (astraSamplingHandler) Run(parent context.Context, task *model.SystemTask, runnerID string) {
	ctx, cancel := context.WithTimeout(parent, 25*time.Minute)
	defer cancel()
	payload := astraSamplingTaskPayload{}
	if err := task.DecodePayload(&payload); err != nil {
		finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
		return
	}
	channels, err := service.AstraSamplingChannelsDue(ctx, time.Now())
	if payload.Force {
		channels, err = service.AstraSamplingChannels(ctx)
	}
	if err != nil {
		finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, errors.New("channel query failed"))
		return
	}
	state := service.AstraSamplingState{Rounds: []service.AstraSamplingRound{}}
	var mu sync.Mutex
	var persistErr error
	// Persist before sending, and after every probe. A failed lease or restart
	// never resumes a paid probe; unfinished records remain visible as errors.
	persist := func() bool {
		if persistErr != nil {
			return false
		}
		persistErr = model.UpdateSystemTaskState(task.TaskID, runnerID, state)
		if persistErr != nil {
			cancel()
		}
		return persistErr == nil
	}
	jobs := make(chan int)
	var workers sync.WaitGroup
	for worker := 0; worker < 2; worker++ {
		workers.Add(1)
		go func() {
			defer workers.Done()
			for id := range jobs {
				if ctx.Err() != nil {
					continue
				}
				mu.Lock()
				index := len(state.Rounds)
				state.Rounds = append(state.Rounds, service.AstraSamplingRound{ID: common.GetRandomString(16), ChannelID: id, StartedAt: time.Now().Unix(), Version: fingerprint.Version, Result: fingerprint.Result{Status: "running"}})
				ok := persist()
				mu.Unlock()
				if !ok {
					continue
				}
				channel, loadErr := model.GetChannelById(id, true)
				samples := []fingerprint.Sample{}
				reason := ""
				usage := []fingerprint.ProbeResponse{}
				if loadErr != nil || channel.Status != common.ChannelStatusEnabled {
					reason = "channel_unavailable"
				} else {
					for probe := 0; probe < fingerprint.ProbeCount && ctx.Err() == nil; probe++ {
						mu.Lock()
						state.Rounds[index].Sent++
						ok = persist()
						mu.Unlock()
						if !ok {
							break
						}
						answer, count, probeErr := runAstraFingerprintProbe(ctx, channel, probe)
						usage = append(usage, answer)
						if probeErr != nil {
							if errors.Is(probeErr, fingerprint.ErrInvalidSample) {
								reason = "invalid_sample"
							} else {
								reason = "transport_error"
								break
							}
						} else {
							samples = append(samples, fingerprint.Sample{Text: answer.Text, Count: count})
						}
						mu.Lock()
						state.Rounds[index].Usage = append([]fingerprint.ProbeResponse(nil), usage...)
						ok = persist()
						mu.Unlock()
						if !ok {
							break
						}
					}
				}
				result, scoreErr := fingerprint.Score(samples)
				if scoreErr != nil {
					reason = "scoring_error"
				}
				if ctx.Err() != nil {
					reason = "interrupted"
				}
				if reason != "" && reason != "invalid_sample" {
					result.Status = "error"
				}
				mu.Lock()
				state.Rounds[index].Result = result
				state.Rounds[index].Reason = reason
				state.Rounds[index].Usage = usage
				state.Rounds[index].FinishedAt = time.Now().Unix()
				if persist() && !payload.Force {
					if err := model.ApplyDetectionScheduleResult(ctx, id, "modeltrace", state.Rounds[index].ID, state.Rounds[index].FinishedAt, service.FingerprintScheduleOutcome(service.AstraSamplingResultStatus(result))); err != nil {
						persistErr = err
						cancel()
					}
				}
				mu.Unlock()
			}
		}()
	}
	for _, channel := range channels {
		select {
		case jobs <- channel.Id:
		case <-ctx.Done():
		}
	}
	close(jobs)
	workers.Wait()
	if persistErr != nil || ctx.Err() != nil {
		finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, errors.New("sampling interrupted"))
		return
	}
	finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, map[string]any{"channels": len(state.Rounds), "model": fingerprint.Model}, nil)
}

func runAstraFingerprintProbe(parent context.Context, channel *model.Channel, index int) (fingerprint.ProbeResponse, int, error) {
	return runModelDetectionProbe(parent, channel, index, "fingerprint")
}

func runModelDetectionProbe(parent context.Context, channel *model.Channel, index int, kind string) (fingerprint.ProbeResponse, int, error) {
	ctx, cancel := context.WithTimeout(parent, 180*time.Second)
	defer cancel()
	var probe fingerprint.Probe
	var err error
	if kind == "fingerprint" {
		probe, err = fingerprint.ProbeAt(index)
	}
	if err != nil {
		return fingerprint.ProbeResponse{}, index, err
	}
	count, output := probe.Count, fingerprint.OutputTokens
	messages := []map[string]any{{"role": "user", "content": probe.Prompt}}
	if kind == "benchmark" {
		secondary, err := benchmark.ProbeAt(index)
		if err != nil {
			return fingerprint.ProbeResponse{}, index, err
		}
		count, output = index, benchmark.OutputTokens
		messages = []map[string]any{{"role": "system", "content": secondary.System}, {"role": "user", "content": secondary.Prompt}}
	}
	if kind == "candy" {
		count, output = 1, service.CandyOutputTokens
		messages = []map[string]any{{"role": "user", "content": service.CandyPrompt}}
	}
	input, _ := common.Marshal(messages)
	stream := true
	limit := uint(output)
	request := &dto.OpenAIResponsesRequest{Model: fingerprint.Model, Input: input, Stream: &stream, MaxOutputTokens: &limit, Store: []byte("false"), Reasoning: &dto.Reasoning{Effort: "low"}}
	encoded, err := common.Marshal(request)
	if err != nil {
		return fingerprint.ProbeResponse{}, count, err
	}
	c, _ := gin.CreateTestContext(httptest.NewRecorder())
	c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, "/v1/responses", bytes.NewReader(encoded))
	c.Request.Header.Set("Content-Type", "application/json")
	if apiErr := middleware.SetupContextForSelectedChannel(c, channel, fingerprint.Model); apiErr != nil {
		return fingerprint.ProbeResponse{}, count, errors.New("channel setup failed")
	}
	info, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAIResponses, request, nil)
	if err != nil {
		return fingerprint.ProbeResponse{}, count, err
	}
	info.IsChannelTest = true
	info.DisablePing = true
	info.AstraDetectionDisabled = true
	info.AstraAutoSwitch = false
	adaptor, body, closer, apiErr := relay.PrepareResponsesRequest(c, info, request)
	if closer != nil {
		defer closer.Close()
	}
	if apiErr != nil {
		return fingerprint.ProbeResponse{}, count, errors.New("probe preparation failed")
	}
	{
		bounded, boundedCloser, boundErr := boundDetectionBody(body, request)
		if boundErr != nil {
			return fingerprint.ProbeResponse{}, count, boundErr
		}
		defer boundedCloser.Close()
		body = bounded
	}
	response, err := adaptor.DoRequest(c, info, body)
	if err != nil {
		return fingerprint.ProbeResponse{}, count, errors.New("probe transport failed")
	}
	resp, ok := response.(*http.Response)
	if !ok || resp == nil {
		return fingerprint.ProbeResponse{}, count, errors.New("unsupported probe transport")
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fingerprint.ProbeResponse{}, count, errors.New("probe HTTP failure")
	}
	// Deliberately bypass DoResponse: this is a platform-funded probe, with no
	// customer token, wallet charge, session binding or generated response relay.
	answer, err := fingerprint.ReadResponse(resp.Body, resp.Header.Get("Content-Type"))
	return answer, count, err
}

func TestAstraFingerprint(c *gin.Context) {
	if !service.AstraSamplingEnabled() {
		c.JSON(http.StatusConflict, gin.H{"success": false, "message": "Astra sampling is disabled"})
		return
	}
	latest, err := model.GetLatestSystemTasks([]string{service.AstraSamplingTaskType})
	if err != nil {
		common.ApiError(c, err)
		return
	}
	if last := latest[service.AstraSamplingTaskType]; last != nil && time.Now().Unix()-last.CreatedAt < 300 {
		c.JSON(http.StatusTooManyRequests, gin.H{"success": false, "message": "Please wait five minutes between sampling rounds"})
		return
	}
	task, created, err := service.EnqueueSystemTask(service.AstraSamplingTaskType, astraSamplingTaskPayload{Force: true})
	if err != nil {
		common.ApiError(c, err)
		return
	}
	c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"task_id": task.TaskID, "created": created}})
}