package main

import (
	"context"
	"fmt"
	app "k8s.io/api/apps/v1"
	core "k8s.io/api/core/v1"
	qty "k8s.io/apimachinery/pkg/api/resource"
	meta "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/util/intstr"
	"k8s.io/utils/ptr"
	"os"
	"reflect"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/cache"
	"sigs.k8s.io/controller-runtime/pkg/client"
	cu "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
)

var gvk = schema.GroupVersionKind{Group: "lab.modeldock.dev", Version: "v1alpha1", Kind: "ModelService"}

func resource() *unstructured.Unstructured {
	u := &unstructured.Unstructured{}
	u.SetGroupVersionKind(gvk)
	return u
}

type Controller struct{ client.Client }

const finalizer = "lab.modeldock.dev/children"

func (c Controller) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	u := resource()
	done := ctrl.Result{}
	if err := c.Get(ctx, req.NamespacedName, u); err != nil {
		return done, client.IgnoreNotFound(err)
	}
	d := &app.Deployment{ObjectMeta: meta.ObjectMeta{Name: u.GetName(), Namespace: u.GetNamespace()}}
	s := &core.Service{ObjectMeta: *d.ObjectMeta.DeepCopy()}
	if !u.GetDeletionTimestamp().IsZero() {
		for _, obj := range []client.Object{d, s} {
			err := c.Get(ctx, req.NamespacedName, obj)
			if client.IgnoreNotFound(err) != nil {
				return done, err
			}
			if err == nil && meta.IsControlledBy(obj, u) {
				uid := obj.GetUID()
				if obj.GetDeletionTimestamp() == nil {
					err = c.Delete(ctx, obj, client.Preconditions{UID: &uid}, client.PropagationPolicy(meta.DeletePropagationForeground))
				}
				return ctrl.Result{RequeueAfter: 1000000000}, client.IgnoreNotFound(err)
			}
		}
		if cu.RemoveFinalizer(u, finalizer) {
			return done, c.Update(ctx, u)
		}
		return done, nil
	}
	if cu.AddFinalizer(u, finalizer) {
		return done, c.Update(ctx, u)
	}
	n, _, _ := unstructured.NestedInt64(u.Object, "spec", "replicas")
	img, _, _ := unstructured.NestedString(u.Object, "spec", "image")
	auto, _, _ := unstructured.NestedBool(u.Object, "spec", "autoscale")
	runtimeName, _, _ := unstructured.NestedString(u.Object, "spec", "runtime")
	target := 80
	path := "/"
	if runtimeName == "vllm" {
		target = 8000
		path = "/health"
	}
	labels := map[string]string{"modeldock-owner": string(u.GetUID())}
	own := func(obj client.Object) error {
		if obj.GetResourceVersion() != "" && !meta.IsControlledBy(obj, u) {
			return fmt.Errorf("name collision: %s", obj.GetName())
		}
		obj.SetOwnerReferences([]meta.OwnerReference{*meta.NewControllerRef(u, gvk)})
		return nil
	}
	_, err := cu.CreateOrUpdate(ctx, c.Client, d, func() error {
		if e := own(d); e != nil {
			return e
		}
		if !auto || d.Spec.Replicas == nil {
			d.Spec.Replicas = ptr.To(int32(n))
		}
		d.Spec.Selector = &meta.LabelSelector{MatchLabels: labels}
		d.Spec.Template.Labels = labels
		if len(d.Spec.Template.Spec.Containers) == 0 {
			d.Spec.Template.Spec.Containers = []core.Container{{Name: "runtime"}}
		}
		x := &d.Spec.Template.Spec.Containers[0]
		x.Image = img
		if runtimeName == "vllm" {
			model, ok, _ := unstructured.NestedString(u.Object, "spec", "model")
			if !ok {
				return fmt.Errorf("vllm requires spec.model")
			}
			count, ok, _ := unstructured.NestedInt64(u.Object, "spec", "gpuPerReplica")
			if !ok {
				count = 1
			}
			x.Args = []string{"--model", model, "--served-model-name", u.GetName(), "--max-model-len", "4096", "--tensor-parallel-size", fmt.Sprint(count)}
			x.Resources.Limits = core.ResourceList{"nvidia.com/gpu": *qty.NewQuantity(count, qty.DecimalSI)}
			x.Resources.Requests = core.ResourceList{"nvidia.com/gpu": *qty.NewQuantity(count, qty.DecimalSI)}
			pvc, ok, _ := unstructured.NestedString(u.Object, "spec", "cachePVC")
			if ok {
				d.Spec.Template.Spec.Volumes = []core.Volume{{Name: "weights", VolumeSource: core.VolumeSource{PersistentVolumeClaim: &core.PersistentVolumeClaimVolumeSource{ClaimName: pvc}}}}
				x.VolumeMounts = []core.VolumeMount{{Name: "weights", MountPath: "/root/.cache/huggingface"}}
			} else {
				x.VolumeMounts = nil
				d.Spec.Template.Spec.Volumes = nil
			}
		} else {
			x.Args = nil
			x.Resources = core.ResourceRequirements{}
			x.VolumeMounts = nil
			d.Spec.Template.Spec.Volumes = nil
		}
		if x.ReadinessProbe == nil || x.ReadinessProbe.HTTPGet == nil || x.ReadinessProbe.HTTPGet.Port.IntVal != int32(target) {
			x.ReadinessProbe = &core.Probe{ProbeHandler: core.ProbeHandler{HTTPGet: &core.HTTPGetAction{Path: path, Port: intstr.FromInt(target)}}}
		}
		return nil
	})
	if err != nil {
		return done, err
	}
	_, err = cu.CreateOrUpdate(ctx, c.Client, s, func() error {
		if e := own(s); e != nil {
			return e
		}
		s.Spec.Selector = labels
		if len(s.Spec.Ports) == 0 || s.Spec.Ports[0].TargetPort.IntVal != int32(target) {
			s.Spec.Ports = []core.ServicePort{{Name: "http", Port: 80, Protocol: core.ProtocolTCP, TargetPort: intstr.FromInt(target)}}
		}
		return nil
	})
	if err != nil {
		return done, err
	}
	ready := d.Status.ObservedGeneration == d.Generation && d.Status.AvailableReplicas == *d.Spec.Replicas && d.Status.UpdatedReplicas == *d.Spec.Replicas && *d.Spec.Replicas > 0
	phase, truth := "Progressing", "False"
	if ready {
		phase, truth = "Running", "True"
	}
	prior, _, _ := unstructured.NestedMap(u.Object, "status")
	transition := meta.Now().Format("2006-01-02T15:04:05Z07:00")
	conditions, _, _ := unstructured.NestedSlice(u.Object, "status", "conditions")
	if len(conditions) > 0 {
		if old, ok := conditions[0].(map[string]any); ok && old["status"] == truth {
			if t, ok := old["lastTransitionTime"].(string); ok {
				transition = t
			}
		}
	}
	status := map[string]any{"phase": phase, "readyReplicas": int64(d.Status.ReadyReplicas), "observedGeneration": u.GetGeneration(), "conditions": []any{map[string]any{"type": "Ready", "status": truth, "reason": phase, "lastTransitionTime": transition, "observedGeneration": u.GetGeneration()}}}
	if reflect.DeepEqual(status, prior) {
		return done, nil
	}
	u.Object["status"] = status
	return done, c.Status().Update(ctx, u)
}
func main() {
	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{Cache: cache.Options{DefaultNamespaces: map[string]cache.Config{"modeldock-lab": {}}}, LeaderElection: os.Getenv("LEADER_ELECTION") == "true", LeaderElectionID: "modeldock.lab", LeaderElectionNamespace: "modeldock-lab"})
	if err != nil {
		panic(err)
	}
	err = ctrl.NewControllerManagedBy(mgr).Named("modelservice").For(resource()).Owns(&app.Deployment{}).Owns(&core.Service{}).Complete(Controller{mgr.GetClient()})
	if err != nil {
		panic(err)
	}
	if err = mgr.Start(ctrl.SetupSignalHandler()); err != nil {
		panic(err)
	}
}