package workbench

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/Tencent/WeKnora/internal/agent/token"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/google/uuid"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// This rate preserves the imported ledger. It is not exposed in member APIs.
const pointsPerToken int64 = 10

type meteringContextKey struct{}
type meteringAttribution struct{ AgentID, JobID, Action string }

func (s *Service) WrapChat(model chat.Chat) chat.Chat { return &meteredChat{model: model, service: s} }

type meteredChat struct {
	model   chat.Chat
	service *Service
}

func (m *meteredChat) GetModelID() string   { return m.model.GetModelID() }
func (m *meteredChat) GetModelName() string { return m.model.GetModelName() }

func (s *Service) reserveBalance(ctx context.Context, amount int64) (*BalanceHold, error) {
	t, ok := types.TenantIDFromContext(ctx)
	if !ok || t == 0 {
		return nil, ErrForbidden
	}
	u, _ := types.UserIDFromContext(ctx)
	// Scheduled maintenance is attributed to the workspace, never a random member.
	if u == "" || types.IsSyntheticUserID(u) {
		return nil, nil
	}
	if amount <= 0 || amount > 1_000_000_000_000 {
		return nil, ErrInvalid
	}
	hold := &BalanceHold{ID: uuid.NewString(), TenantID: t, UserID: u, Amount: amount, Status: "reserved"}
	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		q := tx.Model(&Balance{}).Where("tenant_id=? AND user_id=? AND balance-reserved>=?", t, u, amount).Updates(map[string]any{"reserved": gorm.Expr("reserved+?", amount), "updated_at": time.Now()})
		if q.Error != nil {
			return q.Error
		}
		if q.RowsAffected != 1 {
			return fmt.Errorf("%w: 余额不足以完成本次生成，请联系管理员补充余额或调低输出长度", ErrInvalid)
		}
		hold.CreatedAt = time.Now().UTC()
		hold.ExpiresAt = hold.CreatedAt.Add(30 * time.Minute)
		if e := checkQuotaReservation(tx, hold); e != nil {
			return e
		}
		return tx.Create(hold).Error
	})
	return hold, err
}

func (s *Service) settleBalance(ctx context.Context, hold *BalanceHold, usage Usage) error {
	if usage.CreatedAt.IsZero() {
		usage.CreatedAt = time.Now().UTC()
	} else {
		usage.CreatedAt = usage.CreatedAt.UTC()
	}
	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if hold != nil {
			var current BalanceHold
			if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND tenant_id=? AND user_id=?", hold.ID, hold.TenantID, hold.UserID).Take(&current).Error; e != nil {
				return e
			}
			if current.Status == "settled" {
				return nil
			}
			if current.Status != "reserved" {
				return ErrConflict
			}
			var account Balance
			if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("tenant_id=? AND user_id=?", hold.TenantID, hold.UserID).Take(&account).Error; e != nil {
				return e
			}
			if account.Reserved < current.Amount {
				return errors.New("余额预留记录不一致")
			}
			// Actual usage above the conservative reservation is retained as debt,
			// never silently dropped or charged to another user's account.
			if e := tx.Model(&account).Updates(map[string]any{"balance": account.Balance - usage.Points, "reserved": account.Reserved - current.Amount, "updated_at": time.Now()}).Error; e != nil {
				return e
			}
			if usage.Points > 0 {
				row := Transaction{ID: uuid.NewString(), TenantID: hold.TenantID, UserID: hold.UserID, Delta: -usage.Points, Kind: "debit", Reason: usage.Action, IdempotencyKey: "call:" + hold.ID, Metadata: JSON{"call_id": usage.CallID, "job_id": usage.JobID, "model_id": usage.ModelID}, CreatedAt: time.Now()}
				if e := tx.Create(&row).Error; e != nil {
					return e
				}
			}
			if e := tx.Model(&current).Update("status", "settled").Error; e != nil {
				return e
			}
		}
		return tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&usage).Error
	})
}

