package forwarder

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/store/plan"
)

const meshXResourceQuotaDelta = int64(1)

type meshXResourceQuotaTarget struct {
	metric    plan.MetricKey
	countPath string
}

func checkLoadedWorkspaceResourceQuota(limit plan.WorkspaceResourceLimit, current int64, delta int64) error {
	if limit.Limit < 0 {
		return nil
	}
	if current+delta <= limit.Limit {
		return nil
	}
	return plan.QuotaExceededError{
		OrganizationID: limit.OrganizationID,
		PlanKey:        limit.PlanKey,
		PlanName:       limit.PlanName,
		Metric:         limit.Metric,
		Limit:          limit.Limit,
		Current:        current,
		Delta:          delta,
		Scope:          limit.Scope,
	}
}

func (handler *Handler) requireMeshXResourceQuota(ctx *gin.Context, authenticatedUserID string, jwt string) bool {
	if handler == nil || !handler.billingEnabled {
		return true
	}
	target, ok := meshXResourceQuotaTargetForRequest(ctx.Request)
	if !ok {
		return true
	}
	workspaceID, ok := validatedCurrentWorkspaceID(ctx)
	if !ok {
		api.Error(ctx, http.StatusBadRequest, "No workspace selected")
		return false
	}
	limit, err := handler.planStore.LoadWorkspaceResourceLimit(ctx.Request.Context(), workspaceID, authenticatedUserID, target.metric)
	if err != nil {
		api.StoreError(ctx, err)
		return false
	}
	if limit.Limit < 0 {
		return true
	}
	current, err := handler.countMeshXWorkspaceResource(ctx.Request.Context(), target.countPath, workspaceID, jwt)
	if err != nil {
		if handler.logger != nil {
			handler.logger.Warn("count MeshX workspace resource quota usage failed", "metric", target.metric, "workspaceId", workspaceID, "error", err)
		}
		api.Error(ctx, http.StatusBadGateway, "Failed to check resource quota")
		return false
	}
	if err := checkLoadedWorkspaceResourceQuota(limit, current, meshXResourceQuotaDelta); err == nil {
		return true
	} else {
		api.StoreError(ctx, err)
		return false
	}
}

func meshXResourceQuotaTargetForRequest(request *http.Request) (meshXResourceQuotaTarget, bool) {
	if request == nil || !strings.EqualFold(request.Method, http.MethodPost) || request.URL == nil {
		return meshXResourceQuotaTarget{}, false
	}
	path := strings.TrimRight(strings.TrimSpace(request.URL.Path), "/")
	switch {
	case path == "/api/v1/memory" || (hasForwarderPathPrefix(path, "/api/v1/memory") && strings.HasSuffix(path, "/fork")):
		return meshXResourceQuotaTarget{
			metric:    plan.MetricManagedMemories,
			countPath: "/api/v1/internal/memory/count",
		}, true
	case path == "/api/v1/skills":
		return meshXResourceQuotaTarget{
			metric:    plan.MetricSkillRepositories,
			countPath: "/api/v1/internal/skills/count",
		}, true
	default:
		return meshXResourceQuotaTarget{}, false
	}
}

func (handler *Handler) countMeshXWorkspaceResource(ctx context.Context, path string, workspaceID string, jwt string) (int64, error) {
	if handler == nil || handler.baseURL == nil {
		return 0, fmt.Errorf("API forwarder is not configured")
	}
	path = strings.TrimSpace(path)
	if path == "" {
		return 0, api.ErrBadRequest
	}
	endpoint := *handler.baseURL
	endpoint.Path = singleJoiningSlash(handler.baseURL.Path, path)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint.String(), nil)
	if err != nil {
		return 0, fmt.Errorf("build resource quota count request: %w", err)
	}
	req.Header.Set("Accept", "application/json")
	req.Header.Set(api.HeaderWorkspaceID, strings.TrimSpace(workspaceID))
	if strings.TrimSpace(jwt) != "" {
		authHeader := "Bearer " + jwt
		req.Header.Set("Authorization", authHeader)
		req.Header.Set(trustedJWTHeader, authHeader)
	}

	client := handler.httpClient
	if client == nil {
		client = &http.Client{Timeout: defaultProxyTimeout}
	}
	res, err := client.Do(req)
	if err != nil {
		return 0, fmt.Errorf("send resource quota count request: %w", err)
	}
	defer res.Body.Close()
	if res.StatusCode < http.StatusOK || res.StatusCode >= http.StatusMultipleChoices {
		return 0, upstreamHTTPError(res)
	}
	var payload struct {
		Total int64 `json:"total"`
	}
	if err := json.NewDecoder(res.Body).Decode(&payload); err != nil && !errors.Is(err, io.EOF) {
		return 0, fmt.Errorf("decode resource quota count response: %w", err)
	}
	return payload.Total, nil
}