package rewind

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

	"github.com/tidbcloud/AgentTeam/internal/core/rewindplan"
	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/store/message"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
	"github.com/tidbcloud/AgentTeam/internal/store/storeerr"
)

type Queryer interface {
	Query(context.Context, string, ...any) (coredb.Rows, error)
	QueryRow(context.Context, string, ...any) coredb.Row
}

type TargetMessage struct {
	ID                      string    `db:"id"`
	Seq                     int64     `db:"seq"`
	GlobalSeq               int64     `db:"global_seq"`
	AuthorMemberID          string    `db:"author_member_id"`
	NotifyPrincipalMemberID string    `db:"notify_principal_member_id"`
	AuthorType              string    `db:"author_type"`
	MessageType             string    `db:"message_type"`
	CreatedAt               time.Time `db:"created_at"`
}

type SuffixMessage struct {
	ID                      string    `db:"id"`
	Seq                     int64     `db:"seq"`
	GlobalSeq               int64     `db:"global_seq"`
	AuthorMemberID          string    `db:"author_member_id"`
	NotifyPrincipalMemberID string    `db:"notify_principal_member_id"`
	AllowAuthorUnread       bool      `db:"allow_author_unread"`
	AuthorType              string    `db:"author_type"`
	MessageType             string    `db:"message_type"`
	CreatedAt               time.Time `db:"created_at"`
	AuthorAgentID           string    `db:"author_agent_id"`
	AuthorMachineID         string    `db:"author_machine_id"`
	AuthorRuntime           string    `db:"author_runtime"`
	AuthorModel             string    `db:"author_model"`
	AuthorSessionID         string    `db:"author_session_id"`
	AuthorAgentLive         bool      `db:"author_agent_live"`
}

type targetMessageRow struct {
	TargetMessage
	ArchivedUntilSeq int64 `db:"archived_until_seq"`
}

type suffixDeliveryRow struct {
	MessageID   string    `db:"message_id"`
	GlobalSeq   int64     `db:"global_seq"`
	DeliveredAt time.Time `db:"delivered_at"`
	Payload     []byte    `db:"payload"`
}

type recoveryAgentRow struct {
	AgentMemberID    string `db:"agent_member_id"`
	AgentID          string `db:"agent_id"`
	MachineID        string `db:"machine_id"`
	Runtime          string `db:"runtime"`
	Model            string `db:"model"`
	SessionID        string `db:"session_id"`
	AgentRuntimeLive bool   `db:"agent_runtime_live"`
}

func recoveryDeliveryRow(messageID string, globalSeq int64, deliveredAt time.Time, agent recoveryAgentRow) model.RewindDeliveryRow {
	return model.RewindDeliveryRow{
		MessageID:        messageID,
		GlobalSeq:        globalSeq,
		DeliveredAt:      deliveredAt,
		AgentMemberID:    agent.AgentMemberID,
		AgentID:          agent.AgentID,
		MachineID:        agent.MachineID,
		Runtime:          agent.Runtime,
		Model:            agent.Model,
		SessionID:        agent.SessionID,
		AgentRuntimeLive: agent.AgentRuntimeLive,
	}
}

func LoadTargetMessage(ctx context.Context, queryer Queryer, workspaceID string, scopeID string, messageID string) (TargetMessage, int64, error) {
	var row targetMessageRow
	if err := coredb.Get(ctx, queryer, &row, `
SELECT
    m.id AS id,
    m.seq AS seq,
    m.workspace_global_seq AS global_seq,
    m.author_member_id AS author_member_id,
    COALESCE(NULLIF(m.notify_principal_member_id, ''), m.author_member_id) AS notify_principal_member_id,
    m.author_type AS author_type,
    COALESCE(m.message_type, '') AS message_type,
    m.created_at AS created_at,
    s.archived_until_seq AS archived_until_seq
FROM scope_messages m
JOIN scopes s
  ON s.id = m.scope_id
 AND s.workspace_id = m.workspace_id
 AND s.status = 'active'
 AND s.archived_at IS NULL
WHERE m.workspace_id = $1
  AND m.scope_id = $2
  AND m.id = $3
  AND `+message.MessageNotRewoundSQL("m")+`
`, strings.TrimSpace(workspaceID), strings.TrimSpace(scopeID), strings.TrimSpace(messageID)); err != nil {
		if coredb.IsNotFound(err) {
			return TargetMessage{}, 0, storeerr.ErrNotFound
		}
		return TargetMessage{}, 0, fmt.Errorf("load rewind target message: %w", err)
	}
	return row.TargetMessage, row.ArchivedUntilSeq, nil
}

