// MySQL 8+ local admission worker. Outbox delivery uses a deterministic CR identity.
package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"time"
)

func submit(db *sql.DB, id, tenant string, cost, pri int) error {
	if !regexp.MustCompile("^[a-z][a-z0-9-]{0,30}$").MatchString(id) || cost < 1 || cost > 4 || pri < 0 || pri > 100 {
		return fmt.Errorf("invalid request")
	}
	_, e := db.Exec("INSERT INTO reservation(id,tenant,cost,priority,state) VALUES(?,?,?,?,'waiting') ON DUPLICATE KEY UPDATE id=id", id, tenant, cost, pri)
	if e != nil {
		return e
	}
	var t string
	var c, p int
	e = db.QueryRow("SELECT tenant,cost,priority FROM reservation WHERE id=?", id).Scan(&t, &c, &p)
	if e == nil && (t != tenant || c != cost || p != pri) {
		return fmt.Errorf("request_id payload mismatch")
	}
	return e
}
func admit(db *sql.DB) error {
	tx, e := db.BeginTx(context.Background(), &sql.TxOptions{Isolation: sql.LevelReadCommitted})
	if e != nil {
		return e
	}
	defer tx.Rollback()
	if _, e = tx.Exec("UPDATE reservation SET state='expired' WHERE state='waiting' AND created < NOW()-INTERVAL 1 HOUR"); e != nil {
		return e
	}
	var id, t, node string
	var n int
	e = tx.QueryRow("SELECT r.id,r.tenant,r.cost FROM reservation r JOIN tenant_quota q ON q.tenant=r.tenant WHERE r.state='waiting' AND q.used+r.cost<=q.cap AND EXISTS(SELECT 1 FROM gpu_pool p WHERE p.used+r.cost<=p.cap) ORDER BY priority+TIMESTAMPDIFF(SECOND,created,NOW()) DIV 30 DESC,created,id LIMIT 1 FOR UPDATE SKIP LOCKED").Scan(&id, &t, &n)
	if errors.Is(e, sql.ErrNoRows) {
		return tx.Commit()
	}
	if e != nil {
		return e
	}
	e = tx.QueryRow("SELECT node FROM gpu_pool WHERE cap-used>=? ORDER BY cap-used,node LIMIT 1 FOR UPDATE", n).Scan(&node)
	if e != nil {
		return e
	}
	for _, q := range []struct{ sql, key string }{{"UPDATE tenant_quota SET used=used+? WHERE tenant=? AND used+?<=cap", t}, {"UPDATE gpu_pool SET used=used+? WHERE node=? AND used+?<=cap", node}} {
		r, e := tx.Exec(q.sql, n, q.key, n)
		if e != nil {
			return e
		}
		a, _ := r.RowsAffected()
		if a != 1 {
			return fmt.Errorf("capacity changed; retry")
		}
	}
	if _, e = tx.Exec("UPDATE reservation SET state='reserved',node=? WHERE id=?", node, id); e != nil {
		return e
	}
	if _, e = tx.Exec("INSERT INTO outbox(id) VALUES(?)", id); e != nil {
		return e
	}
	return tx.Commit()
}
func kube(args ...string) ([]byte, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 12*time.Second)
	defer cancel()
	b, e := exec.CommandContext(ctx, "kubectl", append([]string{"--context", "kind-handbook-lab", "-n", "modeldock-lab"}, args...)...).CombinedOutput()
	if e != nil {
		return nil, fmt.Errorf("%w: %s", e, b)
	}
	return b, nil
}
func deliver(db *sql.DB) error {
	tx, e := db.Begin()
	if e != nil {
		return e
	}
	defer tx.Rollback()
	var id string
	e = tx.QueryRow("SELECT id FROM outbox WHERE NOT delivered ORDER BY created LIMIT 1 FOR UPDATE SKIP LOCKED").Scan(&id)
	if errors.Is(e, sql.ErrNoRows) {
		return nil
	}
	if e != nil {
		return e
	}
	name := "r-" + id
	obj := map[string]any{"apiVersion": "lab.modeldock.dev/v1alpha1", "kind": "ModelService", "metadata": map[string]any{"name": name, "labels": map[string]string{"admission-id": id}}, "spec": map[string]any{"image": "modeldock-cpu:lab", "replicas": 1}}
	b, _ := json.Marshal(obj)
	f, e := os.CreateTemp("", "outbox-*.json")
	if e != nil {
		return e
	}
	defer os.Remove(f.Name())
	if _, e = f.Write(b); e != nil {
		f.Close()
		return e
	}
	f.Close()
	if _, e = kube("create", "-f", f.Name()); e != nil {
		observed, ge := kube("get", "msvc", name, "-o", "jsonpath={.metadata.labels.admission-id}")
		if ge != nil || string(observed) != id {
			return e
		}
	}
	cr, ce := getCR(id)
	if ce != nil {
		return ce
	}
	if cr == nil {
		return fmt.Errorf("CR disappeared before acknowledgement")
	}
	if _, e = tx.Exec("INSERT INTO binding(id,uid) VALUES(?,?)", id, cr.Metadata.UID); e != nil {
		return e
	}
	_, e = tx.Exec("UPDATE outbox SET delivered=true WHERE id=?", id)
	if e != nil {
		return e
	}
	return tx.Commit()
}

