package agentruntime

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/service/maas"
	"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/storeerr"
)

type AgentRuntimeStore interface {
	AuthenticateMachine(ctx context.Context, apiKey string) (agent.MachineIdentity, error)
	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
	IssueAgentAuthTokenForMachine(ctx context.Context, machineID string, agentID string, allowWorkspaceWrite bool) (string, error)
	AgentRuntimeForStart(ctx context.Context, agentID string) (string, error)
	PrepareAgentStart(ctx context.Context, agentID string, allowWorkspaceWrite bool, includeLegacySession bool) (agent.PreparedAgentStart, error)
	LoadAgentWorkspaceMaterialization(ctx context.Context, agentID string) (*agent.WorkspaceMaterialization, error)
}

type AgentUnreadStore interface {
	ListAgentUnreadMessageThreads(ctx context.Context, workspaceID string, memberID string) ([]message.AgentUnreadThread, error)
}

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

var ErrRuntimeOffline = errors.New("agent runtime is offline")

type RuntimeControl interface {
	StartAgent(ctx context.Context, machineID string, agentID string) error
	StopAgent(machineID string, agentID string) error
	AgentRuntimeRunning(machineID string, agentID string) bool
	ResetAgentSession(machineID string, agentID string) error
	ResetAgentWorkspace(machineID string, agentID string) error
	GCRecoveryPoints(machineID string, agentID string, points []rewind.MessageRecoveryPoint) error
}

type WebAgentStore interface {
	SetAgentAutoWakeForUser(ctx context.Context, userID string, workspaceID string, agentID string, enabled bool) (agent.AgentView, error)
	GetAgentWorkspaceReaderForUser(ctx context.Context, userID string, workspaceID string, agentID string) (agent.AgentView, error)
	ClearAgentSessionForWorkspace(ctx context.Context, userID string, workspaceID string, agentID string) error
}

type WebRecoveryStore interface {
	DeleteAllAgentRecoveryPointsForWorkspace(ctx context.Context, userID string, workspaceID string, agentID string) ([]rewind.MessageRecoveryPoint, error)
}

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

type Dependencies struct {
	Agents             AgentRuntimeStore
	Messages           AgentUnreadStore
	Recovery           RecoveryStore
	WebAgents          WebAgentStore
	WebRecovery        WebRecoveryStore
	HostedAgents       HostedAgentValidator
	Runtime            RuntimeControl
	IsRuntimeOffline   func(error) bool
	MaaSRuntimeBaseURL string
}

type Service struct {
	agents             AgentRuntimeStore
	messages           AgentUnreadStore
	recovery           RecoveryStore
	webAgents          WebAgentStore
	webRecovery        WebRecoveryStore
	hostedAgents       HostedAgentValidator
	runtime            RuntimeControl
	isRuntimeOffline   func(error) bool
	maasRuntimeBaseURL string
}

func New(deps Dependencies) *Service {
	return &Service{
		agents:             deps.Agents,
		messages:           deps.Messages,
		recovery:           deps.Recovery,
		webAgents:          deps.WebAgents,
		webRecovery:        deps.WebRecovery,
		hostedAgents:       deps.HostedAgents,
		runtime:            deps.Runtime,
		isRuntimeOffline:   deps.IsRuntimeOffline,
		maasRuntimeBaseURL: strings.TrimSpace(deps.MaaSRuntimeBaseURL),
	}
}

func (service *Service) Agents() AgentRuntimeStore { return service.agents }
func (service *Service) Recovery() RecoveryStore   { return service.recovery }

func (service *Service) SetRuntimeControl(runtime RuntimeControl) {
	if service == nil {
		return
	}
	service.runtime = runtime
}

