NetMesh/application/state/aiDraftState.ts

import type {
  AIDraft,
  AIPanelView,
} from '../../infrastructure/ai/types';

type DraftsByScope = Partial<Record<string, AIDraft>>;
type PanelViewByScope = Partial<Record<string, AIPanelView>>;
type ActiveSessionIdMap = Record<string, string | null>;
type DraftMutationVersionByScope = Record<string, number>;
type DraftUploadGenerationByScope = Record<string, number>;

const DEFAULT_PANEL_VIEW: AIPanelView = { mode: 'draft' };

function isComposerOnlyDraft(draft: AIDraft | undefined): boolean {
  return Boolean(
    draft
    && draft.attachments.length === 0
    && draft.selectedUserSkillSlugs.length === 0,
  );
}

/**
 * First keystroke creates a missing → composer-only draft. That is still just
 * typing. Clearing a draft (`right` missing or emptied) is a real lifecycle
 * change so New Chat / send can reset the uncontrolled composer.
 */
function draftsEquivalentIgnoringComposerText(
  left: AIDraft | undefined,
  right: AIDraft | undefined,
): boolean {
  if (left === right) return true;
  if (!left) return Boolean(isComposerOnlyDraft(right) && right.text.trim().length > 0);
  if (!right) return false;
  if (left.agentId !== right.agentId) return false;
  if (left.attachments !== right.attachments) return false;
  if (left.selectedUserSkillSlugs !== right.selectedUserSkillSlugs) return false;
  const leftEmpty = left.text.trim().length === 0;
  const rightEmpty = right.text.trim().length === 0;
  if (leftEmpty !== rightEmpty) return leftEmpty;
  return true;
}

/** True when every scope is unchanged except composer text/updatedAt. */
export function draftsByScopeEqualIgnoringAllComposerText(
  prev: DraftsByScope,
  next: DraftsByScope,
): boolean {
  if (prev === next) return true;
  const keys = new Set([...Object.keys(prev), ...Object.keys(next)]);
  for (const key of keys) {
    if (!draftsEquivalentIgnoringComposerText(prev[key], next[key])) return false;
  }
  return true;
}

/** Typing only changes text/updatedAt. Sibling empty-draft creates stay ignored. */
export function draftsByScopeEqualIgnoringComposerText(
  prev: DraftsByScope,
  next: DraftsByScope,
  scopeKey: string,
): boolean {
  if (prev === next) return true;
  const keys = new Set([...Object.keys(prev), ...Object.keys(next)]);
  for (const key of keys) {
    const left = prev[key];
    const right = next[key];
    if (key !== scopeKey && !left && isComposerOnlyDraft(right)) continue;
    if (!draftsEquivalentIgnoringComposerText(left, right)) return false;
  }
  return true;
}

export function createEmptyDraft(agentId: string): AIDraft {
  return {
    text: '',
    agentId,
    attachments: [],
    selectedUserSkillSlugs: [],
    updatedAt: Date.now(),
  };
}

export function getDraftMutationVersionState(
  versionsByScope: DraftMutationVersionByScope,
  scopeKey: string,
): number {
  return versionsByScope[scopeKey] ?? 0;
}

export function bumpDraftMutationVersionState(
  versionsByScope: DraftMutationVersionByScope,
  scopeKey: string,
): DraftMutationVersionByScope {
  return {
    ...versionsByScope,
    [scopeKey]: getDraftMutationVersionState(versionsByScope, scopeKey) + 1,
  };
}

export function getDraftUploadGenerationState(
  generationsByScope: DraftUploadGenerationByScope,
  scopeKey: string,
): number {
  return generationsByScope[scopeKey] ?? 0;
}

export function bumpDraftUploadGenerationState(
  generationsByScope: DraftUploadGenerationByScope,
  scopeKey: string,
): DraftUploadGenerationByScope {
  return {
    ...generationsByScope,
    [scopeKey]: getDraftUploadGenerationState(generationsByScope, scopeKey) + 1,
  };
}