func (s *Service) releaseExpiredBalances(ctx context.Context) error {
	var holds []BalanceHold
	if e := s.db.WithContext(ctx).Where("status=? AND expires_at<=?", "reserved", time.Now().UTC()).Limit(100).Find(&holds).Error; e != nil {
		return e
	}
	for _, hold := range holds {
		if e := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
			var current BalanceHold
			if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", hold.ID).Take(&current).Error; e != nil {
				return e
			}
			if current.Status != "reserved" || current.ExpiresAt.After(time.Now()) {
				return nil
			}
			q := tx.Model(&Balance{}).Where("tenant_id=? AND user_id=? AND reserved>=?", current.TenantID, current.UserID, current.Amount).Updates(map[string]any{"reserved": gorm.Expr("reserved-?", current.Amount), "updated_at": time.Now()})
			if q.Error != nil {
				return q.Error
			}
			if q.RowsAffected != 1 {
				return ErrConflict
			}
			return tx.Model(&current).Update("status", "expired").Error
		}); e != nil {
			return e
		}
	}
	return nil
}

type meteredCall struct {
	hold      *BalanceHold
	estimator *token.Estimator
	input     int
	usage     Usage
	started   time.Time
	budgetMS  int64
	purpose   string
}

func (m *meteredChat) start(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*meteredCall, error) {
	estimator, e := token.NewEstimator()
	if e != nil {
		return nil, e
	}
	input := estimator.EstimateMessages(messages)
	output := 8192
	if options != nil {
		input += estimator.EstimateTools(options.Tools)
		if options.MaxCompletionTokens > 0 {
			output = options.MaxCompletionTokens
		} else if options.MaxTokens > 0 {
			output = options.MaxTokens
		}
	}
	if input > 10_000_000 || output > 10_000_000 {
		return nil, ErrInvalid
	}
	hold, e := m.service.reserveBalance(ctx, int64(input+input/5+output)*pointsPerToken)
	if e != nil {
		return nil, e
	}
	t, _ := types.TenantIDFromContext(ctx)
	u, _ := types.UserIDFromContext(ctx)
	callID := uuid.NewString()
	started := time.Now().UTC()
	if hold != nil {
		callID = hold.ID
		started = hold.CreatedAt
	}
	a, _ := ctx.Value(meteringContextKey{}).(meteringAttribution)
	purpose, _ := types.LLMCallMetadataFromContext(ctx)
	if a.Action == "" {
		a.Action = actionName(purpose)
	}
	if a.Action == "" {
		a.Action = "AI 生成"
	}
	call := &meteredCall{hold: hold, estimator: estimator, input: input, started: time.Now(), purpose: purpose, usage: Usage{ID: uuid.NewString(), CallID: callID, TenantID: t, OwnerID: u, JobID: a.JobID, AgentID: a.AgentID, ModelID: m.GetModelID(), Action: a.Action, Status: "succeeded", Metadata: JSON{}, CreatedAt: started}}
	if deadline, ok := ctx.Deadline(); ok {
		call.budgetMS = max(0, time.Until(deadline).Milliseconds())
	}
	return call, nil
}

