package auth

import (
	"strings"
	"time"

	"github.com/gin-gonic/gin"

	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
)

const (
	sessionIDContextKey         = "agentteam.uiSession.id"
	sessionExpiresAtContextKey  = "agentteam.uiSession.expiresAt"
	DesktopUISessionInvalidCode = "DESKTOP_UI_SESSION_INVALID"
)

// Middleware validates renderer session tokens (Bearer or X-Session-Token) and
// local CLI control credentials (X-Api-Key). It aborts with 401 on failure.
func Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		var authToken string

		if h := c.GetHeader("Authorization"); h != "" {
			if tok, ok := strings.CutPrefix(h, "Bearer "); ok && tok != "" {
				authToken = tok
			}
		}

		if authToken == "" {
			authToken = c.GetHeader("X-Session-Token")
		}

		if AuthenticateSession(c, authToken) {
			c.Next()
			return
		}

		apiKey := c.GetHeader("X-Api-Key")
		if apiKey != "" {
			if err := token.Default().ValidateTempKey(apiKey); err == nil {
				c.Next()
				return
			}
		}

		c.Header("X-AgentTeam-Desktop-Auth-Error", "ui-session-invalid")
		c.AbortWithStatusJSON(401, gin.H{
			"code":  DesktopUISessionInvalidCode,
			"error": "Desktop UI session is invalid.",
		})
	}
}

// AuthenticateSession validates and installs one explicit UI session token.
// Callers such as the Desktop gateway use this for transports that cannot set
// request headers and carry their local session in an existing query field.
func AuthenticateSession(c *gin.Context, sessionToken string) bool {
	metadata, ok := token.SessionInfo(strings.TrimSpace(sessionToken))
	if !ok {
		return false
	}
	c.Set(sessionIDContextKey, metadata.LogicalSessionID)
	c.Set(sessionExpiresAtContextKey, metadata.ExpiresAt)
	return true
}

// SessionIdentity returns the authenticated UI session identity installed by
// Middleware. Local CLI control credentials intentionally do not produce a UI session.
func SessionIdentity(c *gin.Context) (string, time.Time, bool) {
	sessionID := c.GetString(sessionIDContextKey)
	expiresAtValue, ok := c.Get(sessionExpiresAtContextKey)
	if sessionID == "" || !ok {
		return "", time.Time{}, false
	}
	expiresAt, ok := expiresAtValue.(time.Time)
	if !ok || expiresAt.IsZero() {
		return "", time.Time{}, false
	}
	return sessionID, expiresAt, true
}