export function resolvePanelView(
  panelViewByScope: PanelViewByScope,
  scopeKey: string,
): AIPanelView {
  return panelViewByScope[scopeKey] ?? DEFAULT_PANEL_VIEW;
}

export function setDraftView(
  panelViewByScope: PanelViewByScope,
  scopeKey: string,
): PanelViewByScope {
  const currentPanelView = panelViewByScope[scopeKey];
  if (currentPanelView?.mode === 'draft') {
    return panelViewByScope;
  }

  return {
    ...panelViewByScope,
    [scopeKey]: DEFAULT_PANEL_VIEW,
  };
}

export function activateDraftView(
  activeSessionIdMap: ActiveSessionIdMap,
  panelViewByScope: PanelViewByScope,
  scopeKey: string,
): {
  activeSessionIdMap: ActiveSessionIdMap;
  panelViewByScope: PanelViewByScope;
} {
  const nextPanelViewByScope = setDraftView(panelViewByScope, scopeKey);
  const hasActiveSession = activeSessionIdMap[scopeKey] != null;

  if (!hasActiveSession) {
    return {
      activeSessionIdMap,
      panelViewByScope: nextPanelViewByScope,
    };
  }

  const nextActiveSessionIdMap = { ...activeSessionIdMap };
  delete nextActiveSessionIdMap[scopeKey];

  return {
    activeSessionIdMap: nextActiveSessionIdMap,
    panelViewByScope: nextPanelViewByScope,
  };
}

export function setSessionView(
  panelViewByScope: PanelViewByScope,
  scopeKey: string,
  sessionId: string,
): PanelViewByScope {
  return {
    ...panelViewByScope,
    [scopeKey]: { mode: 'session', sessionId },
  };
}

export function pruneStaleSessionPanelViews(
  panelViewByScope: PanelViewByScope,
  validSessionIds: Set<string>,
): PanelViewByScope {
  let next = panelViewByScope;

  for (const [scopeKey, panelView] of Object.entries(panelViewByScope)) {
    if (panelView?.mode !== 'session' || validSessionIds.has(panelView.sessionId)) {
      continue;
    }
    const updated = setDraftView(next, scopeKey);
    if (updated !== next) {
      next = updated;
    }
  }

  return next;
}

export function updateDraftForScope(
  draftsByScope: DraftsByScope,
  scopeKey: string,
  fallbackAgentId: string,
  updater: (draft: AIDraft) => AIDraft,
): DraftsByScope {
  const currentDraft = draftsByScope[scopeKey] ?? createEmptyDraft(fallbackAgentId);
  const nextDraft = updater(currentDraft);
  if (nextDraft === currentDraft && draftsByScope[scopeKey] === currentDraft) {
    return draftsByScope;
  }

  return {
    ...draftsByScope,
    [scopeKey]: nextDraft,
  };
}

export function ensureDraftForScopeState(
  draftsByScope: DraftsByScope,
  scopeKey: string,
  agentId: string,
): DraftsByScope {
  if (draftsByScope[scopeKey]) {
    return draftsByScope;
  }

  return {
    ...draftsByScope,
    [scopeKey]: createEmptyDraft(agentId),
  };
}

export function selectDraftForAgentSwitch(
  currentDraft: AIDraft | null | undefined,
  agentId: string,
  startFresh: boolean,
): AIDraft {
  const hasPendingDraftContent = Boolean(
    currentDraft
    && (
      currentDraft.text.length > 0
      || currentDraft.attachments.length > 0
      || currentDraft.selectedUserSkillSlugs.length > 0
    ),
  );

  if (startFresh && !hasPendingDraftContent) {
    return createEmptyDraft(agentId);
  }

  const baseDraft = currentDraft ?? createEmptyDraft(agentId);
  return {
    ...baseDraft,
    agentId,
  };
}

