import { EventType } from '../types'

const SHARED: ReadonlySet<string> = new Set([
  'type',
  'timestamp',
  'rawEvent',
  'metadata',
])

function keys(...fields: Array<string>): ReadonlySet<string> {
  return new Set([...SHARED, ...fields])
}

const SPEC_KEYS: ReadonlyMap<string, ReadonlySet<string>> = new Map([
  [EventType.TEXT_MESSAGE_START, keys('messageId', 'role', 'name')],
  [EventType.TEXT_MESSAGE_CONTENT, keys('messageId', 'delta')],
  [EventType.TEXT_MESSAGE_END, keys('messageId')],
  [EventType.TEXT_MESSAGE_CHUNK, keys('messageId', 'role', 'delta', 'name')],
  [
    EventType.TOOL_CALL_START,
    keys('toolCallId', 'toolCallName', 'parentMessageId'),
  ],
  [EventType.TOOL_CALL_ARGS, keys('toolCallId', 'delta')],
  [EventType.TOOL_CALL_END, keys('toolCallId')],
  [
    EventType.TOOL_CALL_CHUNK,
    keys('toolCallId', 'toolCallName', 'parentMessageId', 'delta'),
  ],
  [
    EventType.TOOL_CALL_RESULT,
    keys('messageId', 'toolCallId', 'content', 'role'),
  ],
  [EventType.STATE_SNAPSHOT, keys('snapshot')],
  [EventType.STATE_DELTA, keys('delta')],
  [EventType.MESSAGES_SNAPSHOT, keys('messages')],
  [
    EventType.ACTIVITY_SNAPSHOT,
    keys('messageId', 'activityType', 'content', 'replace'),
  ],
  [EventType.ACTIVITY_DELTA, keys('messageId', 'activityType', 'patch')],
  [EventType.RAW, keys('event', 'source')],
  [EventType.CUSTOM, keys('name', 'value')],
  [
    EventType.RUN_STARTED,
    keys('threadId', 'runId', 'protocolVersion', 'parentRunId', 'input'),
  ],
  [
    EventType.RUN_FINISHED,
    keys('threadId', 'runId', 'result', 'outcome', 'usage'),
  ],
  [EventType.RUN_ERROR, keys('message', 'code', 'usage')],
  [EventType.STEP_STARTED, keys('stepName')],
  [EventType.STEP_FINISHED, keys('stepName')],
  [EventType.REASONING_START, keys('messageId')],
  [EventType.REASONING_MESSAGE_START, keys('messageId', 'role')],
  [EventType.REASONING_MESSAGE_CONTENT, keys('messageId', 'delta')],
  [EventType.REASONING_MESSAGE_END, keys('messageId')],
  [EventType.REASONING_MESSAGE_CHUNK, keys('messageId', 'delta')],
  [EventType.REASONING_END, keys('messageId')],
  [
    EventType.REASONING_ENCRYPTED_VALUE,
    keys('subtype', 'entityId', 'encryptedValue'),
  ],
  [
    EventType.SUBAGENT_STARTED,
    keys(
      'subagentRunId',
      'name',
      'description',
      'parentSubagentRunId',
      'parentToolCallId',
      'parentMessageId',
    ),
  ],
  [EventType.SUBAGENT_FINISHED, keys('subagentRunId', 'result', 'outcome')],
  [EventType.SUBAGENT_ERROR, keys('subagentRunId', 'message', 'code')],
])

const RUN_SCOPED = new Set<string>([
  EventType.RUN_STARTED,
  EventType.RUN_FINISHED,
  EventType.RUN_ERROR,
  EventType.MESSAGES_SNAPSHOT,
])

const ATTRIBUTABLE_KEYS = new Map(
  [...SPEC_KEYS]
    .filter(([type]) => !RUN_SCOPED.has(type))
    .map(
      ([type, fields]) =>
        [type, new Set([...fields, 'subagentRunId'])] as const,
    ),
)

export function specKeysFor(type: string): ReadonlySet<string> {
  return ATTRIBUTABLE_KEYS.get(type) ?? SPEC_KEYS.get(type) ?? SHARED
}

export function isSpecTopLevelKey(type: string, key: string): boolean {
  return specKeysFor(type).has(key)
}