yao/tai/api/register.go

package api

import (
	"fmt"
	"log/slog"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/yaoapp/yao/openapi/oauth"
	tai "github.com/yaoapp/yao/tai"
	"github.com/yaoapp/yao/tai/registry"
	"github.com/yaoapp/yao/tai/taiid"
)

// authenticateBearer validates a Bearer token and returns the caller's identity.
// Package-level var so tests can inject a mock without an OAuth service.
var authenticateBearer = authenticateBearerDefault

func authenticateBearerDefault(token string) (registry.AuthInfo, error) {
	svc := oauth.OAuth
	if svc == nil {
		return registry.AuthInfo{}, fmt.Errorf("oauth service not initialized")
	}
	result, err := svc.AuthenticateToken(oauth.AuthInput{AccessToken: token})
	if err != nil {
		return registry.AuthInfo{}, err
	}
	info := registry.AuthInfo{}
	if result.Info != nil {
		info.Subject = result.Info.Subject
		info.UserID = result.Info.UserID
		info.ClientID = result.Info.ClientID
		info.Scope = result.Info.Scope
		info.TeamID = result.Info.TeamID
		info.TenantID = result.Info.TenantID
	}

	slog.Info("[auth] buildAuthInfo result",
		"subject", info.Subject, "user_id", info.UserID,
		"client_id", info.ClientID, "team_id", info.TeamID,
		"scope", info.Scope)

	if result.Claims != nil {
		slog.Info("[auth] claims",
			"claims.TeamID", result.Claims.TeamID,
			"claims.TenantID", result.Claims.TenantID,
			"claims.ClientID", result.Claims.ClientID,
			"claims.Subject", result.Claims.Subject)
		if result.Claims.Extra != nil {
			slog.Info("[auth] claims.Extra", "extra", fmt.Sprintf("%+v", result.Claims.Extra))
		} else {
			slog.Info("[auth] claims.Extra is nil")
		}

		if info.TeamID == "" {
			switch v := result.Claims.Extra["team_id"].(type) {
			case string:
				info.TeamID = v
				slog.Info("[auth] team_id from Extra (string)", "team_id", v)
			case float64:
				info.TeamID = fmt.Sprintf("%.0f", v)
				slog.Info("[auth] team_id from Extra (float64)", "team_id", info.TeamID)
			default:
				slog.Info("[auth] team_id not found in Extra or unknown type",
					"type", fmt.Sprintf("%T", result.Claims.Extra["team_id"]),
					"value", fmt.Sprintf("%v", result.Claims.Extra["team_id"]))
			}
		}
		if info.TenantID == "" {
			if v, ok := result.Claims.Extra["tenant_id"].(string); ok {
				info.TenantID = v
			}
		}
	}
	return info, nil
}

func extractBearer(r *http.Request) string {
	auth := r.Header.Get("Authorization")
	if len(auth) > 7 && strings.EqualFold(auth[:7], "bearer ") {
		return auth[7:]
	}
	return ""
}

// registerRequest is the JSON body for POST /tai-nodes/register.
type registerRequest struct {
	NodeID       string              `json:"node_id,omitempty"`
	ClientID     string              `json:"client_id,omitempty"`
	MachineID    string              `json:"machine_id"`
	DisplayName  string              `json:"display_name,omitempty"`
	Version      string              `json:"version"`
	Addr         string              `json:"addr"`
	Ports        map[string]int      `json:"ports"`
	Capabilities map[string]bool     `json:"capabilities"`
	System       registry.SystemInfo `json:"system"`
}

// heartbeatRequest is the JSON body for POST /tai-nodes/heartbeat.
type heartbeatRequest struct {
	TaiID string `json:"tai_id"`
}

// HandleRegister handles POST /tai-nodes/register.
// Validates Bearer token, extracts AuthInfo, and writes the node to the Registry.
func HandleRegister(c *gin.Context) {
	reg := registry.Global()
	if reg == nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "registry not initialized"})
		return
	}

	bearer := extractBearer(c.Request)
	if bearer == "" {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"})
		return
	}

	authInfo, err := authenticateBearer(bearer)
	if err != nil {
		slog.Warn("tai register auth failed", "err", err)
		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
		return
	}

	var req registerRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
		return
	}

	if req.NodeID == "" || req.MachineID == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "node_id and machine_id are required"})
		return
	}

	resolvedTaiID, err := taiid.Generate(req.MachineID, req.NodeID)
	if err != nil {
		slog.Warn("taiid generation failed", "err", err)
		c.JSON(http.StatusBadRequest, gin.H{"error": "failed to generate tai_id"})
		return
	}

	remoteIP := c.ClientIP()
	addr := req.Addr
	if addr == "" && remoteIP != "" {
		grpcPort := req.Ports["grpc"]
		if grpcPort > 0 {
			addr = fmt.Sprintf("tai://%s:%d", remoteIP, grpcPort)
		} else {
			addr = remoteIP
		}
	}

	node := &registry.TaiNode{
		TaiID:        resolvedTaiID,
		MachineID:    req.MachineID,
		Version:      req.Version,
		DisplayName:  req.DisplayName,
		Auth:         authInfo,
		System:       req.System,
		Mode:         "direct",
		Addr:         addr,
		Ports:        req.Ports,
		Capabilities: req.Capabilities,
	}
	reg.Register(node)
	slog.Info("[register] node registered via API",
		"tai_id", resolvedTaiID, "addr", addr, "remote_ip", remoteIP,
		"user_id", authInfo.UserID, "team_id", authInfo.TeamID)

	allBefore := reg.List()
	slog.Info("[register] registry snapshot after Register",
		"total", len(allBefore))
	for _, s := range allBefore {
		slog.Info("[register]   node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr)
	}

	if strings.HasPrefix(addr, "tai://") {
		slog.Info("[register] launching connectRegisteredNode goroutine",
			"tai_id", resolvedTaiID, "addr", addr)
		go connectRegisteredNode(resolvedTaiID, addr, reg)
	}

	c.JSON(http.StatusOK, gin.H{
		"status":    "registered",
		"tai_id":    resolvedTaiID,
		"remote_ip": remoteIP,
	})
}