export function clearScopeDraftState(
  draftsByScope: DraftsByScope,
  panelViewByScope: PanelViewByScope,
  scopeKey: string,
): {
  draftsByScope: DraftsByScope;
  panelViewByScope: PanelViewByScope;
} {
  const hasDraft = Object.prototype.hasOwnProperty.call(draftsByScope, scopeKey);
  const hasPanelView = Object.prototype.hasOwnProperty.call(panelViewByScope, scopeKey);

  if (!hasDraft && !hasPanelView) {
    return {
      draftsByScope,
      panelViewByScope,
    };
  }

  return {
    draftsByScope: hasDraft
      ? (() => {
          const nextDrafts = { ...draftsByScope };
          delete nextDrafts[scopeKey];
          return nextDrafts;
        })()
      : draftsByScope,
    panelViewByScope: hasPanelView
      ? (() => {
          const nextPanelViews = { ...panelViewByScope };
          delete nextPanelViews[scopeKey];
          return nextPanelViews;
        })()
      : panelViewByScope,
  };
}

function isClosedTerminalScope(scopeKey: string, activeTerminalTargetIds: Set<string>) {
  if (!scopeKey.startsWith('terminal:')) return false;

  const targetId = scopeKey.slice('terminal:'.length);
  if (!targetId) return false;

  return !activeTerminalTargetIds.has(targetId);
}

export function pruneTerminalScopeState(
  draftsByScope: DraftsByScope,
  panelViewByScope: PanelViewByScope,
  activeTerminalTargetIds: Set<string>,
): {
  draftsByScope: DraftsByScope;
  panelViewByScope: PanelViewByScope;
} {
  const nextDraftsByScope = { ...draftsByScope };
  const nextPanelViewByScope = { ...panelViewByScope };
  let draftsChanged = false;
  let panelViewsChanged = false;

  for (const scopeKey of Object.keys(nextDraftsByScope)) {
    if (!isClosedTerminalScope(scopeKey, activeTerminalTargetIds)) continue;
    delete nextDraftsByScope[scopeKey];
    draftsChanged = true;
  }

  for (const scopeKey of Object.keys(nextPanelViewByScope)) {
    if (!isClosedTerminalScope(scopeKey, activeTerminalTargetIds)) continue;
    delete nextPanelViewByScope[scopeKey];
    panelViewsChanged = true;
  }

  return {
    draftsByScope: draftsChanged ? nextDraftsByScope : draftsByScope,
    panelViewByScope: panelViewsChanged ? nextPanelViewByScope : panelViewByScope,
  };
}

export function pruneTerminalTransientState(
  activeSessionIdMap: ActiveSessionIdMap,
  draftsByScope: DraftsByScope,
  panelViewByScope: PanelViewByScope,
  activeTerminalTargetIds: Set<string>,
): {
  activeSessionIdMap: ActiveSessionIdMap;
  draftsByScope: DraftsByScope;
  panelViewByScope: PanelViewByScope;
} {
  let activeSessionMapChanged = false;
  const nextActiveSessionIdMap: ActiveSessionIdMap = {};

  for (const [scopeKey, sessionId] of Object.entries(activeSessionIdMap)) {
    if (isClosedTerminalScope(scopeKey, activeTerminalTargetIds)) {
      activeSessionMapChanged = true;
      continue;
    }

    nextActiveSessionIdMap[scopeKey] = sessionId;
  }

  const nextTerminalScopeState = pruneTerminalScopeState(
    draftsByScope,
    panelViewByScope,
    activeTerminalTargetIds,
  );

  return {
    activeSessionIdMap: activeSessionMapChanged ? nextActiveSessionIdMap : activeSessionIdMap,
    draftsByScope: nextTerminalScopeState.draftsByScope,
    panelViewByScope: nextTerminalScopeState.panelViewByScope,
  };
}