func LoadSuffixMessages(ctx context.Context, queryer Queryer, workspaceID string, scopeID string, fromSeq int64) ([]SuffixMessage, error) {
	var messages []SuffixMessage
	if err := coredb.Select(ctx, queryer, &messages, `
SELECT
    m.id AS id,
    m.seq AS seq,
    m.workspace_global_seq AS global_seq,
    m.author_member_id AS author_member_id,
    COALESCE(NULLIF(m.notify_principal_member_id, ''), m.author_member_id) AS notify_principal_member_id,
    CASE
        WHEN EXISTS (
            SELECT 1 FROM workflow_run_effects wre
            WHERE wre.workspace_id = m.workspace_id AND wre.message_id = m.id
        ) OR EXISTS (
            SELECT 1 FROM automation_runs ar
            WHERE ar.workspace_id = m.workspace_id AND ar.message_id = m.id
        ) OR COALESCE(JSON_UNQUOTE(JSON_EXTRACT(m.metadata, '$.workflow_run_id')), '') <> ''
        THEN TRUE
        ELSE FALSE
    END AS allow_author_unread,
    m.author_type AS author_type,
    COALESCE(m.message_type, '') AS message_type,
    m.created_at AS created_at,
    COALESCE(a.id, '') AS author_agent_id,
    COALESCE(a.machine_id, '') AS author_machine_id,
    COALESCE(a.runtime, '') AS author_runtime,
    COALESCE(a.model, '') AS author_model,
    '' AS author_session_id,
    CASE
        WHEN m.author_type = 'agent' AND a.id IS NULL THEN FALSE
        ELSE TRUE
    END AS author_agent_live
FROM scope_messages m
JOIN scopes s
  ON s.id = m.scope_id
 AND s.workspace_id = m.workspace_id
 AND s.status = 'active'
 AND s.archived_at IS NULL
LEFT JOIN agents a
  ON a.workspace_id = m.workspace_id
 AND a.member_id = m.author_member_id
 AND m.author_type = 'agent'
WHERE m.workspace_id = $1
  AND m.scope_id = $2
  AND m.seq >= $3
  AND `+message.MessageNotRewoundSQL("m")+`
  AND `+message.MessageNotRewindSystemNoticeSQL("m")+`
ORDER BY m.seq ASC
`, strings.TrimSpace(workspaceID), strings.TrimSpace(scopeID), fromSeq); err != nil {
		return nil, fmt.Errorf("load rewind suffix messages: %w", err)
	}
	return messages, nil
}