// HandleHeartbeat handles POST /tai-nodes/heartbeat.
// Validates Bearer token and updates the node's last ping timestamp.
func HandleHeartbeat(c *gin.Context) {
	reg := registry.Global()
	if reg == nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "registry not initialized"})
		return
	}

	bearer := extractBearer(c.Request)
	if bearer == "" {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"})
		return
	}

	authInfo, err := authenticateBearer(bearer)
	if err != nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
		return
	}

	var req heartbeatRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
		return
	}
	if req.TaiID == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "tai_id is required"})
		return
	}

	snap, ok := reg.Get(req.TaiID)
	if !ok {
		c.JSON(http.StatusNotFound, gin.H{"error": "tai node not found"})
		return
	}
	if snap.Auth.ClientID != authInfo.ClientID {
		c.JSON(http.StatusForbidden, gin.H{"error": "tai_id does not belong to this client"})
		return
	}

	reg.UpdatePing(req.TaiID)
	c.JSON(http.StatusOK, gin.H{"status": "ok"})
}

// HandleUnregister handles DELETE /tai-nodes/register/:tai_id.
// Validates Bearer token, checks ownership, and removes the node.
func HandleUnregister(c *gin.Context) {
	reg := registry.Global()
	if reg == nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "registry not initialized"})
		return
	}

	bearer := extractBearer(c.Request)
	if bearer == "" {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"})
		return
	}

	authInfo, err := authenticateBearer(bearer)
	if err != nil {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
		return
	}

	taiID := c.Param("tai_id")
	if taiID == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "tai_id is required"})
		return
	}

	snap, ok := reg.Get(taiID)
	if !ok {
		c.JSON(http.StatusNotFound, gin.H{"error": "tai node not found"})
		return
	}
	if snap.Auth.ClientID != authInfo.ClientID {
		c.JSON(http.StatusForbidden, gin.H{"error": "tai_id does not belong to this client"})
		return
	}

	reg.Unregister(taiID)
	slog.Info("tai node unregistered via API", "tai_id", taiID, "user_id", authInfo.UserID)

	c.JSON(http.StatusOK, gin.H{"status": "unregistered"})
}

// connectRegisteredNode dials the self-registered Tai node via gRPC,
// creates a tai.Client, and binds it to the node's TaiID in the registry.
// initRemote internally registers a redundant "host-port" entry; we remove
// it so that the registry contains only the canonical taiID.
func connectRegisteredNode(taiID, addr string, reg *registry.Registry) {
	slog.Info("[connect] start", "tai_id", taiID, "addr", addr)

	client, err := tai.New(addr)
	if err != nil {
		slog.Warn("[connect] tai.New FAILED",
			"tai_id", taiID, "addr", addr, "err", err)

		allAfterFail := reg.List()
		slog.Info("[connect] registry after tai.New failure", "total", len(allAfterFail))
		for _, s := range allAfterFail {
			slog.Info("[connect]   node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr)
		}
		return
	}

	autoID := client.TaiID()
	slog.Info("[connect] tai.New OK", "tai_id", taiID, "autoID", autoID)

	allAfterNew := reg.List()
	slog.Info("[connect] registry after tai.New", "total", len(allAfterNew))
	for _, s := range allAfterNew {
		slog.Info("[connect]   node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr)
	}

	if autoID != "" && autoID != taiID {
		slog.Info("[connect] removing redundant autoID", "autoID", autoID)
		reg.Unregister(autoID)
	}
	reg.SetClient(taiID, client)

	allFinal := reg.List()
	slog.Info("[connect] registry FINAL", "total", len(allFinal))
	for _, s := range allFinal {
		slog.Info("[connect]   node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr)
	}
	slog.Info("[connect] done", "tai_id", taiID)
}