package framework

import (
	"fmt"
	"time"

	"k8s.io/apimachinery/pkg/util/wait"

	"github.com/pingcap/test-infra/sdk/resource"

	apiresource "github.com/pingcap/test-infra/common/model/resource"
	"github.com/pingcap/test-infra/sdk/resource/impl/k8s"
)

func (f *Framework) CreateTiDBCluster(testbed *k8s.Testbed) *k8s.TiDBCluster {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name": "tc",
  "testbed": "%s",
  "type": "%s",
  "spec": {
    "pd": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/pd",
      "version": "v5.4.0"
    },
    "tidb": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tidb",
      "version": "v5.4.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "config": "[log]\nlevel = \"info\"\nnot-exist-config = 1"
    },
    "tikv": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tikv",
      "version": "v5.4.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "evictLeaderTimeout": "10s",
      "config": "[storage]\nreserve-space = \"100MB\"\n[rocksdb]\nmax-open-files = 256"
    },
    "ticdc": {
      "replicas": 0,
      "baseImage": "hub.pingcap.net/qa/ticdc",
      "version": "v5.4.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      }
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeTiDBCluster)))
	ExpectNoError(err)
	return r.(*k8s.TiDBCluster)
}

func (f *Framework) CreateTiDBClusterWithMaxReplicas(testbed *k8s.Testbed) *k8s.TiDBCluster {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
    "name": "tc",
    "testbed": "%s",
    "type": "%s",
    "spec": {
      "pd": {
        "replicas": 1,
        "baseImage": "hub.pingcap.net/qa/pd",
        "version": "v5.4.0"
      },
      "tidb": {
        "replicas": 1,
        "maxReplicas": 2,
        "baseImage": "hub.pingcap.net/qa/tidb",
        "version": "v5.4.0",
        "requests": {
          "cpu": 1,
          "memory": "1Gi"
        },
        "limits": {
          "cpu": 2,
          "memory": "4Gi"
        }
      },
      "tikv": {
        "replicas": 1,
        "baseImage": "hub.pingcap.net/qa/tikv",
        "version": "v5.4.0",
        "requests": {
          "cpu": 1,
          "memory": "1Gi"
        },
        "limits": {
          "cpu": 2,
          "memory": "4Gi"
        }
      }
    }
  }
`, testbed.Name, apiresource.ResourceTypeTiDBCluster)))
	ExpectNoError(err)
	return r.(*k8s.TiDBCluster)
}

func (f *Framework) CreateTiDBClusterWithMPostStartScript(testbed *k8s.Testbed) *k8s.TiDBCluster {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
    "name": "tc",
    "testbed": "%s",
    "type": "%s",
    "spec": {
      "postStart": {
        "sql": "create database test_infra_e2e;\nuse test_infra_e2e;\ncreate table t(id int)"
      },
      "pd": {
        "replicas": 1,
        "baseImage": "hub.pingcap.net/qa/pd",
        "version": "v5.4.0"
      },
      "tidb": {
        "replicas": 1,
        "baseImage": "hub.pingcap.net/qa/tidb",
        "version": "v5.4.0",
        "requests": {
          "cpu": 1,
          "memory": "1Gi"
        },
        "limits": {
          "cpu": 2,
          "memory": "4Gi"
        }
      },
      "tikv": {
        "replicas": 1,
        "baseImage": "hub.pingcap.net/qa/tikv",
        "version": "v5.4.0",
        "requests": {
          "cpu": 1,
          "memory": "1Gi"
        },
        "limits": {
          "cpu": 2,
          "memory": "4Gi"
        }
      }
    }
  }
`, testbed.Name, apiresource.ResourceTypeTiDBCluster)))
	ExpectNoError(err)
	return r.(*k8s.TiDBCluster)
}

func (f *Framework) CreateTiDBClusterWithTiDBHotfix(testbed *k8s.Testbed) *k8s.TiDBCluster {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name": "tc",
  "testbed": "%s",
  "type": "%s",
  "spec": {
    "pd": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/pd",
      "version": "v5.2.0"
    },
    "tidb": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tidb",
      "version": "v5.2.1",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "config": "[log]\nlevel = \"info\"\nnot-exist-config = 1"
    },
    "tikv": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tikv",
      "version": "v5.2.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "evictLeaderTimeout": "10s",
      "config": "[storage]\nreserve-space = \"100MB\"\n[rocksdb]\nmax-open-files = 256"
    },
    "ticdc": {
      "replicas": 0,
      "baseImage": "hub.pingcap.net/qa/ticdc",
      "version": "v5.2.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      }
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeTiDBCluster)))
	ExpectNoError(err)
	return r.(*k8s.TiDBCluster)
}

