/*
 * Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * openFuyao is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

package tokenizer

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/go-logr/logr"
	"sigs.k8s.io/controller-runtime/pkg/log"
	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/common/observability/logging"
	fwkplugin "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/plugin"
	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/requestcontrol"
	fwkscheduling "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/scheduling"

	tokenizedpromptattr "hermes-router/pkg/epp/framework/plugins/datalayer/attribute/tokenizedprompt"
	"hermes-router/pkg/epp/internal/utils"
)

const (
	// PluginType is the registered plugin type identifier.
	PluginType                 = "tokenizer"
	defaultBackend             = "hf"
	defaultTimeout             = 5 * time.Second
	defaultMaxMessageSize      = 4 * 1024 * 1024
	defaultInitRetries         = 1
	prepareRequestWaitAttempts = 3
	prepareRequestWaitInterval = 100 * time.Millisecond
	// backgroundInitInitialRetryInterval is the first attempt's time budget;
	// it doubles after each failure up to backgroundInitMaxRetryInterval.
	backgroundInitInitialRetryInterval = 1 * time.Second
	backgroundInitMaxRetryInterval     = 30 * time.Second
	// backgroundInitTotalRetryWindow bounds the total retry time so a
	// permanently broken sidecar cannot consume unlimited compute.
	backgroundInitTotalRetryWindow = 1 * time.Minute
)

type initState string

const (
	initStateInitializing initState = "initializing"
	initStateReady        initState = "ready"
	initStateFailed       initState = "failed"
)

type tokenizerInitializer func(context.Context, GRPCTokenizerConfig) (Backend, error)
type sleepFunc func(context.Context, time.Duration) error

func defaultTokenizerInitializer(ctx context.Context, cfg GRPCTokenizerConfig) (Backend, error) {
	return NewGRPCTokenizer(ctx, cfg)
}

type pluginParameters struct {
	Model      string `json:"model"`
	SocketPath string `json:"socketPath"`
	Backend    string `json:"backend"`
	// tokenizerSource overrides where the tokenizer is loaded from when it
	// differs from the served model name (e.g. a quantized model served as
	// "glm-51" whose tokenizer lives at a custom path or repo id). Empty falls
	// back to model. Only the operator-facing json key is renamed; the struct
	// field name stays TokenizerName to minimize churn (see design spec, Scope A).
	TokenizerName  string         `json:"tokenizerSource"`
	Timeout        utils.Duration `json:"timeout"`
	MaxMessageSize int            `json:"maxMessageSize"`
	InitRetries    int            `json:"initRetries"`
}

// Plugin populates request.TokenizedPrompt before scheduling.
type Plugin struct {
	typedName       fwkplugin.TypedName
	tokenizer       Backend
	backend         string
	model           string
	initCfg         GRPCTokenizerConfig
	initCtx         context.Context
	initMu          sync.RWMutex
	state           initState
	lastInitErr     error
	newTokenizer    tokenizerInitializer
	sleep           sleepFunc
	initRetryWindow time.Duration
}

var _ requestcontrol.PrepareDataPlugin = (*Plugin)(nil)

// Factory builds a tokenizer plugin and validates its config without blocking EPP startup.
func Factory(name string, parameters json.RawMessage, handle fwkplugin.Handle) (fwkplugin.Plugin, error) {
	params, err := parseAndValidateParams(parameters)
	if err != nil {
		return nil, fmt.Errorf("tokenizer plugin %q: %w", name, err)
	}

	plugin := &Plugin{
		typedName:       fwkplugin.TypedName{Type: PluginType, Name: name},
		backend:         params.Backend,
		model:           params.Model,
		initCtx:         initContext(handle),
		state:           initStateInitializing,
		newTokenizer:    defaultTokenizerInitializer,
		sleep:           waitWithContext,
		initRetryWindow: backgroundInitTotalRetryWindow,
		initCfg: GRPCTokenizerConfig{
			SocketPath:     params.SocketPath,
			Model:          params.Model,
			TokenizerName:  params.TokenizerName,
			Timeout:        params.Timeout.Std(),
			MaxMessageSize: params.MaxMessageSize,
			InitRetries:    params.InitRetries,
		},
	}
	plugin.registerShutdownClose()
	go plugin.runInitLoop()
	return plugin, nil
}

func initContext(handle fwkplugin.Handle) context.Context {
	if handle == nil || handle.Context() == nil {
		return context.Background()
	}
	return handle.Context()
}

func parseAndValidateParams(raw []byte) (pluginParameters, error) {
	if len(raw) == 0 {
		return pluginParameters{}, fmt.Errorf("parameters are required")
	}

	params := pluginParameters{}
	if err := json.Unmarshal(raw, &params); err != nil {
		return pluginParameters{}, fmt.Errorf("parse parameters: %w", err)
	}
	if params.Model == "" {
		return pluginParameters{}, fmt.Errorf("model is required")
	}
	if params.SocketPath == "" {
		return pluginParameters{}, fmt.Errorf("socketPath is required")
	}
	if params.Backend == "" {
		params.Backend = defaultBackend
	}
	if params.Backend != defaultBackend && params.Backend != "vllm" {
		return pluginParameters{}, fmt.Errorf("backend must be 'hf' or 'vllm', got %q", params.Backend)
	}
	if params.Timeout.Std() <= 0 {
		params.Timeout = utils.Duration(defaultTimeout)
	}
	if params.MaxMessageSize <= 0 {
		params.MaxMessageSize = defaultMaxMessageSize
	}
	if params.InitRetries <= 0 {
		params.InitRetries = defaultInitRetries
	}

	return params, nil
}

// TypedName returns the plugin identity.
func (p *Plugin) TypedName() fwkplugin.TypedName {
	return p.typedName
}

// Produces declares the tokenized prompt contract while the runtime payload remains on LLMRequest.
func (p *Plugin) Produces() map[string]any {
	return map[string]any{
		tokenizedpromptattr.TokenizedPromptKey: tokenizedpromptattr.Info{},
	}
}

// Consumes declares no framework-managed upstream keys.
func (p *Plugin) Consumes() map[string]any {
	return map[string]any{}
}

// PrepareRequestData tokenizes supported request shapes and fails open on runtime errors.
func (p *Plugin) PrepareRequestData(
	ctx context.Context,
	request *fwkscheduling.LLMRequest,
	_ []fwkscheduling.Endpoint,
) error {
	if ctx == nil {
		ctx = context.Background()
	}

	logger := log.FromContext(ctx).WithName("Plugin.PrepareRequestData")
	start := time.Now()
	if request == nil || request.Body == nil {
		logger.V(logutil.DEBUG).Info(
			"Tokenizer skipped tokenization",
			"reason", "missing_request_body",
			"durationMs", utils.DurationMillisecondsSince(start),
		)
		return nil
	}
	if request.TokenizedPrompt != nil {
		logger.V(logutil.DEBUG).Info(
			"Tokenizer skipped tokenization",
			"reason", "already_tokenized",
			"durationMs", utils.DurationMillisecondsSince(start),
		)
		return nil
	}

	tokenizer, state, err := p.snapshotInitState()
	if state != initStateReady || tokenizer == nil {
		tokenizer, state, err = p.waitForReadyTokenizer(ctx)
	}
	if state != initStateReady || tokenizer == nil {
		logger.Error(
			unavailableError(state, err),
			"Tokenizer unavailable after bounded readiness wait; continuing without tokenized prompt",
			"requestId", request.RequestId,
			"state", state,
			"requestKind", requestKind(request),
			"model", p.requestModel(request),
			"backend", p.backend,
			"durationMs", utils.DurationMillisecondsSince(start),
		)
		return nil
	}

	result, err := p.prepareTokenization(ctx, tokenizer, request)
	if err != nil {
		logger.Error(
			err,
			"Tokenizer request tokenization failed; continuing without tokenized prompt",
			"requestId", request.RequestId,
			"requestKind", requestKind(request),
			"model", p.requestModel(request),
			"backend", p.backend,
			"durationMs", utils.DurationMillisecondsSince(start),
		)
		return nil
	}

	request.TokenizedPrompt = &fwkscheduling.TokenizedPrompt{
		TokenIDs:           result.TokenIDs,
		MultiModalFeatures: result.MultiModalFeatures,
	}
	logger.V(logutil.DEBUG).Info(
		"Tokenizer generated tokenized prompt",
		"model", p.requestModel(request),
		"tokenCount", len(result.TokenIDs),
		"multimodalFeatureCount", len(result.MultiModalFeatures),
		"durationMs", utils.DurationMillisecondsSince(start),
	)
	return nil
}

func (p *Plugin) snapshotInitState() (Backend, initState, error) {
	if p == nil {
		return nil, initStateFailed, fmt.Errorf("tokenizer plugin is nil")
	}

	p.initMu.RLock()
	defer p.initMu.RUnlock()
	return p.tokenizer, p.state, p.lastInitErr
}

func waitWithContext(ctx context.Context, delay time.Duration) error {
	if ctx == nil {
		ctx = context.Background()
	}

	timer := time.NewTimer(delay)
	defer timer.Stop()

	select {
	case <-timer.C:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (p *Plugin) waitForReadyTokenizer(ctx context.Context) (Backend, initState, error) {
	for attempt := 0; attempt < prepareRequestWaitAttempts; attempt++ {
		tokenizer, state, err := p.snapshotInitState()
		if state == initStateReady && tokenizer != nil {
			return tokenizer, state, nil
		}
		if attempt == prepareRequestWaitAttempts-1 {
			return nil, state, err
		}
		if err := waitWithContext(ctx, prepareRequestWaitInterval); err != nil {
			return nil, state, err
		}
	}

	return nil, initStateFailed, nil
}

func unavailableError(state initState, err error) error {
	if err != nil {
		return err
	}
	return fmt.Errorf("tokenizer state %s", state)
}

func requestKind(request *fwkscheduling.LLMRequest) string {
	if request != nil && request.Body != nil && request.Body.ChatCompletions != nil {
		return "chat"
	}
	return "completion"
}

func (p *Plugin) initLogger() logr.Logger {
	if p == nil || p.initCtx == nil {
		return logr.Discard()
	}
	return log.FromContext(p.initCtx).WithName("Plugin.Init").WithValues(
		"plugin", p.typedName.Name,
		"model", p.model,
		"backend", p.backend,
		"socketPath", p.initCfg.SocketPath,
	)
}

func (p *Plugin) storeReadyTokenizer(tokenizer Backend) {
	p.initMu.Lock()
	defer p.initMu.Unlock()
	p.tokenizer = tokenizer
	p.state = initStateReady
	p.lastInitErr = nil
}

func (p *Plugin) storeInitFailure(err error) {
	p.initMu.Lock()
	defer p.initMu.Unlock()
	p.tokenizer = nil
	p.state = initStateFailed
	p.lastInitErr = err
}

func (p *Plugin) storeInitializing() {
	p.initMu.Lock()
	defer p.initMu.Unlock()
	p.state = initStateInitializing
}

func (p *Plugin) runInitLoop() {
	if p == nil {
		return
	}

	logger := p.initLogger()
	window := p.initRetryWindow
	if window <= 0 {
		window = backgroundInitTotalRetryWindow
	}
	start := time.Now()
	attempt := 0
	attemptWindow := backgroundInitInitialRetryInterval
	for {
		attempt++
		logger.Info("Tokenizer initialization attempt started", "attempt", attempt)
		attemptCtx, cancel := context.WithTimeout(p.initCtx, attemptWindow)
		attemptStart := time.Now()
		tokenizer, err := p.newTokenizer(attemptCtx, p.initCfg)
		cancel()
		if err == nil {
			p.storeReadyTokenizer(tokenizer)
			logger.Info("Tokenizer initialization succeeded", "attempt", attempt)
			return
		}

		// The attempt runs inside its budget window and is declared failed
		// only when the window ends without success, so even instant failures
		// (e.g. the sidecar socket is not up) wait out the backoff window.
		attemptDuration := time.Since(attemptStart)
		if remaining := attemptWindow - attemptDuration; remaining > 0 {
			if err := p.sleep(p.initCtx, remaining); err != nil {
				return
			}
		}

		p.storeInitFailure(err)
		elapsed := time.Since(start)
		if elapsed >= window {
			logger.Error(
				err,
				"Tokenizer initialization failed; giving up after retry window",
				"attempt", attempt,
				"elapsed", elapsed,
			)
			return
		}
		logger.Error(
			err,
			"Tokenizer initialization failed; retrying",
			"attempt", attempt,
			"attemptWindow", attemptWindow,
			"elapsed", elapsed,
		)
		p.storeInitializing()
		attemptWindow = nextBackoff(attemptWindow)
	}
}

// nextBackoff doubles the retry delay until backgroundInitMaxRetryInterval caps it.
func nextBackoff(delay time.Duration) time.Duration {
	if delay >= backgroundInitMaxRetryInterval {
		return backgroundInitMaxRetryInterval
	}
	doubled := delay * 2
	if doubled > backgroundInitMaxRetryInterval {
		return backgroundInitMaxRetryInterval
	}
	return doubled
}

func (p *Plugin) registerShutdownClose() {
	if p == nil || p.initCtx == nil {
		return
	}

	go func() {
		<-p.initCtx.Done()
		tokenizer, _, err := p.snapshotInitState()
		if err != nil {
			p.initLogger().Error(err, "Tokenizer shutdown snapshot failed")
			return
		}
		if tokenizer == nil {
			return
		}
		if err := tokenizer.Close(); err != nil {
			p.initLogger().Error(err, "Tokenizer shutdown close failed")
		}
	}()
}

func (p *Plugin) prepareTokenization(
	ctx context.Context,
	tokenizer Backend,
	request *fwkscheduling.LLMRequest,
) (*TokenizationResult, error) {
	if request.Body.Completions != nil {
		return p.tokenizeCompletion(ctx, tokenizer, request)
	}
	if request.Body.ChatCompletions != nil {
		return p.tokenizeChatCompletion(ctx, tokenizer, request)
	}
	return nil, nil
}

func (p *Plugin) tokenizeCompletion(
	ctx context.Context,
	tokenizer Backend,
	request *fwkscheduling.LLMRequest,
) (*TokenizationResult, error) {
	completions := request.Body.Completions
	if completions == nil || completions.Prompt.IsEmpty() {
		return nil, nil
	}

	if p.backend == defaultBackend && canUseDirectTokenize(completions) {
		prompt := completions.Prompt.PlainText()
		if prompt == "" {
			return nil, nil
		}
		return tokenizer.Tokenize(ctx, Request{
			Model:  p.requestModel(request),
			Prompt: prompt,
		})
	}

	return tokenizer.Complete(ctx, CompleteRequest{
		Model:      p.requestModel(request),
		PromptText: completions.Prompt.PlainText(),
	})
}

func canUseDirectTokenize(request *fwkscheduling.CompletionsRequest) bool {
	return request != nil && request.Prompt.Raw != ""
}

func (p *Plugin) tokenizeChatCompletion(
	ctx context.Context,
	tokenizer Backend,
	request *fwkscheduling.LLMRequest,
) (*TokenizationResult, error) {
	chat := request.Body.ChatCompletions
	if chat == nil || len(chat.Messages) == 0 {
		return nil, nil
	}

	return tokenizer.ChatComplete(ctx, ChatCompleteRequest{
		Model:                  p.requestModel(request),
		Messages:               convertMessages(chat.Messages),
		ChatTemplate:           chat.ChatTemplate,
		ChatTemplateKWargsJSON: marshalJSONString(chat.ChatTemplateKWArgs),
		ToolsJSON:              marshalJSONString(chat.Tools),
		AddGenerationPrompt:    chat.AddGenerationPrompt,
		ContinueFinalMessage:   chat.ContinueFinalMessage,
	})
}

func (p *Plugin) requestModel(request *fwkscheduling.LLMRequest) string {
	if request != nil && request.TargetModel != "" {
		return request.TargetModel
	}
	return p.model
}

func convertMessages(messages []fwkscheduling.Message) []ChatMessage {
	converted := make([]ChatMessage, 0, len(messages))
	for _, message := range messages {
		converted = append(converted, ChatMessage{
			Role:         message.Role,
			Content:      message.Content.Raw,
			ContentParts: convertContentParts(message.Content.Structured),
		})
	}
	return converted
}

func convertContentParts(blocks []fwkscheduling.ContentBlock) []ChatContentPart {
	parts := make([]ChatContentPart, 0, len(blocks))
	for _, block := range blocks {
		switch block.Type {
		case "text":
			parts = append(parts, ChatContentPart{Text: block.Text})
		case "image_url":
			parts = append(parts, ChatContentPart{Media: &MediaRef{
				Modality: string(fwkscheduling.ModalityImage),
				URL:      block.ImageURL.Url,
			}})
		default:
			continue
		}
	}
	return parts
}

func marshalJSONString(value any) string {
	if value == nil {
		return ""
	}
	b, err := json.Marshal(value)
	if err != nil {
		return ""
	}
	return string(b)
}