/*
 * 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 prerequest provides pre-request plugins for GIE.
package prerequest

import (
	"context"
	"encoding/json"
	"fmt"
	"net"

	"sigs.k8s.io/controller-runtime/pkg/log"
	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/common/observability/logging"
	fwkdl "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/datalayer"
	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"

	"hermes-router/pkg/epp/internal/endpointmeta"
	"hermes-router/pkg/epp/internal/pdgroup"
	"hermes-router/pkg/epp/internal/utils"
)

const (
	// HeaderHandlerType is the plugin type identifier for the PD header handler.
	HeaderHandlerType = "pd-header-handler"
)

var _ requestcontrol.PreRequest = &HeaderHandler{}

// HeaderHandlerFactory creates a plugin instance from EPP configuration.
func HeaderHandlerFactory(
	name string,
	_ json.RawMessage,
	_ fwkplugin.Handle,
) (fwkplugin.Plugin, error) {
	return NewHeaderHandler().WithName(name), nil
}

// NewHeaderHandler initializes a new HeaderHandler.
func NewHeaderHandler() *HeaderHandler {
	return &HeaderHandler{
		typedName: fwkplugin.TypedName{Type: HeaderHandlerType},
	}
}

// HeaderHandler injects PD route headers during the pre-request hook.
type HeaderHandler struct {
	typedName fwkplugin.TypedName
}

// TypedName returns the plugin's typed name.
func (handler *HeaderHandler) TypedName() fwkplugin.TypedName {
	return handler.typedName
}

// WithName sets the plugin's configured name.
func (handler *HeaderHandler) WithName(name string) *HeaderHandler {
	handler.typedName.Name = name
	return handler
}

// Route describes the leader and the selected prefill/decode members resolved from the primary profile result.
type Route struct {
	Leader  *fwkdl.EndpointMetadata
	Prefill *fwkdl.EndpointMetadata
	Decode  *fwkdl.EndpointMetadata
}

// ExtractPrimaryRoute resolves the selected leader endpoint into a PD route.
func ExtractPrimaryRoute(ctx context.Context, schedulingResult *fwkscheduling.SchedulingResult) *Route {
	logger := log.FromContext(ctx)
	if schedulingResult == nil {
		return nil
	}

	primaryProfileResult := schedulingResult.ProfileResults[schedulingResult.PrimaryProfileName]
	if primaryProfileResult == nil ||
		primaryProfileResult.TargetEndpoints == nil ||
		len(primaryProfileResult.TargetEndpoints) == 0 {
		logger.Error(
			fmt.Errorf("nil target endpoint in primaryProfileResult"),
			"primaryProfileResult", primaryProfileResult,
		)
		return nil
	}

	logger.V(logutil.DEBUG).Info("prerequest_pd: Read scheduling result successfully")
	leader := primaryLeaderEndpoint(primaryProfileResult)
	if route := selectedLeaderRoute(leader); route != nil {
		return route
	}

	logger.V(logutil.DEFAULT).Info("primary scheduling result does not contain a complete PD route",
		"targets", len(primaryProfileResult.TargetEndpoints))
	return nil
}

func selectedLeaderRoute(leader fwkscheduling.Endpoint) *Route {
	if leader == nil {
		return nil
	}
	info, ok := pdgroup.GetDirect(leader)
	if !ok || info == nil {
		return nil
	}
	prefillEP, decodeEP, ok := selectedPDRouteEndpoints(info)
	if !ok {
		return nil
	}
	prefill := endpointmeta.MetadataOf(prefillEP)
	decode := endpointmeta.MetadataOf(decodeEP)
	if prefill == nil || decode == nil {
		return nil
	}
	return &Route{
		Leader:  endpointmeta.MetadataOf(leader),
		Prefill: prefill,
		Decode:  decode,
	}
}

func selectedPDRouteEndpoints(info *pdgroup.Info) (fwkscheduling.Endpoint, fwkscheduling.Endpoint, bool) {
	if info.SelectedPrefillPod == nil || info.SelectedDecodePod == nil {
		return nil, nil, false
	}
	prefillEP := info.SelectedPrefillPod.Endpoint
	decodeEP := info.SelectedDecodePod.Endpoint
	return prefillEP, decodeEP, prefillEP != nil && decodeEP != nil
}

func primaryLeaderEndpoint(primaryProfileResult *fwkscheduling.ProfileRunResult) fwkscheduling.Endpoint {
	if primaryProfileResult == nil || len(primaryProfileResult.TargetEndpoints) == 0 {
		return nil
	}
	return primaryProfileResult.TargetEndpoints[0]
}

// InjectHeadersFromRoute writes the prefill/decode target addresses into the request headers.
func InjectHeadersFromRoute(
	ctx context.Context,
	request *fwkscheduling.LLMRequest,
	route *Route,
) {
	if request == nil || route == nil || route.Prefill == nil || route.Decode == nil {
		return
	}

	logger := log.FromContext(ctx).V(logutil.DEBUG)
	prefillHostPort := net.JoinHostPort(route.Prefill.Address, route.Prefill.Port)
	decodeHostPort := net.JoinHostPort(route.Decode.Address, route.Decode.Port)

	if request.Headers == nil {
		request.Headers = make(map[string]string)
	}
	request.Headers[utils.PrefillPodHeader] = prefillHostPort
	request.Headers[utils.DecodePodHeader] = decodeHostPort

	logger.Info("injected PD prefill/decode headers",
		"prefill", prefillHostPort,
		"decode", decodeHostPort,
		"prefillPod", route.Prefill.Address,
		"decodePod", route.Decode.Address,
	)
}

// InjectHeaders extracts the primary route and writes the corresponding request headers.
func InjectHeaders(
	ctx context.Context,
	request *fwkscheduling.LLMRequest,
	schedulingResult *fwkscheduling.SchedulingResult,
) {
	route := ExtractPrimaryRoute(ctx, schedulingResult)
	InjectHeadersFromRoute(ctx, request, route)
}

// PreRequest injects PD routing headers before the request is forwarded.
func (handler *HeaderHandler) PreRequest(
	ctx context.Context,
	request *fwkscheduling.LLMRequest,
	schedulingResult *fwkscheduling.SchedulingResult,
) {
	InjectHeaders(ctx, request, schedulingResult)
}