func (f *Framework) CreateTiDBClusterWithMonitorConfigured(testbed *k8s.Testbed) *k8s.TiDBCluster {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name": "tc",
  "testbed": "%s",
  "type": "%s",
  "spec": {
    "pd": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/pd",
      "version": "v5.4.0"
    },
    "tidb": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tidb",
      "version": "v5.4.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "config": "[log]\nlevel = \"info\"\nnot-exist-config = 1"
    },
    "tikv": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tikv",
      "version": "v5.4.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "evictLeaderTimeout": "10s",
      "config": "[storage]\nreserve-space = \"100MB\"\n[rocksdb]\nmax-open-files = 256"
    },
    "monitor": {
        "initializer": {"baseImage": "registry-mirror.pingcap.net/pingcap/tidb-monitor-initializer", "version": "v5.2.1", "imagePullPolicy": "IfNotPresent"}
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeTiDBCluster)))
	ExpectNoError(err)
	return r.(*k8s.TiDBCluster)
}

func (f *Framework) CreateTiDBClusterWithNGMonitorConfigured(testbed *k8s.Testbed) *k8s.TiDBCluster {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name": "tc",
  "testbed": "%s",
  "type": "%s",
  "spec": {
    "pd": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/pd",
      "version": "v5.4.0"
    },
    "tidb": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tidb",
      "version": "v5.4.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "config": "[log]\nlevel = \"info\"\nnot-exist-config = 1"
    },
    "tikv": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tikv",
      "version": "v5.4.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "evictLeaderTimeout": "10s",
      "config": "[storage]\nreserve-space = \"100MB\"\n[rocksdb]\nmax-open-files = 256"
    },
    "ngMonitoring": {
        "baseImage": "registry-mirror.pingcap.net/pingcap/ng-monitoring", "version": "v5.4.1", "create": true
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeTiDBCluster)))
	ExpectNoError(err)
	return r.(*k8s.TiDBCluster)
}

func (f *Framework) CreateTiDBClusterOOM(testbed *k8s.Testbed) *k8s.TiDBCluster {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)

	// there isn't memory enough for tikv, so it will be OOMKilled during bootstraping
	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name": "tc",
  "testbed": "%s",
  "type": "%s",
  "spec": {
    "pd": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/pd",
      "version": "v5.4.0"
    },
    "tidb": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tidb",
      "version": "v5.4.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "config": "[log]\nlevel = \"info\"\nnot-exist-config = 1"
    },
    "tikv": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tikv",
      "version": "v5.4.0",
      "requests": {
        "cpu": 1,
        "memory": "500Mi"
      },
      "limits": {
        "cpu": 1,
        "memory": "500Mi"
      },
      "evictLeaderTimeout": "10s",
      "config": "[storage]\nreserve-space = \"100MB\"\n[rocksdb]\nmax-open-files = 256"
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeTiDBCluster)))
	ExpectNoError(err)
	return r.(*k8s.TiDBCluster)
}

func (f *Framework) CreateTiDBClusterWithTiCDC(testbed *k8s.Testbed) *k8s.TiDBCluster {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name": "tc",
  "testbed": "%s",
  "type": "%s",
  "spec": {
    "imagePullPolicy": "IfNotPresent",
    "version": "v5.0.0",
    "pd": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/pd"
    },
    "tidb": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tidb",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      }
    },
    "tikv": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tikv",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "evictLeaderTimeout": "10s"
    },
    "ticdc": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/ticdc",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      }
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeTiDBCluster)))
	ExpectNoError(err)
	return r.(*k8s.TiDBCluster)
}

func (f *Framework) CreateTiDBClusterWithTiflash(testbed *k8s.Testbed) *k8s.TiDBCluster {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name": "tc",
  "testbed": "%s",
  "type": "%s",
  "spec": {
    "pd": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/pd",
      "version": "v5.0.0"
    },
    "tidb": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tidb",
      "version": "v5.0.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      }
    },
    "tikv": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tikv",
      "version": "v5.0.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "limits": {
        "cpu": 2,
        "memory": "4Gi"
      },
      "evictLeaderTimeout": "10s"
    },
    "tiflash": {
      "replicas": 1,
      "baseImage": "hub.pingcap.net/qa/tiflash",
      "version": "v5.0.0",
      "requests": {
        "cpu": 1,
        "memory": "1Gi"
      },
      "config": "log-level = \"info\""
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeTiDBCluster)))
	ExpectNoError(err)
	return r.(*k8s.TiDBCluster)
}

