package daemon

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"

	"github.com/tidbcloud/AgentTeam/internal/runtimecontrol"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/versioning"
	"github.com/tidbcloud/AgentTeam/internal/workspace/owner"
)

var upgrader = websocket.Upgrader{
	CheckOrigin:       func(_ *http.Request) bool { return true },
	EnableCompression: true,
}

const daemonTokenRefreshProofHeader = "X-AgentTeam-Daemon-Connection-Token"
const headerRouteWorkspaceID = "X-AgentTeam-Route-Workspace-Id"

// MachineAuthenticator validates daemon websocket API keys.
type MachineAuthenticator interface {
	AuthenticateMachine(ctx context.Context, apiKey string) (agent.MachineIdentity, error)
	IssueAgentAuthTokenForMachine(ctx context.Context, machineID string, agentID string, allowWorkspaceWrite bool) (string, error)
}

// SessionTokenIssuer mints a user session token for a machine identity. Used
// by the daemon to call browser-facing APIs (forwarder/memory) that require
// user-session authentication.
type SessionTokenIssuer interface {
	IssueSessionTokenForUser(ctx context.Context, userID string, ttl time.Duration) (string, error)
}

// RegisterRoutesOptions carries optional daemon route dependencies. Zero-value
// fields disable the corresponding behavior so existing callers keep working.
type RegisterRoutesOptions struct {
	// LatestDesktopVersion supplies the recommended desktop/daemon version for
	// this deployment. When nil, the upgrade-info endpoint omits the field.
	LatestDesktopVersion versioning.Provider
	// MinDesktopVersion is the hard minimum desktop version gate. Empty means
	// no minimum is enforced.
	MinDesktopVersion string
}

