package daemon

import (
	"context"
	"fmt"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/runtimecontrol"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/store/message"
	"github.com/tidbcloud/AgentTeam/internal/store/rewind"
	"github.com/tidbcloud/AgentTeam/internal/store/workspace"
)

type PrepareAgentStartFunc func(ctx context.Context, agentID string, publicBaseURL string, wakeMessage *message.ReceiveMessage, allowWorkspaceWrite bool, includeLegacySession bool, useDirectMaaS bool) (agent.StartCommand, error)

type AgentStore interface {
	MarkMachineReady(ctx context.Context, identity agent.MachineIdentity, payload agent.ReadyPayload) error
	MarkMachineDisconnected(ctx context.Context, machineID string) error
	ListAgentsForMachine(ctx context.Context, machineID string) ([]agent.MachineAgent, error)
	UpdateAgentStatusFromMachine(ctx context.Context, machineID string, agentID string, status string) (bool, error)
	UpdateAgentActivityFromMachine(ctx context.Context, machineID string, agentID string, activity string, detail string) (bool, error)
	UpdateAgentSessionFromMachine(ctx context.Context, machineID string, agentID string, sessionID string, persistLegacySession bool) (bool, error)
	AgentBelongsToMachine(ctx context.Context, machineID string, agentID string) (bool, error)
	AgentNamesForMachine(ctx context.Context, machineID string, agentID string) (string, string, bool, error)
	AgentPermissionModeForMachine(ctx context.Context, machineID string, agentID string) (string, bool, error)
	MachineIDForAgent(ctx context.Context, workspaceID string, agentID string) (string, error)
	ClearWorkspaceWriteAgentAuthTokensForMachine(ctx context.Context, machineID string) error
}

type RecoveryStore interface {
	RecordMessageRecoveryPointAndPruneOtherThreads(ctx context.Context, workspaceID string, messageID string, agentID string, threadID string, strategy string, recoveryRef string) ([]rewind.MessageRecoveryPoint, error)
}

type TaskActiveTargetStore interface {
	GetTaskActiveTarget(ctx context.Context, workspaceID string, scope string, targetKey string) (message.ReceiveActiveTaskTarget, bool, error)
}

type WorkspaceStore interface {
	ListUserWorkspaces(ctx context.Context, userID string) ([]workspace.Summary, error)
}

type composedStore struct {
	AgentStore
	RecoveryStore
	prepareStart PrepareAgentStartFunc
	workspaces   WorkspaceStore
}

func NewStore(agents AgentStore, recovery RecoveryStore, prepareStart PrepareAgentStartFunc, workspaces ...WorkspaceStore) Store {
	var workspaceLister WorkspaceStore
	if len(workspaces) > 0 {
		workspaceLister = workspaces[0]
	}
	return &composedStore{
		AgentStore:    agents,
		RecoveryStore: recovery,
		prepareStart:  prepareStart,
		workspaces:    workspaceLister,
	}
}

func (s *composedStore) PrepareAgentStart(ctx context.Context, agentID string, publicBaseURL string, wakeMessage *message.ReceiveMessage, allowWorkspaceWrite bool, includeLegacySession bool, useDirectMaaS bool) (agent.StartCommand, error) {
	return s.prepareStart(ctx, agentID, publicBaseURL, wakeMessage, allowWorkspaceWrite, includeLegacySession, useDirectMaaS)
}

func (s *composedStore) GetTaskActiveTarget(ctx context.Context, workspaceID string, scope string, targetKey string) (message.ReceiveActiveTaskTarget, bool, error) {
	activeTargets, ok := s.AgentStore.(TaskActiveTargetStore)
	if !ok || activeTargets == nil {
		return message.ReceiveActiveTaskTarget{}, false, nil
	}
	return activeTargets.GetTaskActiveTarget(ctx, workspaceID, scope, targetKey)
}

func (s *composedStore) ListMachineOwnerWorkspaces(ctx context.Context, machine agent.MachineIdentity) ([]runtimecontrol.MachineWorkspace, error) {
	if s.workspaces == nil {
		return nil, runtimecontrol.ErrMachineWorkspaceListUnavailable
	}
	userID := strings.TrimSpace(machine.CreatedByUserID)
	if userID == "" {
		return nil, fmt.Errorf("machine owner is missing")
	}
	summaries, err := s.workspaces.ListUserWorkspaces(ctx, userID)
	if err != nil {
		return nil, err
	}
	items := make([]runtimecontrol.MachineWorkspace, 0, len(summaries))
	for _, summary := range summaries {
		id := strings.TrimSpace(summary.WorkspaceID)
		if id == "" {
			continue
		}
		name := strings.TrimSpace(summary.Name)
		if name == "" {
			name = strings.TrimSpace(summary.Slug)
		}
		if name == "" {
			name = id
		}
		items = append(items, runtimecontrol.MachineWorkspace{
			ID:   id,
			Name: name,
			Slug: strings.TrimSpace(summary.Slug),
		})
	}
	return items, nil
}