package dataplane

import (
	"testing"

	pipetype "github.com/tidbcloud/metrics-integration/global/pkg/helm/pipeline"
	"sigs.k8s.io/yaml"
)

func TestPatchDedicatedClusterCPUConfigYAML(t *testing.T) {
	original := `
server:
  logLevel: info
o11yGlobalServer:
  url: https://example.com
targetClusters:
  - id: 123
    name: demo
    components:
      - componentType: COMPONENT_TYPE_TIDB
        cpuNum: 2
      - componentType: COMPONENT_TYPE_TIKV
        cpuNum: 4
`

	patched, err := PatchDedicatedClusterCPUConfigYAML(original, []pipetype.ClusterComponent{
		{ComponentType: "COMPONENT_TYPE_TIDB", CpuNum: 8},
		{ComponentType: "COMPONENT_TYPE_TIKV", CpuNum: 16},
	})
	if err != nil {
		t.Fatalf("PatchDedicatedClusterCPUConfigYAML() error = %v", err)
	}

	var cfg map[string]interface{}
	if err := yaml.Unmarshal([]byte(patched), &cfg); err != nil {
		t.Fatalf("yaml.Unmarshal() error = %v", err)
	}

	targetClusters := cfg["targetClusters"].([]interface{})
	cluster := targetClusters[0].(map[string]interface{})
	components := cluster["components"].([]interface{})

	got := map[string]float64{}
	for _, item := range components {
		component := item.(map[string]interface{})
		got[component["componentType"].(string)] = component["cpuNum"].(float64)
	}

	if got["COMPONENT_TYPE_TIDB"] != 8 {
		t.Fatalf("TIDB cpuNum = %v, want 8", got["COMPONENT_TYPE_TIDB"])
	}
	if got["COMPONENT_TYPE_TIKV"] != 16 {
		t.Fatalf("TIKV cpuNum = %v, want 16", got["COMPONENT_TYPE_TIKV"])
	}
	if _, ok := cfg["server"]; !ok {
		t.Fatalf("server section should be preserved")
	}
	if _, ok := cfg["o11yGlobalServer"]; !ok {
		t.Fatalf("o11yGlobalServer section should be preserved")
	}
}

func TestPatchDedicatedClusterCPUConfigYAMLAppendsNewComponents(t *testing.T) {
	original := `
targetClusters:
  - id: 123
    name: demo
    components:
      - componentType: COMPONENT_TYPE_TIDB
        cpuNum: 2
`

	patched, err := PatchDedicatedClusterCPUConfigYAML(original, []pipetype.ClusterComponent{
		{ComponentType: "COMPONENT_TYPE_TIDB", CpuNum: 8},
		{ComponentType: "COMPONENT_TYPE_TIFLASH", CpuNum: 4},
	})
	if err != nil {
		t.Fatalf("PatchDedicatedClusterCPUConfigYAML() error = %v", err)
	}

	var cfg map[string]interface{}
	if err := yaml.Unmarshal([]byte(patched), &cfg); err != nil {
		t.Fatalf("yaml.Unmarshal() error = %v", err)
	}

	components := cfg["targetClusters"].([]interface{})[0].(map[string]interface{})["components"].([]interface{})
	got := map[string]float64{}
	for _, item := range components {
		component := item.(map[string]interface{})
		got[component["componentType"].(string)] = component["cpuNum"].(float64)
	}

	if got["COMPONENT_TYPE_TIDB"] != 8 {
		t.Fatalf("TIDB cpuNum = %v, want 8", got["COMPONENT_TYPE_TIDB"])
	}
	if got["COMPONENT_TYPE_TIFLASH"] != 4 {
		t.Fatalf("TIFLASH cpuNum = %v, want 4", got["COMPONENT_TYPE_TIFLASH"])
	}
	if len(components) != 2 {
		t.Fatalf("components len = %d, want 2", len(components))
	}
}