package llm

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

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/apr/runtime"
	"github.com/tidbcloud/AgentTeam/internal/store/organization"
	"github.com/tidbcloud/AgentTeam/internal/workspace/owner"

	"github.com/gin-gonic/gin"
)

// LLMHandler serves AgentTeam APR-backed proxy HTTP endpoints.
type LLMHandler struct {
	aprDispatcher         *runtime.Dispatcher
	logger                *slog.Logger
	proxyKeyAuthenticator ProxyKeyAuthenticator
	usageTracker          UsageTracker
	managedWorkspaces     owner.ManagedWorkspaceSet
}

type llmProxyAuthorization struct {
	organization.OrganizationProviderProxyAccess
}

type contextKey string

const ctxKeyLLMProxyAuthorization contextKey = "llm.proxy.authorization"

func (handler *LLMHandler) RequireProxyAuthorization() func(ctx *gin.Context) {
	return func(ctx *gin.Context) {
		token, ok := api.ProxyToken(ctx)
		if !ok {
			ctx.JSON(401, gin.H{"error": "missing proxy token"})
			ctx.Abort()
			return
		}
		if handler.proxyKeyAuthenticator == nil {
			ctx.JSON(503, gin.H{"error": "llm proxy authentication is unavailable"})
			ctx.Abort()
			return
		}
		access, err := handler.proxyKeyAuthenticator.AuthenticateOrganizationLLMProxyKey(ctx.Request.Context(), token)
		if err != nil {
			ctx.JSON(401, gin.H{"error": "invalid proxy token"})
			ctx.Abort()
			return
		}
		routeWorkspaceID := strings.TrimSpace(ctx.Param("workspaceId"))
		if workspaceID := strings.TrimSpace(access.WorkspaceID); workspaceID != "" && workspaceID != routeWorkspaceID {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "workspaceId does not match proxy token workspace"})
			ctx.Abort()
			return
		}
		if workspaceID := strings.TrimSpace(access.WorkspaceID); workspaceID != "" {
			if !api.EnsureWorkspaceOwner(ctx, handler.managedWorkspaces, workspaceID) {
				ctx.Abort()
				return
			}
		}
		ctx.Request = ctx.Request.WithContext(context.WithValue(ctx.Request.Context(), ctxKeyLLMProxyAuthorization, llmProxyAuthorization{OrganizationProviderProxyAccess: access}))
		ctx.Next()
	}
}

func (handler *LLMHandler) RequireWorkspaceRoute() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		workspaceID, ok := api.CanonicalRouteWorkspaceID(ctx)
		if !ok {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "workspaceId is invalid"})
			ctx.Abort()
			return
		}
		if !api.EnsureWorkspaceOwner(ctx, handler.managedWorkspaces, workspaceID) {
			ctx.Abort()
			return
		}
		ctx.Next()
	}
}