func (service *Service) PrepareAgentStart(ctx context.Context, agentID string, publicBaseURL string, wakeMessage *message.ReceiveMessage, allowWorkspaceWrite bool, includeLegacySession bool, useDirectMaaS bool) (agent.StartCommand, error) {
	if service == nil || service.agents == nil || service.messages == nil {
		return agent.StartCommand{}, fmt.Errorf("agent runtime start service is not configured")
	}
	prepared, err := service.agents.PrepareAgentStart(ctx, agentID, allowWorkspaceWrite, includeLegacySession)
	if err != nil {
		return agent.StartCommand{}, err
	}
	if err := service.validateHostedAgentStart(ctx, prepared, useDirectMaaS); err != nil {
		return agent.StartCommand{}, err
	}

	unreadThreads, err := service.messages.ListAgentUnreadMessageThreads(ctx, prepared.WorkspaceID, prepared.MemberID)
	if err != nil {
		return agent.StartCommand{}, fmt.Errorf("load agent unread summary: %w", err)
	}
	materialization, err := service.agents.LoadAgentWorkspaceMaterialization(ctx, prepared.AgentID)
	if err != nil {
		return agent.StartCommand{}, fmt.Errorf("load agent materialization: %w", err)
	}
	llm, err := service.resolvePreparedAgentLLM(ctx, prepared, publicBaseURL, useDirectMaaS)
	if err != nil {
		return agent.StartCommand{}, fmt.Errorf("resolve agent %s llm for start: %w", strings.TrimSpace(prepared.AgentID), err)
	}
	llmMode := ""
	baseURL := ""
	apiKey := ""
	headers := map[string]string{}
	wireAPI := ""
	credentialKind := ""
	if llm != nil {
		llmMode = strings.TrimSpace(llm.Mode)
		baseURL = strings.TrimSpace(llm.BaseURL)
		apiKey = strings.TrimSpace(llm.APIKey)
		headers = make(map[string]string, len(llm.Headers))
		for key, value := range llm.Headers {
			headers[key] = value
		}
		wireAPI = strings.TrimSpace(llm.WireAPI)
		credentialKind = strings.TrimSpace(llm.CredentialKind)
	}
	permissionMode := agent.DefaultAgentPermissionMode(prepared.PermissionMode)

	return agent.StartCommand{
		AgentID: prepared.AgentID,
		Config: agent.StartConfig{
			Name:                        strings.TrimSpace(prepared.Name),
			DisplayName:                 prepared.DisplayName,
			Description:                 prepared.Description,
			Runtime:                     prepared.Runtime,
			PermissionMode:              permissionMode,
			RuntimePolicy:               prepared.RuntimePolicy,
			VisibleReplyContractEnabled: prepared.VisibleReplyContractEnabled,
			RuntimeDebugEnabled:         prepared.RuntimeDebugEnabled,
			Model:                       prepared.Model,
			ReasoningEffort:             prepared.ReasoningEffort,
			LLMMode:                     llmMode,
			BaseURL:                     baseURL,
			APIKey:                      apiKey,
			Headers:                     headers,
			WireAPI:                     wireAPI,
			CredentialKind:              credentialKind,
			ServerURL:                   publicBaseURL,
			AuthToken:                   prepared.AuthToken,
			SessionID:                   prepared.SessionID,
			Materialization:             materialization,
		},
		WakeMessage:   wakeMessage,
		UnreadSummary: message.BuildAgentUnreadSummary(unreadThreads),
	}, nil
}

func (service *Service) validateHostedAgentStart(ctx context.Context, prepared agent.PreparedAgentStart, useDirectMaaS bool) error {
	if agent.NormalizeAgentLLMMode(prepared.LLMMode) != agent.AgentLLMModeAgentTeamHosted {
		return nil
	}
	if service == nil || service.hostedAgents == nil || strings.TrimSpace(service.maasRuntimeBaseURL) == "" {
		return storeerr.Conflict(maas.ReasonRuntimeUnavailable)
	}
	return service.hostedAgents.ValidateAgentBinding(ctx, maas.AgentBindingValidation{
		OrganizationID:            prepared.OrganizationID,
		WorkspaceID:               prepared.WorkspaceID,
		ModelReference:            prepared.Model,
		LLMMode:                   prepared.LLMMode,
		Runtime:                   prepared.Runtime,
		MachineSupportsDirectMaaS: &useDirectMaaS,
	})
}

