/*
 * Copyright (c) 2026 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 main

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

	"many-core-orchestrator/api/v1alpha1"
	"many-core-orchestrator/pkg/constant"

	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/dynamic/dynamicinformer"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
	"k8s.io/klog/v2"
	"volcano.sh/volcano/pkg/scheduler/api"
	"volcano.sh/volcano/pkg/scheduler/framework"
)

const (
	// 注册在 volcano-scheduler-configmap 中的插件名称
	PluginName = "mco-plugin"

	// CR 数据过期的默认 TTL（秒）
	defaultCRStatusTTLSeconds = 30

	// 无可用的干扰数据时回退的中性分
	neutralScore = 50.0
)

var (
	cacheMu          sync.RWMutex
	cacheReports     = make(map[string]*v1alpha1.NodeInterferenceNodeReport)
	cacheLastRefresh time.Time
	cacheLastMaxSeq  int64
	informerOnce     sync.Once
	informerErr      error
)

// mcoPlugin 实现 framework.Plugin，提供干扰感知调度
type mcoPlugin struct {
	args *pluginArgs
}

// pluginArgs 对应 volcano-scheduler-configmap 中 mco-plugin 的 arguments 配置
type pluginArgs struct {
	// NodeInterferenceReport CR 名称（固定 "cluster"）
	CRName string `json:"crName"`
	// 条目过期时间（秒），默认 30
	CRStatusTTLSeconds int `json:"crStatusTTLSeconds"`
	// Filter 阶段硬约束阈值
	HardConstraints hardConstraints `json:"hardConstraints"`
	// CR 过期时是否降级
	FallbackWhenCROutdated bool `json:"fallbackWhenCROutdated"`
}

// hardConstraints 定义 Filter 阶段硬约束阈值
type hardConstraints struct {
	// psiMemoryFull avg10 超过此值则排除节点（默认 0.0）
	PsiMemoryFullAvg10 float64 `json:"psiMemoryFullAvg10"`
	// psiIoSome avg10 超过此值则排除节点（默认 20.0）
	PsiIoSomeAvg10Threshold float64 `json:"psiIoSomeAvg10Threshold"`
}

// ── 插件生命周期 ──

func (p *mcoPlugin) Name() string {
	return PluginName
}

// OnSessionOpen 注册 Filter 和 Score 回调
func (p *mcoPlugin) OnSessionOpen(ssn *framework.Session) {
	klog.V(4).Infof("[%s] session opened, registering filter/score", PluginName)

	p.ensureInformerStarted()

	// Filter：硬约束（predicate）
	ssn.AddPredicateFn(p.Name(), func(task *api.TaskInfo, node *api.NodeInfo) ([]*api.Status, error) {
		return nil, p.filterNode(task, node)
	})

	// Score：干扰等级评分
	ssn.AddNodeOrderFn(p.Name(), func(task *api.TaskInfo, node *api.NodeInfo) (float64, error) {
		return p.scoreNode(task, node)
	})

	// 批量评分
	ssn.AddBatchNodeOrderFn(p.Name(), func(task *api.TaskInfo, nodes []*api.NodeInfo) (map[string]float64, error) {
		scores := make(map[string]float64, len(nodes))
		for _, node := range nodes {
			s, err := p.scoreNode(task, node)
			if err != nil {
				klog.Warningf("[%s] batch score failed for node %s: %v", PluginName, node.Name, err)
				scores[node.Name] = neutralScore
				continue
			}
			scores[node.Name] = s
		}
		return scores, nil
	})
}

func (p *mcoPlugin) OnSessionClose(ssn *framework.Session) {
	klog.V(4).Infof("[%s] session closed", PluginName)
}

// ── Filter ──

// filterNode 执行硬约束过滤：Kata 节点检查 → PSI Memory Full → PSI I/O
func (p *mcoPlugin) filterNode(task *api.TaskInfo, node *api.NodeInfo) error {
	// Kata 运行时检查：需同时满足运行时已注册且 kata-check 验证通过
	if needsKata(task) {
		if node.Node.Labels[constant.KataRuntimeLabel] != "true" ||
			node.Node.Labels[constant.KataReadyLabel] != "true" {
			klog.V(5).Infof("[%s] filter: node %s lacks kata runtime or not kata-ready", PluginName, node.Name)
			return fmt.Errorf("node %s does not have kata runtime ready", node.Name)
		}
	}

	report, ok := p.getCachedReport(node.Name)
	if !ok {
		return nil // 无数据 → 不排除（无证据默认放行）
	}

	// PSI 不可用 → 跳过 PSI 约束
	if !report.PSIAvailable {
		klog.V(3).Infof("[%s] filter: PSI unavailable for node %s, skipping PSI constraints",
			PluginName, node.Name)
		return nil
	}

	psiMemFull := report.Metrics.PSIMemFullAvg10
	if psiMemFull > p.args.HardConstraints.PsiMemoryFullAvg10 {
		klog.V(3).Infof("[%s] filter: node %s excluded (psiMemoryFull=%.2f > %.2f)",
			PluginName, node.Name, psiMemFull, p.args.HardConstraints.PsiMemoryFullAvg10)
		return fmt.Errorf("node %s memory full pressure (psiMemoryFull=%.2f)", node.Name, psiMemFull)
	}

	psiIoSome := report.Metrics.PSIIOAvg10
	if psiIoSome > p.args.HardConstraints.PsiIoSomeAvg10Threshold {
		klog.V(3).Infof("[%s] filter: node %s excluded (psiIoSome=%.2f > %.2f)",
			PluginName, node.Name, psiIoSome, p.args.HardConstraints.PsiIoSomeAvg10Threshold)
		return fmt.Errorf("node %s I/O pressure (psiIoSome=%.2f)", node.Name, psiIoSome)
	}

	return nil
}

// ── Score ──

// scoreNode 根据干扰等级评分：score = 100 × (1 - interferenceLevel)
// 降级情况（节点不在 CR / 数据过期 / status≠ok）返回中性分 50
func (p *mcoPlugin) scoreNode(task *api.TaskInfo, node *api.NodeInfo) (float64, error) {
	_ = task

	report, ok := p.getCachedReport(node.Name)
	if !ok {
		klog.V(3).Infof("[%s] score: node %s not in CR or CR stale, returning neutral", PluginName, node.Name)
		return neutralScore, nil
	}

	if report.Status != "ok" {
		klog.V(3).Infof("[%s] score: node %s status=%q, using neutral score",
			PluginName, node.Name, report.Status)
		return neutralScore, nil
	}

	score := 100.0 * (1.0 - report.InterferenceLevel)
	if score < 0 {
		score = 0
	}
	if score > 100 {
		score = 100
	}

	klog.V(3).Infof("[%s] score: node %s status=%s interferenceLevel=%.3f → score=%.1f",
		PluginName, node.Name, report.Status, report.InterferenceLevel, score)

	return score, nil
}

// ── CR 缓存 ──

// getCachedReport 查询节点的缓存报告，TTL 过期或不存在时返回 false
func (p *mcoPlugin) getCachedReport(nodeName string) (*v1alpha1.NodeInterferenceNodeReport, bool) {
	cacheMu.RLock()
	defer cacheMu.RUnlock()

	ttl := time.Duration(p.args.CRStatusTTLSeconds) * time.Second
	if ttl <= 0 {
		ttl = defaultCRStatusTTLSeconds * time.Second
	}

	elapsed := time.Since(cacheLastRefresh)
	if elapsed > ttl {
		klog.Warningf("[%s] CR data stale (last update %v ago, TTL=%v), falling back to neutral score — check Analyzer health",
			PluginName, elapsed.Round(time.Second), ttl)
		return nil, false
	}

	report, ok := cacheReports[nodeName]
	if !ok {
		return nil, false
	}
	return report, true
}

// updateCache 替换本地缓存，仅在数据变化时刷新时间戳
func (p *mcoPlugin) updateCache(nir *v1alpha1.NodeInterferenceReport) {
	cacheMu.Lock()
	defer cacheMu.Unlock()

	var maxSeq int64
	newReports := make(map[string]*v1alpha1.NodeInterferenceNodeReport, len(nir.Status.Nodes))
	for nodeName, report := range nir.Status.Nodes {
		reportCopy := report
		newReports[nodeName] = &reportCopy
		if report.Sequence > maxSeq {
			maxSeq = report.Sequence
		}
	}

	if maxSeq != cacheLastMaxSeq {
		cacheLastRefresh = time.Now()
		klog.V(5).Infof("[%s] cache updated: %d nodes (maxSeq=%d)", PluginName, len(newReports), maxSeq)
		cacheLastMaxSeq = maxSeq
	}
	cacheReports = newReports
}

// clearCache 清空本地缓存（CR 被删除时触发）
func (p *mcoPlugin) clearCache() {
	cacheMu.Lock()
	defer cacheMu.Unlock()

	cacheReports = make(map[string]*v1alpha1.NodeInterferenceNodeReport)
	cacheLastRefresh = time.Time{}
	cacheLastMaxSeq = 0
	klog.V(4).Infof("[%s] cache cleared (CR deleted)", PluginName)
}

// ── CR Informer ──

// ensureInformerStarted 懒启动 NodeInterferenceReport CR 的 informer，可多次调用
func (p *mcoPlugin) ensureInformerStarted() {
	informerOnce.Do(func() {
		stopCh := make(chan struct{})
		informerErr = p.startInformer(stopCh)
		if informerErr != nil {
			klog.Errorf("[%s] CR informer failed to start: %v — scheduling will use fallback", PluginName, informerErr)
		}
	})
}

// startInformer 使用 dynamic client 监听 cluster 范围的 NodeInterferenceReport CR
func (p *mcoPlugin) startInformer(stopCh <-chan struct{}) error {
	config, err := rest.InClusterConfig()
	if err != nil {
		return fmt.Errorf("in-cluster config: %w", err)
	}

	dynClient, err := dynamic.NewForConfig(config)
	if err != nil {
		return fmt.Errorf("dynamic client: %w", err)
	}

	gvr := schema.GroupVersionResource{
		Group:    v1alpha1.GroupVersion.Group,
		Version:  v1alpha1.GroupVersion.Version,
		Resource: "nodeinterferencereports",
	}

	factory := dynamicinformer.NewFilteredDynamicSharedInformerFactory(dynClient, 5*time.Second, "", nil)
	informer := factory.ForResource(gvr).Informer()

	informer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			nir := p.unstructuredToReport(obj)
			if nir == nil || nir.Name != "cluster" {
				return
			}
			klog.V(4).Infof("[%s] CR 'cluster' created, updating cache", PluginName)
			p.updateCache(nir)
		},
		UpdateFunc: func(_, obj interface{}) {
			nir := p.unstructuredToReport(obj)
			if nir == nil || nir.Name != "cluster" {
				return
			}
			klog.V(6).Infof("[%s] CR 'cluster' updated, refreshing cache", PluginName)
			p.updateCache(nir)
		},
		DeleteFunc: func(obj interface{}) {
			if tombstone, ok := obj.(cache.DeletedFinalStateUnknown); ok {
				obj = tombstone.Obj
			}
			nir := p.unstructuredToReport(obj)
			if nir == nil || nir.Name != "cluster" {
				return
			}
			klog.V(4).Infof("[%s] CR 'cluster' deleted, clearing cache", PluginName)
			p.clearCache()
		},
	})

	go informer.Run(stopCh)

	klog.Infof("[%s] waiting for CR informer cache sync...", PluginName)
	if !cache.WaitForCacheSync(stopCh, informer.HasSynced) {
		return fmt.Errorf("timed out waiting for CR informer cache sync")
	}

	klog.Infof("[%s] CR informer synced successfully", PluginName)
	return nil
}

// unstructuredToReport 将 dynamic informer 的 Unstructured 对象转为 NodeInterferenceReport
func (p *mcoPlugin) unstructuredToReport(obj interface{}) *v1alpha1.NodeInterferenceReport {
	u, ok := obj.(*unstructured.Unstructured)
	if !ok {
		return nil
	}
	nir := &v1alpha1.NodeInterferenceReport{}
	if err := runtime.DefaultUnstructuredConverter.FromUnstructured(u.Object, nir); err != nil {
		klog.Warningf("[%s] failed to convert CR: %v", PluginName, err)
		return nil
	}
	return nir
}

// ── 辅助函数 ──

// needsKata 判断 Pod 是否要求 Kata 运行时（runtimeClassName 以 kata- 开头）
func needsKata(task *api.TaskInfo) bool {
	if task == nil || task.Pod == nil {
		return false
	}
	rc := task.Pod.Spec.RuntimeClassName
	if rc == nil {
		return false
	}
	return len(*rc) >= 5 && (*rc)[:5] == "kata-"
}

// parseArgs 解析 Volcano ConfigMap 中 mco-plugin 的参数，处理 YAML map[interface{}] 类型
func parseArgs(arguments framework.Arguments) (*pluginArgs, error) {
	args := &pluginArgs{
		CRName:                 "cluster",
		CRStatusTTLSeconds:     defaultCRStatusTTLSeconds,
		FallbackWhenCROutdated: true,
		HardConstraints: hardConstraints{
			PsiMemoryFullAvg10:      0.0,
			PsiIoSomeAvg10Threshold: 20.0,
		},
	}

	raw, ok := arguments[PluginName]
	if !ok || raw == nil {
		klog.Infof("[%s] no plugin arguments provided, using defaults", PluginName)
		return args, nil
	}

	// YAML 解析出的 map[interface{}]interface{} 不能直接 JSON 序列化，先递归转换
	normalized := normalizeMap(raw)
	b, err := json.Marshal(normalized)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal plugin arguments: %w", err)
	}
	if err := json.Unmarshal(b, args); err != nil {
		return nil, fmt.Errorf("failed to parse plugin arguments: %w", err)
	}
	return args, nil
}

// normalizeMap 递归转换 map[interface{}]interface{} → map[string]interface{}
func normalizeMap(v interface{}) interface{} {
	switch val := v.(type) {
	case map[interface{}]interface{}:
		out := make(map[string]interface{}, len(val))
		for k, vv := range val {
			out[fmt.Sprint(k)] = normalizeMap(vv)
		}
		return out
	case map[string]interface{}:
		out := make(map[string]interface{}, len(val))
		for k, vv := range val {
			out[k] = normalizeMap(vv)
		}
		return out
	case []interface{}:
		out := make([]interface{}, len(val))
		for i, vv := range val {
			out[i] = normalizeMap(vv)
		}
		return out
	default:
		return v
	}
}

// ── 插件入口（导出） ──

// New 是 Volcano 插件构建函数，框架通过 .so 动态加载调用此函数
func New(arguments framework.Arguments) framework.Plugin {
	klog.Infof("[%s] initializing MCO interference-aware scheduling plugin (dynamic .so)", PluginName)

	args, err := parseArgs(arguments)
	if err != nil {
		klog.Errorf("[%s] failed to parse arguments: %v, using defaults", PluginName, err)
		args = &pluginArgs{
			CRName:                 "cluster",
			CRStatusTTLSeconds:     defaultCRStatusTTLSeconds,
			FallbackWhenCROutdated: true,
			HardConstraints: hardConstraints{
				PsiMemoryFullAvg10:      0.0,
				PsiIoSomeAvg10Threshold: 20.0,
			},
		}
	}

	return &mcoPlugin{
		args: args,
	}
}

// main 是 .so 插件的空入口，Volcano 框架通过 .so 加载后直接调用导出的 New()，无需在此执行逻辑
func main() {
	// .so 插件无需在 main 中启动任何逻辑
}