func LoadSuffixDeliveries(ctx context.Context, queryer Queryer, workspaceID string, scopeID string, fromSeq int64) (map[string][]model.RewindDeliveryRow, error) {
	var rows []suffixDeliveryRow
	if err := coredb.Select(ctx, queryer, &rows, `
SELECT
    m.id AS message_id,
    m.workspace_global_seq AS global_seq,
    COALESCE(sfo.delivered_at, m.created_at) AS delivered_at,
    COALESCE(CAST(sfo.payload AS CHAR), '') AS payload
FROM scope_messages m
JOIN scopes s
  ON s.id = m.scope_id
 AND s.workspace_id = m.workspace_id
 AND s.status = 'active'
 AND s.archived_at IS NULL
JOIN scope_fanout_outbox sfo
  ON sfo.workspace_id = m.workspace_id
 AND sfo.scope_id = m.scope_id
 AND sfo.message_id = m.id
 AND sfo.fanout_type = 'scope_fanout'
 AND sfo.status = 'delivered'
 AND sfo.suppressed_at IS NULL
WHERE m.workspace_id = $1
  AND m.scope_id = $2
  AND m.seq >= $3
  AND `+message.MessageNotRewoundSQL("m")+`
  AND `+message.MessageNotRewindSystemNoticeSQL("m")+`
ORDER BY m.seq ASC
`, strings.TrimSpace(workspaceID), strings.TrimSpace(scopeID), fromSeq); err != nil {
		return nil, fmt.Errorf("load rewind suffix delivery payloads: %w", err)
	}

	recipientAgents := make(map[string]struct{})
	recipientsByMessageID := make(map[string][]message.AgentRecipient, len(rows))
	rowByMessageID := make(map[string]suffixDeliveryRow, len(rows))
	for _, row := range rows {
		wakeRecipients, err := wakeRecipientsFromScopeFanoutPayload(row.Payload)
		if err != nil {
			return nil, fmt.Errorf("load rewind suffix delivery payload for %s: %w", strings.TrimSpace(row.MessageID), err)
		}
		if len(wakeRecipients) == 0 {
			continue
		}
		rowByMessageID[strings.TrimSpace(row.MessageID)] = row
		for _, recipient := range wakeRecipients {
			agentID := strings.TrimSpace(recipient.AgentID)
			if agentID == "" {
				continue
			}
			recipientAgents[agentID] = struct{}{}
			recipientsByMessageID[strings.TrimSpace(row.MessageID)] = append(recipientsByMessageID[strings.TrimSpace(row.MessageID)], recipient)
		}
	}
	if len(recipientAgents) == 0 {
		return map[string][]model.RewindDeliveryRow{}, nil
	}

	agentIDs := make([]string, 0, len(recipientAgents))
	for agentID := range recipientAgents {
		agentIDs = append(agentIDs, agentID)
	}
	sort.Strings(agentIDs)
	agents, err := loadRecoveryAgentRows(ctx, queryer, strings.TrimSpace(workspaceID), agentIDs)
	if err != nil {
		return nil, err
	}

	deliveries := make(map[string][]model.RewindDeliveryRow)
	seen := make(map[string]struct{})
	for messageID, recipients := range recipientsByMessageID {
		row := rowByMessageID[messageID]
		for _, recipient := range recipients {
			agentID := strings.TrimSpace(recipient.AgentID)
			if agentID == "" {
				continue
			}
			key := messageID + "\x00" + agentID
			if _, ok := seen[key]; ok {
				continue
			}
			seen[key] = struct{}{}
			agent, ok := agents[agentID]
			if !ok {
				agent = recoveryAgentRow{AgentID: agentID, AgentMemberID: strings.TrimSpace(recipient.AgentMemberID), MachineID: strings.TrimSpace(recipient.MachineID)}
			}
			deliveries[messageID] = append(deliveries[messageID], recoveryDeliveryRow(messageID, row.GlobalSeq, row.DeliveredAt, agent))
		}
	}
	return deliveries, nil
}

func wakeRecipientsFromScopeFanoutPayload(payload []byte) ([]message.AgentRecipient, error) {
	if len(payload) == 0 {
		return nil, nil
	}
	fanoutPayload, err := message.UnmarshalScopeFanoutPayload(payload)
	if err == nil {
		return message.MergeAgentRecipients(fanoutPayload.Delivery.WakeRecipients), nil
	}
	intentPayload, intentErr := message.UnmarshalScopeDeliveryIntentPayload(payload)
	if intentErr == nil {
		return message.MergeAgentRecipients(intentPayload.WakeRecipients), nil
	}
	return nil, err
}