// Only call after CR and pods have disappeared; scope is the dedicated lab (no external writers).
func release(db *sql.DB, id string) error {
	tx, e := db.Begin()
	if e != nil {
		return e
	}
	defer tx.Rollback()
	var state, tenant, node string
	var cost int
	e = tx.QueryRow("SELECT state,tenant,COALESCE(node,''),cost FROM reservation WHERE id=? FOR UPDATE", id).Scan(&state, &tenant, &node, &cost)
	if e != nil {
		return e
	}
	if state == "released" {
		return nil
	}
	if state != "reserved" {
		return fmt.Errorf("not reserved")
	}
	var sent bool
	if e = tx.QueryRow("SELECT delivered FROM outbox WHERE id=?", id).Scan(&sent); e != nil || !sent {
		return fmt.Errorf("await outbox acknowledgement")
	}
	b, e := kube("get", "msvc", "r-"+id, "--ignore-not-found", "-o", "name")
	if e != nil || len(b) > 0 {
		return fmt.Errorf("CR not confirmed absent")
	}
	// Check all terminating pods too, not just Ready pods. Pods use controller UID labels;
	// without a stored UID use the conservative bound: refuse while any owned pod is terminating.
	b, e = kube("get", "pods", "-l", "modeldock-owner", "-o", "json")
	if e != nil {
		return e
	}
	var pods struct {
		Items []struct {
			Metadata struct {
				Name              string
				DeletionTimestamp *time.Time
			}
		}
	}
	if e = json.Unmarshal(b, &pods); e != nil {
		return e
	}
	for _, p := range pods.Items {
		if p.Metadata.DeletionTimestamp != nil || regexp.MustCompile("^r-"+regexp.QuoteMeta(id)+"-").MatchString(p.Metadata.Name) {
			return fmt.Errorf("pod drain incomplete")
		}
	}
	for _, q := range []struct{ sql, key string }{{"UPDATE tenant_quota SET used=used-? WHERE tenant=?", tenant}, {"UPDATE gpu_pool SET used=used-? WHERE node=?", node}} {
		if _, e = tx.Exec(q.sql, cost, q.key); e != nil {
			return e
		}
	}
	if _, e = tx.Exec("UPDATE reservation SET state='released' WHERE id=?", id); e != nil {
		return e
	}
	return tx.Commit()
}
func main() {
	db, e := sql.Open("mysql", os.Getenv("LAB_MYSQL_DSN"))
	if e != nil {
		panic(e)
	}
	defer db.Close()
	a := os.Args[1:]
	switch {
	case len(a) == 5 && a[0] == "submit":
		c, e1 := strconv.Atoi(a[3])
		p, e2 := strconv.Atoi(a[4])
		if e1 != nil || e2 != nil {
			panic("integer required")
		}
		e = submit(db, a[1], a[2], c, p)
	case len(a) == 1 && a[0] == "serve":
		e = serve(db)
	case len(a) == 1 && a[0] == "tick":
		e = admit(db)
		if e == nil {
			e = deliver(db)
		}
	case len(a) == 2 && a[0] == "elastic":
		e = elastic(db, a[1])
	case len(a) == 2 && a[0] == "reconcile":
		e = reconcile(db, a[1])
	case len(a) == 2 && a[0] == "release":
		e = release(db, a[1])
	default:
		panic("submit ID TENANT COST PRIORITY | tick | release ID")
	}
	if e != nil {
		panic(e)
	}
}