// RegisterRoutes mounts the daemon websocket routes.
func RegisterRoutes(router gin.IRoutes, hub *Hub, auth MachineAuthenticator, sessionIssuer SessionTokenIssuer, opts RegisterRoutesOptions) {
	registerDaemonRoutes := func(connectPath string, agentTokenPath string, sessionTokenPath string, upgradeInfoPath string) {
		router.GET(connectPath, func(ctx *gin.Context) {
			apiKey := strings.TrimSpace(ctx.Query("key"))
			if apiKey == "" {
				ctx.JSON(http.StatusUnauthorized, gin.H{"error": "missing machine api key"})
				return
			}
			identity, err := auth.AuthenticateMachine(ctx.Request.Context(), apiKey)
			if err != nil {
				writeMachineAuthError(ctx, err)
				return
			}
			if !requireMatchingRouteWorkspace(ctx, identity.WorkspaceID) {
				return
			}
			responseHeader := http.Header{}
			if workspaceID := strings.TrimSpace(identity.WorkspaceID); workspaceID != "" {
				responseHeader.Set(headerRouteWorkspaceID, workspaceID)
			}
			connection, err := upgrader.Upgrade(ctx.Writer, ctx.Request, responseHeader)
			if err != nil {
				return
			}
			serveCtx, cancel := context.WithCancel(context.Background())
			defer cancel()
			if err := hub.ServeConnection(serveCtx, identity, connection); err != nil && !errors.Is(err, context.Canceled) {
				hub.logger.Warn("daemon websocket closed", slog.String("machine_id", identity.MachineID), slog.String("error_class", "websocket_closed"))
			}
		})

		if elasticPath := elasticRequestLinkConnectPath(connectPath); elasticPath != "" {
			router.GET(elasticPath, func(ctx *gin.Context) {
				if hub == nil || hub.elasticRequestLinks == nil {
					ctx.JSON(http.StatusServiceUnavailable, gin.H{"error": "desktop_unavailable"})
					return
				}
				workspaceID := strings.TrimSpace(ctx.Param("workspaceId"))
				linkID := strings.TrimSpace(ctx.Param("linkId"))
				token := bearerTokenFromHeader(ctx.GetHeader("Authorization"))
				if token == "" {
					hub.elasticRequestLinks.recordElasticEvent("link_rejected", string(ElasticRequestLinkRejectMissingToken))
					ctx.JSON(http.StatusUnauthorized, gin.H{"error": "desktop_unavailable"})
					return
				}
				if err := hub.elasticRequestLinks.PreviewLinkToken(workspaceID, linkID, token); err != nil {
					if hub.logger != nil {
						hub.logger.Warn("reject elastic request link", slog.String("reason", redactedElasticErrorReason(err)))
					}
					ctx.JSON(http.StatusUnauthorized, gin.H{"error": "desktop_unavailable"})
					return
				}
				connection, err := upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
				if err != nil {
					return
				}
				if _, err := hub.elasticRequestLinks.AcceptLink(workspaceID, linkID, token, websocketElasticRequestLinkSender{conn: connection}); err != nil {
					if hub.logger != nil {
						hub.logger.Warn("reject elastic request link after upgrade", slog.String("reason", redactedElasticErrorReason(err)))
					}
					_ = connection.Close()
					return
				}
				hub.serveElasticRequestLink(ctx.Request.Context(), linkID, connection)
			})
		}

		router.POST(agentTokenPath, func(ctx *gin.Context) {
			apiKey := machineAPIKeyFromRequest(ctx)
			if apiKey == "" {
				ctx.JSON(http.StatusUnauthorized, gin.H{"error": "missing machine api key"})
				return
			}
			identity, err := auth.AuthenticateMachine(ctx.Request.Context(), apiKey)
			if err != nil {
				writeMachineAuthError(ctx, err)
				return
			}
			if !requireMatchingRouteWorkspace(ctx, identity.WorkspaceID) {
				return
			}
			allowWorkspaceWrite := hub != nil && hub.agentTokenRefreshWorkspaceWriteAllowed(identity.MachineID, ctx.GetHeader(daemonTokenRefreshProofHeader))
			authToken, err := auth.IssueAgentAuthTokenForMachine(ctx.Request.Context(), identity.MachineID, ctx.Param("agentID"), allowWorkspaceWrite)
			if err != nil {
				if errors.Is(err, owner.ErrOwnerMismatch) {
					writeOwnerMismatchResponse(ctx, identity.WorkspaceID)
					return
				}
				if errors.Is(err, agent.ErrUnauthorized) {
					ctx.JSON(http.StatusUnauthorized, gin.H{"error": "agent is not owned by machine"})
					return
				}
				if errors.Is(err, agent.ErrPermissionModeCapabilityRequired) {
					ctx.JSON(http.StatusConflict, gin.H{"error": runtimecontrol.ErrDaemonPermissionModeUpgradeRequired.Error()})
					return
				}
				if errors.Is(err, agent.ErrPermissionModeStartRequired) {
					ctx.JSON(http.StatusConflict, gin.H{"error": agent.ErrPermissionModeStartRequired.Error()})
					return
				}
				ctx.JSON(http.StatusInternalServerError, gin.H{"error": "issue agent auth token"})
				return
			}
			ctx.Header("Cache-Control", "no-store")
			ctx.JSON(http.StatusOK, gin.H{"authToken": authToken})
		})

		router.POST(sessionTokenPath, func(ctx *gin.Context) {
			apiKey := machineAPIKeyFromRequest(ctx)
			if apiKey == "" {
				ctx.JSON(http.StatusUnauthorized, gin.H{"error": "missing machine api key"})
				return
			}
			identity, err := auth.AuthenticateMachine(ctx.Request.Context(), apiKey)
			if err != nil {
				writeMachineAuthError(ctx, err)
				return
			}
			if !requireMatchingRouteWorkspace(ctx, identity.WorkspaceID) {
				return
			}
			if identity.CreatedByUserID == "" {
				ctx.JSON(http.StatusForbidden, gin.H{"error": "machine has no associated user"})
				return
			}
			token, err := sessionIssuer.IssueSessionTokenForUser(ctx.Request.Context(), identity.CreatedByUserID, daemonSessionTokenTTL)
			if err != nil {
				ctx.JSON(http.StatusInternalServerError, gin.H{"error": "issue session token"})
				return
			}
			ctx.JSON(http.StatusOK, gin.H{"token": token})
		})

		router.GET(upgradeInfoPath, func(ctx *gin.Context) {
			apiKey := machineAPIKeyFromRequest(ctx)
			if apiKey == "" {
				ctx.JSON(http.StatusUnauthorized, gin.H{"error": "missing machine api key"})
				return
			}
			identity, err := auth.AuthenticateMachine(ctx.Request.Context(), apiKey)
			if err != nil {
				writeMachineAuthError(ctx, err)
				return
			}
			if !requireMatchingRouteWorkspace(ctx, identity.WorkspaceID) {
				return
			}
			response := gin.H{
				"serverVersion": versioning.ServerVersion(),
			}
			if opts.LatestDesktopVersion != nil {
				if latest, err := opts.LatestDesktopVersion.Latest(ctx.Request.Context()); err == nil && strings.TrimSpace(latest) != "" {
					response["recommendedDesktopVersion"] = latest
				} else if err != nil && hub != nil {
					hub.logger.Warn("fetch latest desktop version for upgrade-info", slog.Any("error", err))
				}
			}
			if min := strings.TrimSpace(opts.MinDesktopVersion); min != "" {
				response["minDesktopVersion"] = min
			}
			ctx.JSON(http.StatusOK, response)
		})
	}

	registerDaemonRoutes(
		"/api/v1/workspaces/:workspaceId/daemon/connect",
		"/api/v1/workspaces/:workspaceId/daemon/agents/:agentID/auth-token",
		"/api/v1/workspaces/:workspaceId/daemon/session-token",
		"/api/v1/workspaces/:workspaceId/daemon/upgrade-info",
	)
}