func loadRecoveryAgentRows(ctx context.Context, queryer Queryer, workspaceID string, agentIDs []string) (map[string]recoveryAgentRow, error) {
	if len(agentIDs) == 0 {
		return map[string]recoveryAgentRow{}, nil
	}
	agentIDsCTE, agentIDArgs, _, err := coredb.BuildStringValuesCTE("requested_agent_ids", "agent_id", agentIDs, 2)
	if err != nil {
		return nil, err
	}
	query := fmt.Sprintf(`
WITH %s
SELECT
    COALESCE(a.member_id, '') AS agent_member_id,
    requested_agent_ids.agent_id AS agent_id,
    COALESCE(a.machine_id, '') AS machine_id,
    COALESCE(a.runtime, '') AS runtime,
    COALESCE(a.model, '') AS model,
    '' AS session_id,
    CASE
        WHEN a.id IS NOT NULL AND a.status <> 'archived' AND mac.id IS NOT NULL AND mac.status = 'active' THEN TRUE
        ELSE FALSE
    END AS agent_runtime_live
FROM requested_agent_ids
LEFT JOIN agents a
  ON a.workspace_id = $1
 AND a.id = requested_agent_ids.agent_id
LEFT JOIN machines mac
  ON mac.workspace_id = a.workspace_id
 AND mac.id = a.machine_id
ORDER BY requested_agent_ids.agent_id ASC
`, agentIDsCTE)
	var rows []recoveryAgentRow
	if err := coredb.Select(ctx, queryer, &rows, query, append([]any{strings.TrimSpace(workspaceID)}, agentIDArgs...)...); err != nil {
		return nil, fmt.Errorf("load rewind recovery agents: %w", err)
	}
	agents := make(map[string]recoveryAgentRow, len(rows))
	for _, row := range rows {
		agents[strings.TrimSpace(row.AgentID)] = row
	}
	return agents, nil
}

type messageRecoveryPointRow struct {
	ThreadID    string `db:"thread_id"`
	Strategy    string `db:"recovery_strategy"`
	RecoveryRef string `db:"recovery_ref"`
}

func LoadMessageRecoveryPoint(ctx context.Context, queryer Queryer, workspaceID string, messageID string, agentID string) (string, string, string, error) {
	var row messageRecoveryPointRow
	if err := coredb.Get(ctx, queryer, &row, `
SELECT thread_id, recovery_strategy, recovery_ref
FROM message_recovery_points
WHERE workspace_id = $1 AND message_id = $2 AND agent_id = $3
`, strings.TrimSpace(workspaceID), strings.TrimSpace(messageID), strings.TrimSpace(agentID)); err != nil {
		if coredb.IsNotFound(err) {
			return "", "", "", storeerr.ErrNotFound
		}
		return "", "", "", fmt.Errorf("load message recovery point: %w", err)
	}
	return strings.TrimSpace(row.ThreadID), strings.TrimSpace(row.Strategy), strings.TrimSpace(row.RecoveryRef), nil
}

