package agent

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/core/identifiers"
	"github.com/tidbcloud/AgentTeam/internal/maas/capability"
	"github.com/tidbcloud/AgentTeam/internal/service/maas"
	storeagent "github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/store/storeerr"
)

const (
	maasLLMProxySessionTTL          = 15 * time.Minute
	maasLLMProxySessionRefreshAfter = 10 * time.Minute
	maasRuntimeSessionReadyWait     = 1500 * time.Millisecond
	maasRuntimeSessionReadyPoll     = 50 * time.Millisecond

	maasRuntimeSessionNotReadyCode    = "agent_runtime_session_not_ready"
	maasRuntimeSessionNotReadyMessage = "agent runtime session is not ready"
)

type MaaSLLMProxyScopeStore interface {
	LoadMaaSLLMProxyScope(ctx context.Context, authenticated storeagent.AuthenticatedAgent) (storeagent.MaaSLLMProxyScope, error)
}

type RuntimeSessionLookup interface {
	AgentRuntimeSessionID(machineID string, agentID string) (string, bool)
	AgentRuntimeRunning(machineID string, agentID string) bool
}

type MaaSAuthorityProjectionPusher interface {
	UpsertAuthorityProjection(ctx context.Context, projection capability.AuthorityProjection) error
}

type HostedAgentValidator interface {
	ValidateAgentBinding(ctx context.Context, params maas.AgentBindingValidation) error
}

type MaaSLLMProxySessionConfig struct {
	Scopes          MaaSLLMProxyScopeStore
	HostedAgents    HostedAgentValidator
	RuntimeSessions RuntimeSessionLookup
	Signer          capability.Signer
	ProjectionPush  MaaSAuthorityProjectionPusher
	RuntimeBaseURL  string
	Clock           func() time.Time
	ReadyWait       time.Duration
	ReadyPoll       time.Duration
}

type MaaSLLMProxySessionService struct {
	scopes          MaaSLLMProxyScopeStore
	hostedAgents    HostedAgentValidator
	runtimeSessions RuntimeSessionLookup
	signer          capability.Signer
	projectionPush  MaaSAuthorityProjectionPusher
	runtimeBaseURL  string
	clock           func() time.Time
	readyWait       time.Duration
	readyPoll       time.Duration
}

func NewMaaSLLMProxySessionService(config MaaSLLMProxySessionConfig) *MaaSLLMProxySessionService {
	readyWait := config.ReadyWait
	if readyWait <= 0 {
		readyWait = maasRuntimeSessionReadyWait
	}
	readyPoll := config.ReadyPoll
	if readyPoll <= 0 {
		readyPoll = maasRuntimeSessionReadyPoll
	}
	if readyPoll > readyWait {
		readyPoll = readyWait
	}
	return &MaaSLLMProxySessionService{
		scopes:          config.Scopes,
		hostedAgents:    config.HostedAgents,
		runtimeSessions: config.RuntimeSessions,
		signer:          config.Signer,
		projectionPush:  config.ProjectionPush,
		runtimeBaseURL:  strings.TrimSpace(config.RuntimeBaseURL),
		clock:           config.Clock,
		readyWait:       readyWait,
		readyPoll:       readyPoll,
	}
}

func (service *MaaSLLMProxySessionService) IssueMaaSLLMProxySession(ctx context.Context, authenticated storeagent.AuthenticatedAgent, selectorSessionID string) (storeagent.LLMProxySession, error) {
	return service.issueMaaSLLMProxySession(ctx, authenticated, selectorSessionID, false)
}

func (service *MaaSLLMProxySessionService) RefreshMaaSLLMProxySession(ctx context.Context, authenticated storeagent.AuthenticatedAgent, selectorSessionID string) (storeagent.LLMProxySession, error) {
	return service.issueMaaSLLMProxySession(ctx, authenticated, selectorSessionID, true)
}

