package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"time"

	corev1 "k8s.io/api/core/v1"
	policyv1 "k8s.io/api/policy/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"sigs.k8s.io/yaml"
)

// k8s clients are rebuilt per call so a refreshed token file (make token + make vm-sync) is picked up without restart.
func (a *App) restConfig() (*rest.Config, error) {
	tok, err := readFileTrim(a.cfg.K8sTokenFile)
	if err != nil {
		return nil, fmt.Errorf("k8s token: %w", err)
	}
	ca, err := os.ReadFile(a.cfg.K8sCAFile)
	if err != nil {
		return nil, fmt.Errorf("k8s ca: %w", err)
	}
	return &rest.Config{Host: a.cfg.K8sAPI, BearerToken: tok, TLSClientConfig: rest.TLSClientConfig{CAData: ca}, Timeout: 30 * time.Second}, nil
}

func (a *App) k8sClients() (kubernetes.Interface, dynamic.Interface, error) {
	rc, err := a.restConfig()
	if err != nil {
		return nil, nil, err
	}
	cs, err := kubernetes.NewForConfig(rc)
	if err != nil {
		return nil, nil, err
	}
	dc, err := dynamic.NewForConfig(rc)
	if err != nil {
		return nil, nil, err
	}
	return cs, dc, nil
}

func (a *App) k8sPodCount(ctx context.Context, ns string) (int, error) {
	cs, _, err := a.k8sClients()
	if err != nil {
		return 0, err
	}
	l, err := cs.CoreV1().Pods(ns).List(ctx, metav1.ListOptions{})
	if err != nil {
		return 0, err
	}
	return len(l.Items), nil
}

// chaos-mesh.org resources are named like their kind in lower case (podchaos, networkchaos, stresschaos ...)
func chaosGVR(kind, version string) schema.GroupVersionResource {
	return schema.GroupVersionResource{Group: "chaos-mesh.org", Version: version, Resource: strings.ToLower(kind)}
}

// applyChaos creates the CR from the scenario manifest with a run-unique name; returns kind/name for revert.
func (a *App) applyChaos(ctx context.Context, manifest, name string) (map[string]any, error) {
	_, dc, err := a.k8sClients()
	if err != nil {
		return nil, err
	}
	var obj unstructured.Unstructured
	if err := yaml.Unmarshal([]byte(manifest), &obj.Object); err != nil {
		return nil, fmt.Errorf("manifest: %w", err)
	}
	gv, err := schema.ParseGroupVersion(obj.GetAPIVersion())
	if err != nil {
		return nil, err
	}
	if gv.Group != "chaos-mesh.org" {
		return nil, fmt.Errorf("only chaos-mesh.org manifests are allowed, got %s", obj.GetAPIVersion())
	}
	ns := obj.GetNamespace()
	if ns == "" {
		ns = "chaos-mesh"
	}
	obj.SetNamespace(ns)
	obj.SetName(name)
	gvr := chaosGVR(obj.GetKind(), gv.Version)
	created, err := dc.Resource(gvr).Namespace(ns).Create(ctx, &obj, metav1.CreateOptions{})
	if err != nil {
		return nil, fmt.Errorf("create %s/%s: %w", obj.GetKind(), name, err)
	}
	return map[string]any{"kind": obj.GetKind(), "version": gv.Version, "namespace": ns, "name": created.GetName()}, nil
}

func (a *App) deleteChaos(ctx context.Context, ref map[string]any) error {
	_, dc, err := a.k8sClients()
	if err != nil {
		return err
	}
	kind, _ := ref["kind"].(string)
	version, _ := ref["version"].(string)
	ns, _ := ref["namespace"].(string)
	name, _ := ref["name"].(string)
	if kind == "" || name == "" {
		return fmt.Errorf("no chaos reference to delete")
	}
	err = dc.Resource(chaosGVR(kind, version)).Namespace(ns).Delete(ctx, name, metav1.DeleteOptions{})
	if err != nil && !strings.Contains(err.Error(), "not found") {
		return err
	}
	return nil
}