func LoadEarliestMessageForRecoveryPoint(ctx context.Context, queryer Queryer, workspaceID string, scopeID string, agentID string, threadID string, strategy string, recoveryRef string, atOrBeforeSeq int64) (TargetMessage, error) {
	var row TargetMessage
	if err := coredb.Get(ctx, queryer, &row, `
SELECT
    m.id AS id,
    m.seq AS seq,
    m.workspace_global_seq AS global_seq,
    m.author_member_id AS author_member_id,
    COALESCE(NULLIF(m.notify_principal_member_id, ''), m.author_member_id) AS notify_principal_member_id,
    m.author_type AS author_type,
    COALESCE(m.message_type, '') AS message_type,
    m.created_at AS created_at
FROM message_recovery_points mrp
JOIN scope_messages m
  ON m.workspace_id = mrp.workspace_id
 AND m.id = mrp.message_id
JOIN scopes s
  ON s.id = m.scope_id
 AND s.workspace_id = m.workspace_id
 AND s.status = 'active'
 AND s.archived_at IS NULL
WHERE mrp.workspace_id = $1
  AND mrp.agent_id = $2
  AND mrp.thread_id = $3
  AND mrp.recovery_strategy = $4
  AND mrp.recovery_ref = $5
  AND m.scope_id = $6
  AND m.seq <= $7
  AND `+message.MessageNotRewoundSQL("m")+`
ORDER BY m.seq ASC
LIMIT 1
`, strings.TrimSpace(workspaceID), strings.TrimSpace(agentID), strings.TrimSpace(threadID), strings.TrimSpace(strategy), strings.TrimSpace(recoveryRef), strings.TrimSpace(scopeID), atOrBeforeSeq); err != nil {
		if coredb.IsNotFound(err) {
			return TargetMessage{}, storeerr.ErrNotFound
		}
		return TargetMessage{}, fmt.Errorf("load earliest message for recovery point: %w", err)
	}
	return row, nil
}

func LoadLatestRecoveryPointBeforeSeq(ctx context.Context, queryer Queryer, workspaceID string, scopeID string, agentID string, beforeSeq int64) (string, string, string, error) {
	var row messageRecoveryPointRow
	if err := coredb.Get(ctx, queryer, &row, `
SELECT mrp.thread_id, mrp.recovery_strategy, mrp.recovery_ref
FROM message_recovery_points mrp
JOIN scope_messages m
  ON m.workspace_id = mrp.workspace_id
 AND m.id = mrp.message_id
JOIN scopes s
  ON s.id = m.scope_id
 AND s.workspace_id = m.workspace_id
 AND s.status = 'active'
 AND s.archived_at IS NULL
WHERE mrp.workspace_id = $1
  AND mrp.agent_id = $2
  AND m.scope_id = $3
  AND m.seq < $4
  AND `+message.MessageNotRewoundSQL("m")+`
ORDER BY m.seq DESC
LIMIT 1
`, strings.TrimSpace(workspaceID), strings.TrimSpace(agentID), strings.TrimSpace(scopeID), beforeSeq); err != nil {
		if coredb.IsNotFound(err) {
			return "", "", "", storeerr.ErrNotFound
		}
		return "", "", "", fmt.Errorf("load latest recovery point before seq: %w", err)
	}
	return strings.TrimSpace(row.ThreadID), strings.TrimSpace(row.Strategy), strings.TrimSpace(row.RecoveryRef), nil
}