func (f *Framework) CreateTiDBClusterWithError(testbed *k8s.Testbed) error {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)

	_, err := client.Create([]byte(fmt.Sprintf(`{
    "name": "tc",
    "testbed": "%s",
    "type": "%s",
    "spec": {
        "pd": {
            "replicas": 1,
			"baseImage": "hub.pingcap.net/qa/pd",
			"version": "v5.0.0"
        },
        "tidb": {
            "replicas": 1,
			"baseImage": "hub.pingcap.net/qa/tidb",
			"version": "v5.0.0",
			"requests": {
				"cpu": 1,
				"memory": "1Gi",
				"topologySpreadConstraints": {
					"topologyKey": "zone"
				}
			},
			"limits": {
				"cpu": 2,
				"memory": "4Gi"
			}
        },
        "tikv": {
            "replicas": 1,
			"baseImage": "hub.pingcap.net/qa/tikv",
			"version": "v5.0.0",
			"requests": {
				"cpu": 1,
				"memory": "1Gi"
			},
			"limits": {
				"cpu": 2,
				"memory": "4Gi"
			}
        }
    }
}
`, testbed.Name, apiresource.ResourceTypeTiDBCluster)))
	ExpectError(err)
	return err
}

func (f *Framework) UpdateTiDBClusterConfig(cluster *k8s.TiDBCluster, opts resource.TiDBClusterUpdateConfigOptions) error {
	err := cluster.UpdateConfig(opts)
	return err
}

func (f *Framework) UpdateTiDBClusterConfigs(cluster *k8s.TiDBCluster, opts resource.TiDBClusterUpdateConfigsOptions) error {
	err := cluster.UpdateConfigs(opts)
	return err
}

func (f *Framework) GetTiDBCluster(handle apiresource.ResourceHandle) *k8s.TiDBCluster {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)
	r, err := client.Get(handle)
	ExpectNoError(err)
	return r.(*k8s.TiDBCluster)
}

func (f *Framework) GetTiDBClusterSpec(cluster *k8s.TiDBCluster, compType resource.TiDBCompType) apiresource.TiDBClusterComponentSpec {
	spec, err := cluster.GetSpec(compType)
	ExpectNoError(err)
	return *spec
}

func (f *Framework) GetTiDBClusterStatus(cluster *k8s.TiDBCluster, compType resource.TiDBCompType) apiresource.TiDBClusterComponentStatus {
	status, err := cluster.GetStatus(compType)
	ExpectNoError(err)
	return *status
}

func (f *Framework) GetTiDBClusterEndpoints(handle apiresource.ResourceHandle) apiresource.ResourceEndpoints {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)
	r, err := client.GetEndpoints(handle)
	ExpectNoError(err)
	return r
}

func (f *Framework) WaitForTiDBClusterReadyWithoutPing(tidb *k8s.TiDBCluster, timeout time.Duration) {
	k8s.SetConfig(k8s.Config{
		TiDBCluster: k8s.TiDBClusterConfig{
			ReadyCheck: k8s.TiDBClusterReadyCheckConfig{
				PingDB: false,
			},
		},
	})
	f.WaitForTiDBClusterReady(tidb, timeout)
}

func (f *Framework) WaitForTiDBClusterReady(tidb *k8s.TiDBCluster, timeout time.Duration) {
	err := tidb.WaitForReady(timeout)
	ExpectNoError(err)
}

func (f *Framework) ArchiveTiDBCluster(handle apiresource.ResourceHandle, archiver resource.ArchiverFunc) {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)
	_, err := client.Artifact(handle, apiresource.GcArtifactAction{Monitor: true, Log: true}, archiver)
	ExpectNoError(err)
}

func (f *Framework) DeleteTiDBClusterWaited(tidbcluster *k8s.TiDBCluster) {
	client := k8s.NewTiDBClusterClient(f.KubeClientSetPack)
	ExpectNoError(client.Delete(tidbcluster.ResourceHandle))

	err := wait.PollImmediate(poll, f.DeletionTimeout, func() (done bool, err error) {
		_, err = client.Get(tidbcluster.ResourceHandle)
		if resource.IsResourceNotFoundError(err) {
			return true, nil
		}
		return false, nil
	})
	ExpectNoError(err)
}