func (m *meteredChat) finish(ctx context.Context, call *meteredCall, reported types.TokenUsage, output string, failure error) error {
	call.usage.InputTokens = int64(max(0, reported.PromptTokens))
	call.usage.OutputTokens = int64(max(0, reported.CompletionTokens))
	measured := reported.TotalTokens > 0 || reported.PromptTokens > 0 || reported.CompletionTokens > 0
	if !measured && (failure == nil || output != "") {
		call.usage.InputTokens = int64(call.input)
		call.usage.OutputTokens = int64(call.estimator.EstimateString(output))
	}
	if measured && int64(reported.TotalTokens) > call.usage.InputTokens+call.usage.OutputTokens {
		call.usage.OutputTokens = int64(reported.TotalTokens) - call.usage.InputTokens
	}
	call.usage.Metadata = JSON{"estimated": !measured, "model_name": m.GetModelName(), "elapsed_ms": time.Since(call.started).Milliseconds(), "timeout_budget_ms": call.budgetMS, "purpose": call.purpose}
	failure = phaseFailure(failure, "model", call.started)
	if failure != nil {
		call.usage.Status = "failed"
		category, message := jobErrorDetails(failure)
		call.usage.Metadata["error_class"], call.usage.Metadata["error_message"] = category, message
	}
	if call.hold != nil {
		call.usage.Points = (call.usage.InputTokens + call.usage.OutputTokens) * pointsPerToken
	}
	finish, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
	defer cancel()
	if e := m.service.settleBalance(finish, call.hold, call.usage); e != nil {
		return fmt.Errorf("模型用量结算失败，请联系管理员核对调用 %s: %w", call.usage.CallID, e)
	}
	if failure != nil {
		_, message := jobErrorDetails(failure)
		diagnostic := jobErrorMetadata(failure)
		diagnostic["model_id"], diagnostic["agent_id"], diagnostic["purpose"] = m.GetModelID(), call.usage.AgentID, call.purpose
		diagnostic["elapsed_ms"], diagnostic["timeout_budget_ms"] = call.usage.Metadata["elapsed_ms"], call.budgetMS
		if e := m.service.recordOperation(finish, OperationLog{Kind: "system", Level: "error", Action: call.usage.Action, Message: message, TraceID: call.usage.CallID, ResourceID: call.usage.JobID, Metadata: diagnostic}); e != nil {
			logJobFailure(finish, nil, "model_error_log", e)
		}
	}
	return nil
}

func (m *meteredChat) Chat(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, 25*time.Minute)
	defer cancel()
	call, e := m.start(ctx, messages, options)
	if e != nil {
		return nil, e
	}
	response, err := m.model.Chat(ctx, messages, options)
	err = phaseFailure(err, "model", call.started)
	usage := types.TokenUsage{}
	output := ""
	if response != nil {
		usage = response.Usage
		output = response.Content + response.ReasoningContent
		if len(response.ToolCalls) > 0 {
			output += compact(response.ToolCalls)
		}
	}
	if settleErr := m.finish(ctx, call, usage, output, err); settleErr != nil {
		return nil, settleErr
	}
	return response, err
}

func (m *meteredChat) ChatStream(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (<-chan types.StreamResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, 25*time.Minute)
	call, e := m.start(ctx, messages, options)
	if e != nil {
		cancel()
		return nil, e
	}
	upstream, err := m.model.ChatStream(ctx, messages, options)
	err = phaseFailure(err, "model", call.started)
	if upstream == nil && err == nil {
		err = errors.New("AI 返回了空的流式响应")
	}
	if err != nil {
		defer cancel()
		if e = m.finish(ctx, call, types.TokenUsage{}, "", err); e != nil {
			return nil, e
		}
		return nil, err
	}
	out := make(chan types.StreamResponse)
	go func() {
		defer close(out)
		defer cancel()
		var usage types.TokenUsage
		var output strings.Builder
		var failure error
		var final *types.StreamResponse
		active := true
	stream:
		for {
			var chunk types.StreamResponse
			select {
			case <-ctx.Done():
				failure = ctx.Err()
				active = false
				break stream
			case next, ok := <-upstream:
				if !ok {
					break stream
				}
				chunk = next
			}
			if chunk.Usage != nil {
				usage = *chunk.Usage
			}
			output.WriteString(chunk.Content)
			if chunk.FinishReason == types.FinishReasonIncomplete || chunk.ResponseType == types.ResponseTypeError {
				failure = errors.New("AI 响应中断")
			}
			if chunk.Done {
				copy := chunk
				final = &copy
				continue
			}
			if active {
				select {
				case out <- chunk:
				case <-ctx.Done():
					active = false
					failure = ctx.Err()
				}
			}
		}
		if final == nil && failure == nil {
			failure = errors.New("AI 未返回完成标记")
		}
		failure = phaseFailure(failure, "model", call.started)
		if e := m.finish(ctx, call, usage, output.String(), failure); e != nil {
			failure = e
		}
		if final == nil {
			final = &types.StreamResponse{Done: true}
		}
		if failure != nil {
			final.FinishReason = types.FinishReasonIncomplete
			final.Data = map[string]any{"error": publicJobError(failure)}
		}
		if active {
			select {
			case out <- *final:
			case <-ctx.Done():
			}
		}
	}()
	return out, nil
}