func elasticRequestLinkConnectPath(connectPath string) string {
	connectPath = strings.TrimSpace(connectPath)
	if strings.Contains(connectPath, ":workspaceId") {
		return "/api/v1/workspaces/:workspaceId/daemon/elastic-request-links/:linkId/connect"
	}
	return ""
}

func (hub *Hub) serveElasticRequestLink(ctx context.Context, linkID string, connection *websocket.Conn) {
	defer func() {
		if hub != nil && hub.elasticRequestLinks != nil {
			hub.elasticRequestLinks.CloseLink(linkID, "link_disconnected")
		}
	}()
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		_, payload, err := connection.ReadMessage()
		if err != nil {
			return
		}
		var message RemoteRequestResult[json.RawMessage]
		if err := json.Unmarshal(payload, &message); err != nil {
			if hub.logger != nil {
				hub.logger.Warn("decode elastic request link message", slog.String("reason", "protocol_error"))
			}
			continue
		}
		if hub.elasticRequestLinks != nil {
			hub.elasticRequestLinks.CompleteRemoteRequest(message)
		}
	}
}

func bearerTokenFromHeader(authHeader string) string {
	authHeader = strings.TrimSpace(authHeader)
	if len(authHeader) >= len("Bearer ") && strings.EqualFold(strings.TrimSpace(authHeader[:len("Bearer ")]), "Bearer") {
		return strings.TrimSpace(authHeader[len("Bearer "):])
	}
	return ""
}

const daemonSessionTokenTTL = 30 * 24 * time.Hour

func machineAPIKeyFromRequest(ctx *gin.Context) string {
	authHeader := strings.TrimSpace(ctx.GetHeader("Authorization"))
	if len(authHeader) >= len("Bearer ") && strings.EqualFold(strings.TrimSpace(authHeader[:len("Bearer ")]), "Bearer") {
		if token := strings.TrimSpace(authHeader[len("Bearer "):]); token != "" {
			return token
		}
	}
	return strings.TrimSpace(ctx.Query("key"))
}

func writeOwnerMismatchResponse(ctx *gin.Context, workspaceID string) {
	body := gin.H{
		"error": "Workspace owner mismatch",
		"code":  "workspace_owner_mismatch",
	}
	if trimmed := strings.TrimSpace(workspaceID); trimmed != "" {
		body["workspaceId"] = trimmed
	}
	ctx.JSON(http.StatusMisdirectedRequest, body)
}

func writeMachineAuthError(ctx *gin.Context, err error) {
	var ownerMismatch owner.OwnerMismatchError
	if errors.As(err, &ownerMismatch) {
		writeOwnerMismatchResponse(ctx, ownerMismatch.WorkspaceID)
		return
	}
	if errors.Is(err, owner.ErrOwnerMismatch) {
		writeOwnerMismatchResponse(ctx, "")
		return
	}
	if errors.Is(err, agent.ErrUnauthorized) {
		ctx.JSON(http.StatusUnauthorized, gin.H{"error": "invalid machine api key"})
		return
	}
	ctx.JSON(http.StatusServiceUnavailable, gin.H{"error": "machine authentication unavailable"})
}

func requireMatchingRouteWorkspace(ctx *gin.Context, authenticatedWorkspaceID string) bool {
	routeWorkspaceID := strings.TrimSpace(ctx.Param("workspaceId"))
	if routeWorkspaceID == "" {
		routeWorkspaceID = strings.TrimSpace(ctx.GetHeader(headerRouteWorkspaceID))
	}
	authenticated := strings.TrimSpace(authenticatedWorkspaceID)
	if routeWorkspaceID == "" || routeWorkspaceID == authenticated {
		return true
	}
	writeOwnerMismatchResponse(ctx, authenticated)
	return false
}