func AgentHasCrossScopeActivityAfterExcluding(ctx context.Context, queryer Queryer, workspaceID string, agentID string, excludeScopeIDs []string, since time.Time) (bool, error) {
	trimmedWorkspaceID := strings.TrimSpace(workspaceID)
	trimmedAgentID := strings.TrimSpace(agentID)
	normalizedExcludeScopeIDs := rewindplan.NormalizeNonEmptyStrings(excludeScopeIDs)

	baseExistsSQL := `
SELECT EXISTS (
    SELECT 1
    FROM scope_messages m
    JOIN agents a
      ON a.workspace_id = m.workspace_id
     AND a.member_id = m.author_member_id
    WHERE m.workspace_id = $1
      AND a.id = $2
      AND m.author_type = 'agent'
      AND m.created_at > $3
      AND ` + message.MessageNotRewoundSQL("m") + `
      %s
    UNION
    SELECT 1
    FROM scope_activity_events ae
    JOIN agents a
      ON a.workspace_id = ae.workspace_id
     AND a.member_id = ae.actor_member_id
    WHERE ae.workspace_id = $1
      AND a.id = $2
      AND ae.created_at > $3
      AND ae.suppressed_at IS NULL
      %s
)
`

	if len(normalizedExcludeScopeIDs) == 0 {
		row := queryer.QueryRow(ctx, fmt.Sprintf(baseExistsSQL, "", ""), trimmedWorkspaceID, trimmedAgentID, since)
		var conflicted bool
		if err := row.Scan(&conflicted); err != nil {
			return false, fmt.Errorf("check rewind cross-scope activity excluding: %w", err)
		}
		return conflicted, nil
	}

	excludedScopesCTE, excludedScopeArgs, _, err := coredb.BuildStringValuesCTE("excluded_scopes", "scope_id", normalizedExcludeScopeIDs, 4)
	if err != nil {
		return false, err
	}
	messageExclusion := `
      AND NOT EXISTS (
          SELECT 1
          FROM excluded_scopes
          WHERE excluded_scopes.scope_id = m.scope_id
      )`
	activityExclusion := `
      AND NOT EXISTS (
          SELECT 1
          FROM excluded_scopes
          WHERE excluded_scopes.scope_id = ae.scope_id
      )`
	query := fmt.Sprintf("WITH %s\n%s", excludedScopesCTE, fmt.Sprintf(baseExistsSQL, messageExclusion, activityExclusion))
	args := append([]any{trimmedWorkspaceID, trimmedAgentID, since}, excludedScopeArgs...)
	row := queryer.QueryRow(ctx, query, args...)

	var conflicted bool
	if err := row.Scan(&conflicted); err != nil {
		return false, fmt.Errorf("check rewind cross-scope activity excluding: %w", err)
	}
	return conflicted, nil
}

func LoadCurrentSuffixMessageIDs(ctx context.Context, queryer Queryer, workspaceID string, scopeID string, targetMessageID string) ([]string, error) {
	var rawIDs []string
	var err error
	if strings.TrimSpace(targetMessageID) == "" {
		// Load all messages in the scope (for cascade reply thread validation).
		err = coredb.Select(ctx, queryer, &rawIDs, `
SELECT m.id
FROM scope_messages m
JOIN scopes s
  ON s.id = m.scope_id
 AND s.workspace_id = m.workspace_id
 AND s.status = 'active'
 AND s.archived_at IS NULL
WHERE m.workspace_id = $1
  AND m.scope_id = $2
  AND `+message.MessageNotRewoundSQL("m")+`
  AND `+message.MessageNotRewindSystemNoticeSQL("m")+`
ORDER BY m.seq ASC
`, strings.TrimSpace(workspaceID), strings.TrimSpace(scopeID))
	} else {
		err = coredb.Select(ctx, queryer, &rawIDs, `
WITH target AS (
    SELECT m.seq
    FROM scope_messages m
    JOIN scopes s
      ON s.id = m.scope_id
     AND s.workspace_id = m.workspace_id
     AND s.status = 'active'
     AND s.archived_at IS NULL
    WHERE m.workspace_id = $1
      AND m.scope_id = $2
      AND m.id = $3
      AND `+message.MessageNotRewoundSQL("m")+`
)
SELECT m.id
FROM target
JOIN scope_messages m
JOIN scopes s
  ON s.id = m.scope_id
 AND s.workspace_id = m.workspace_id
 AND s.status = 'active'
 AND s.archived_at IS NULL
WHERE m.workspace_id = $1
  AND m.scope_id = $2
  AND m.seq >= target.seq
  AND `+message.MessageNotRewoundSQL("m")+`
  AND `+message.MessageNotRewindSystemNoticeSQL("m")+`
ORDER BY m.seq ASC
`, strings.TrimSpace(workspaceID), strings.TrimSpace(scopeID), strings.TrimSpace(targetMessageID))
	}
	if err != nil {
		return nil, fmt.Errorf("load current rewind suffix message ids: %w", err)
	}
	ids := make([]string, 0, len(rawIDs))
	for _, id := range rawIDs {
		ids = append(ids, strings.TrimSpace(id))
	}
	return ids, nil
}