func (service *Service) resolvePreparedAgentLLM(ctx context.Context, prepared agent.PreparedAgentStart, publicBaseURL string, useDirectMaaS bool) (*agent.DaemonAgentLLM, error) {
	_ = ctx
	llm := agent.NewAgentLLM(prepared.LLMMode, prepared.BaseURL, prepared.APIKey, prepared.Headers, prepared.WireAPI)
	if llm == nil {
		return nil, nil
	}
	return agent.ProjectAgentLLMForDaemonWithOptions(llm, publicBaseURL, prepared.Runtime, agent.DaemonAgentLLMProjectionOptions{
		UseDirectMaaS:            useDirectMaaS,
		DirectMaaSRuntimeBaseURL: service.maasRuntimeBaseURL,
		RouteWorkspaceID:         prepared.WorkspaceID,
	})
}

func (service *Service) StartAgentRuntimeForUser(ctx context.Context, userID string, workspaceID string, agentID string) (agent.AgentView, error) {
	if service == nil || service.webAgents == nil {
		return agent.AgentView{}, fmt.Errorf("web agent runtime store is not configured")
	}
	agentView, err := service.webAgents.SetAgentAutoWakeForUser(ctx, userID, workspaceID, agentID, true)
	if err != nil {
		return agent.AgentView{}, err
	}
	if service.runtime != nil && !service.runtime.AgentRuntimeRunning(agentView.MachineID, agentView.AgentID) {
		if err := service.runtime.StartAgent(ctx, agentView.MachineID, agentView.AgentID); err != nil && !service.runtimeOffline(err) {
			return agent.AgentView{}, err
		}
	}
	return agentView, nil
}

func (service *Service) StopAgentRuntimeForUser(ctx context.Context, userID string, workspaceID string, agentID string) (agent.AgentView, error) {
	if service == nil || service.webAgents == nil {
		return agent.AgentView{}, fmt.Errorf("web agent runtime store is not configured")
	}
	agentView, err := service.webAgents.SetAgentAutoWakeForUser(ctx, userID, workspaceID, agentID, false)
	if err != nil {
		return agent.AgentView{}, err
	}
	if service.runtime != nil && service.runtime.AgentRuntimeRunning(agentView.MachineID, agentView.AgentID) {
		if err := service.runtime.StopAgent(agentView.MachineID, agentView.AgentID); err != nil && !service.runtimeOffline(err) {
			return agent.AgentView{}, err
		}
	}
	return agentView, nil
}

func (service *Service) ResetAgentRuntimeForUser(ctx context.Context, userID string, workspaceID string, agentID string, clearWorkspace bool) error {
	if service == nil || service.webAgents == nil {
		return fmt.Errorf("web agent runtime store is not configured")
	}
	if clearWorkspace && service.runtime == nil {
		return ErrRuntimeOffline
	}
	agentView, err := service.webAgents.GetAgentWorkspaceReaderForUser(ctx, userID, workspaceID, agentID)
	if err != nil {
		return err
	}
	if service.runtime != nil {
		_ = service.runtime.ResetAgentSession(agentView.MachineID, agentView.AgentID)
	}
	if service.runtime != nil && clearWorkspace {
		if err := service.runtime.ResetAgentWorkspace(agentView.MachineID, agentView.AgentID); err != nil {
			if service.runtimeOffline(err) {
				return ErrRuntimeOffline
			}
			return err
		}
	}
	if service.webRecovery == nil {
		return fmt.Errorf("web agent recovery store is not configured")
	}
	recoveryPoints, err := service.webRecovery.DeleteAllAgentRecoveryPointsForWorkspace(ctx, userID, workspaceID, agentID)
	if err != nil {
		return err
	}
	if service.runtime != nil && len(recoveryPoints) > 0 {
		_ = service.runtime.GCRecoveryPoints(agentView.MachineID, agentView.AgentID, recoveryPoints)
	}
	return service.webAgents.ClearAgentSessionForWorkspace(ctx, userID, workspaceID, agentID)
}

func (service *Service) runtimeOffline(err error) bool {
	return err != nil && service != nil && service.isRuntimeOffline != nil && service.isRuntimeOffline(err)
}