func (service *MaaSLLMProxySessionService) issueMaaSLLMProxySession(ctx context.Context, authenticated storeagent.AuthenticatedAgent, selectorSessionID string, requireSelector bool) (storeagent.LLMProxySession, error) {
	if service == nil || service.scopes == nil || service.projectionPush == nil {
		return storeagent.LLMProxySession{}, errMaaSLLMProxySessionsUnavailable
	}
	runtimeBaseURL := strings.TrimSpace(service.runtimeBaseURL)
	if runtimeBaseURL == "" {
		return storeagent.LLMProxySession{}, errMaaSLLMProxySessionsUnavailable
	}
	trimmedSelector := strings.TrimSpace(selectorSessionID)
	if requireSelector && trimmedSelector == "" {
		return storeagent.LLMProxySession{}, storeerr.Conflict("llm proxy session selector did not match current session")
	}

	scope, err := service.scopes.LoadMaaSLLMProxyScope(ctx, authenticated)
	if err != nil {
		return storeagent.LLMProxySession{}, err
	}
	runtimeSessionID, err := service.waitRuntimeSessionID(ctx, scope.MachineID, scope.AgentID)
	if err != nil {
		return storeagent.LLMProxySession{}, err
	}
	scope.RuntimeSessionID = runtimeSessionID
	if storeagent.NormalizeAgentLLMMode(scope.LLMMode) != storeagent.AgentLLMModeAgentTeamHosted {
		return storeagent.LLMProxySession{}, storeerr.Conflict(maas.ReasonRuntimeUnavailable)
	}
	if service.hostedAgents == nil {
		return storeagent.LLMProxySession{}, storeerr.Conflict(maas.ReasonRuntimeUnavailable)
	}
	if err := service.hostedAgents.ValidateAgentBinding(ctx, maas.AgentBindingValidation{
		OrganizationID: scope.OrganizationID,
		WorkspaceID:    scope.WorkspaceID,
		ModelReference: scope.ModelReference,
		LLMMode:        scope.LLMMode,
		Runtime:        scope.Runtime,
	}); err != nil {
		return storeagent.LLMProxySession{}, err
	}

	now := service.now()
	refreshAt := now.Add(maasLLMProxySessionRefreshAfter)
	expiresAt := now.Add(maasLLMProxySessionTTL)
	sessionID, err := identifiers.NewEntityID("lps")
	if err != nil {
		return storeagent.LLMProxySession{}, fmt.Errorf("generate maas llm proxy session id: %w", err)
	}

	projection := capability.AuthorityProjection{
		OrganizationID:      scope.OrganizationID,
		WorkspaceID:         scope.WorkspaceID,
		AgentID:             scope.AgentID,
		MemberID:            scope.MemberID,
		MachineID:           scope.MachineID,
		RuntimeSessionID:    scope.RuntimeSessionID,
		AuthorityGeneration: scope.AuthorityGeneration,
		OrganizationStatus:  scope.OrganizationStatus,
		WorkspaceStatus:     scope.WorkspaceStatus,
		AgentStatus:         scope.AgentStatus,
		RuntimeStatus:       scope.RuntimeStatus,
	}
	projection.ID = capability.ProjectionID(projection)
	if err := service.projectionPush.UpsertAuthorityProjection(ctx, projection); err != nil {
		return storeagent.LLMProxySession{}, fmt.Errorf("projection_push: %w", err)
	}

	proxyToken, err := service.signer.Sign(capability.Claims{
		Issuer:              capability.DefaultIssuer,
		Subject:             sessionID,
		Audience:            capability.Audience{capability.DefaultAudience},
		ExpiresAt:           expiresAt.Unix(),
		NotBefore:           now.Unix(),
		CapabilitySessionID: sessionID,
		OrganizationID:      strings.TrimSpace(scope.OrganizationID),
		WorkspaceID:         strings.TrimSpace(scope.WorkspaceID),
		AgentID:             strings.TrimSpace(scope.AgentID),
		MemberID:            strings.TrimSpace(scope.MemberID),
		RuntimeSessionID:    strings.TrimSpace(scope.RuntimeSessionID),
		AuthorityGeneration: scope.AuthorityGeneration,
		WorkloadKind:        capability.WorkloadAgentRuntime,
		IssuedAt:            now.Unix(),
	})
	if err != nil {
		return storeagent.LLMProxySession{}, fmt.Errorf("sign maas llm proxy session: %w", err)
	}

	return storeagent.LLMProxySession{
		SessionID:           sessionID,
		ProxyToken:          proxyToken,
		ExpiresAt:           expiresAt,
		RefreshAt:           &refreshAt,
		SupersedesSessionID: trimmedSelector,
		MaasBaseURL:         runtimeBaseURL,
		CredentialKind:      LLMProxyCredentialKindMaaSCapability,
	}, nil
}

func (service *MaaSLLMProxySessionService) currentRuntimeSessionID(machineID string, agentID string) (string, bool) {
	if service == nil || service.runtimeSessions == nil {
		return "", false
	}
	sessionID, ok := service.runtimeSessions.AgentRuntimeSessionID(machineID, agentID)
	sessionID = strings.TrimSpace(sessionID)
	return sessionID, ok && sessionID != ""
}

func (service *MaaSLLMProxySessionService) waitRuntimeSessionID(ctx context.Context, machineID string, agentID string) (string, error) {
	if sessionID, ok := service.currentRuntimeSessionID(machineID, agentID); ok {
		return sessionID, nil
	}
	if !service.runtimeRunning(machineID, agentID) {
		return "", maasRuntimeSessionNotReadyError{}
	}

	timer := time.NewTimer(service.readyWait)
	defer timer.Stop()
	ticker := time.NewTicker(service.readyPoll)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-ticker.C:
			if sessionID, ok := service.currentRuntimeSessionID(machineID, agentID); ok {
				return sessionID, nil
			}
			if !service.runtimeRunning(machineID, agentID) {
				return "", maasRuntimeSessionNotReadyError{}
			}
		case <-timer.C:
			return "", maasRuntimeSessionNotReadyError{}
		}
	}
}

func (service *MaaSLLMProxySessionService) runtimeRunning(machineID string, agentID string) bool {
	if service == nil || service.runtimeSessions == nil {
		return false
	}
	return service.runtimeSessions.AgentRuntimeRunning(machineID, agentID)
}

func (service *MaaSLLMProxySessionService) now() time.Time {
	if service != nil && service.clock != nil {
		return service.clock().UTC()
	}
	return time.Now().UTC()
}

type maasRuntimeSessionNotReadyError struct{}

func (maasRuntimeSessionNotReadyError) Error() string {
	return maasRuntimeSessionNotReadyMessage
}

func (maasRuntimeSessionNotReadyError) Unwrap() error {
	return storeerr.ErrConflict
}