func (a *App) scaleDeployment(ctx context.Context, ns, name string, replicas int32) error {
	cs, _, err := a.k8sClients()
	if err != nil {
		return err
	}
	patch := fmt.Sprintf(`{"spec":{"replicas":%d}}`, replicas)
	_, err = cs.AppsV1().Deployments(ns).Patch(ctx, name, types.MergePatchType, []byte(patch), metav1.PatchOptions{})
	return err
}

func (a *App) rolloutRestart(ctx context.Context, ns, name string) error {
	cs, _, err := a.k8sClients()
	if err != nil {
		return err
	}
	patch := fmt.Sprintf(`{"spec":{"template":{"metadata":{"annotations":{"kubectl.kubernetes.io/restartedAt":%q}}}}}`, time.Now().UTC().Format(time.RFC3339))
	_, err = cs.AppsV1().Deployments(ns).Patch(ctx, name, types.StrategicMergePatchType, []byte(patch), metav1.PatchOptions{})
	return err
}

// rolloutUndo copies the pod template of the previous ReplicaSet revision into the deployment (what kubectl rollout undo does).
func (a *App) rolloutUndo(ctx context.Context, ns, name string) error {
	cs, _, err := a.k8sClients()
	if err != nil {
		return err
	}
	dep, err := cs.AppsV1().Deployments(ns).Get(ctx, name, metav1.GetOptions{})
	if err != nil {
		return err
	}
	sel := metav1.FormatLabelSelector(dep.Spec.Selector)
	rss, err := cs.AppsV1().ReplicaSets(ns).List(ctx, metav1.ListOptions{LabelSelector: sel})
	if err != nil {
		return err
	}
	cur := dep.Annotations["deployment.kubernetes.io/revision"]
	var bestRev int64 = -1
	var best int = -1
	for i, rs := range rss.Items {
		rev := rs.Annotations["deployment.kubernetes.io/revision"]
		if rev == cur || rev == "" {
			continue
		}
		var n int64
		fmt.Sscanf(rev, "%d", &n)
		if n > bestRev {
			bestRev, best = n, i
		}
	}
	if best < 0 {
		return fmt.Errorf("no previous revision for %s/%s", ns, name)
	}
	tpl := rss.Items[best].Spec.Template
	delete(tpl.Labels, "pod-template-hash")
	body, _ := json.Marshal(map[string]any{"spec": map[string]any{"template": tpl}})
	_, err = cs.AppsV1().Deployments(ns).Patch(ctx, name, types.StrategicMergePatchType, body, metav1.PatchOptions{})
	return err
}

func (a *App) cordon(ctx context.Context, node string, unschedulable bool) error {
	cs, _, err := a.k8sClients()
	if err != nil {
		return err
	}
	patch := fmt.Sprintf(`{"spec":{"unschedulable":%t}}`, unschedulable)
	_, err = cs.CoreV1().Nodes().Patch(ctx, node, types.MergePatchType, []byte(patch), metav1.PatchOptions{})
	return err
}

// drain = cordon + evict every non-DaemonSet, non-mirror pod on the node (PDBs are honoured by the eviction API)
func (a *App) drain(ctx context.Context, node string) (int, error) {
	if err := a.cordon(ctx, node, true); err != nil {
		return 0, err
	}
	cs, _, err := a.k8sClients()
	if err != nil {
		return 0, err
	}
	pods, err := cs.CoreV1().Pods("").List(ctx, metav1.ListOptions{FieldSelector: "spec.nodeName=" + node})
	if err != nil {
		return 0, err
	}
	n := 0
	for _, p := range pods.Items {
		if isDaemonOrMirror(&p) || p.Status.Phase == corev1.PodSucceeded || p.Status.Phase == corev1.PodFailed {
			continue
		}
		ev := &policyv1.Eviction{ObjectMeta: metav1.ObjectMeta{Name: p.Name, Namespace: p.Namespace}}
		if err := cs.CoreV1().Pods(p.Namespace).EvictV1(ctx, ev); err != nil {
			return n, fmt.Errorf("evict %s/%s: %w", p.Namespace, p.Name, err)
		}
		n++
	}
	return n, nil
}

func isDaemonOrMirror(p *corev1.Pod) bool {
	if _, ok := p.Annotations["kubernetes.io/config.mirror"]; ok {
		return true
	}
	for _, o := range p.OwnerReferences {
		if o.Kind == "DaemonSet" {
			return true
		}
	}
	return false
}