social-app/apps/lib/features/chat/data/models/ag_ui_event.dart

import 'package:json_annotation/json_annotation.dart';
import 'tool_result.dart';

part 'ag_ui_event.g.dart';

class AgUiEventTypeWire {
  static const runStarted = 'RUN_STARTED';
  static const runFinished = 'RUN_FINISHED';
  static const runError = 'RUN_ERROR';
  static const textMessageStart = 'TEXT_MESSAGE_START';
  static const textMessageContent = 'TEXT_MESSAGE_CONTENT';
  static const textMessageEnd = 'TEXT_MESSAGE_END';
  static const toolCallStart = 'TOOL_CALL_START';
  static const toolCallArgs = 'TOOL_CALL_ARGS';
  static const toolCallEnd = 'TOOL_CALL_END';
  static const toolCallResult = 'TOOL_CALL_RESULT';
  static const toolCallError = 'TOOL_CALL_ERROR';
}

enum AgUiEventType {
  runStarted,
  runFinished,
  runError,
  textMessageStart,
  textMessageContent,
  textMessageEnd,
  toolCallStart,
  toolCallArgs,
  toolCallEnd,
  toolCallResult,
  toolCallError,
  unknown,
}

const _wireToTypeMap = {
  AgUiEventTypeWire.runStarted: AgUiEventType.runStarted,
  AgUiEventTypeWire.runFinished: AgUiEventType.runFinished,
  AgUiEventTypeWire.runError: AgUiEventType.runError,
  AgUiEventTypeWire.textMessageStart: AgUiEventType.textMessageStart,
  AgUiEventTypeWire.textMessageContent: AgUiEventType.textMessageContent,
  AgUiEventTypeWire.textMessageEnd: AgUiEventType.textMessageEnd,
  AgUiEventTypeWire.toolCallStart: AgUiEventType.toolCallStart,
  AgUiEventTypeWire.toolCallArgs: AgUiEventType.toolCallArgs,
  AgUiEventTypeWire.toolCallEnd: AgUiEventType.toolCallEnd,
  AgUiEventTypeWire.toolCallResult: AgUiEventType.toolCallResult,
  AgUiEventTypeWire.toolCallError: AgUiEventType.toolCallError,
};

const _typeToWireMap = {
  AgUiEventType.runStarted: AgUiEventTypeWire.runStarted,
  AgUiEventType.runFinished: AgUiEventTypeWire.runFinished,
  AgUiEventType.runError: AgUiEventTypeWire.runError,
  AgUiEventType.textMessageStart: AgUiEventTypeWire.textMessageStart,
  AgUiEventType.textMessageContent: AgUiEventTypeWire.textMessageContent,
  AgUiEventType.textMessageEnd: AgUiEventTypeWire.textMessageEnd,
  AgUiEventType.toolCallStart: AgUiEventTypeWire.toolCallStart,
  AgUiEventType.toolCallArgs: AgUiEventTypeWire.toolCallArgs,
  AgUiEventType.toolCallEnd: AgUiEventTypeWire.toolCallEnd,
  AgUiEventType.toolCallResult: AgUiEventTypeWire.toolCallResult,
  AgUiEventType.toolCallError: AgUiEventTypeWire.toolCallError,
  AgUiEventType.unknown: '',
};

AgUiEventType agUiEventTypeFromWire(String wire) =>
    _wireToTypeMap[wire] ?? AgUiEventType.unknown;

String agUiEventTypeToWire(AgUiEventType type) => _typeToWireMap[type] ?? '';

@JsonSerializable()
class AgUiEvent {
  final AgUiEventType type;

  AgUiEvent({required this.type});

  factory AgUiEvent.fromJson(Map<String, dynamic> json) {
    final typeStr = json['type'] as String? ?? '';
    final type = agUiEventTypeFromWire(typeStr);

    switch (type) {
      case AgUiEventType.runStarted:
        return RunStartedEvent.fromJson(json);
      case AgUiEventType.runFinished:
        return RunFinishedEvent.fromJson(json);
      case AgUiEventType.runError:
        return RunErrorEvent.fromJson(json);
      case AgUiEventType.textMessageStart:
        return TextMessageStartEvent.fromJson(json);
      case AgUiEventType.textMessageContent:
        return TextMessageContentEvent.fromJson(json);
      case AgUiEventType.textMessageEnd:
        return TextMessageEndEvent.fromJson(json);
      case AgUiEventType.toolCallStart:
        return ToolCallStartEvent.fromJson(json);
      case AgUiEventType.toolCallArgs:
        return ToolCallArgsEvent.fromJson(json);
      case AgUiEventType.toolCallEnd:
        return ToolCallEndEvent.fromJson(json);
      case AgUiEventType.toolCallResult:
        return ToolCallResultEvent.fromJson(json);
      case AgUiEventType.toolCallError:
        return ToolCallErrorEvent.fromJson(json);
      case AgUiEventType.unknown:
        return UnknownAgUiEvent.fromJson(json);
    }
  }

  Map<String, dynamic> toJson() => _$AgUiEventToJson(this);
}

@JsonSerializable()
class UnknownAgUiEvent extends AgUiEvent {
  final Map<String, dynamic> rawJson;

  UnknownAgUiEvent({required this.rawJson})
    : super(type: AgUiEventType.unknown);

  factory UnknownAgUiEvent.fromJson(Map<String, dynamic> json) =>
      UnknownAgUiEvent(rawJson: json);

  @override
  Map<String, dynamic> toJson() => rawJson;
}

@JsonSerializable()
class RunStartedEvent extends AgUiEvent {
  final String threadId;
  final String runId;

  RunStartedEvent({required this.threadId, required this.runId})
    : super(type: AgUiEventType.runStarted);

  factory RunStartedEvent.fromJson(Map<String, dynamic> json) =>
      _$RunStartedEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$RunStartedEventToJson(this);
}

@JsonSerializable()
class RunFinishedEvent extends AgUiEvent {
  final String threadId;
  final String runId;

  RunFinishedEvent({required this.threadId, required this.runId})
    : super(type: AgUiEventType.runFinished);

  factory RunFinishedEvent.fromJson(Map<String, dynamic> json) =>
      _$RunFinishedEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$RunFinishedEventToJson(this);
}

@JsonSerializable()
class RunErrorEvent extends AgUiEvent {
  final String message;
  final String? code;

  RunErrorEvent({required this.message, this.code})
    : super(type: AgUiEventType.runError);

  factory RunErrorEvent.fromJson(Map<String, dynamic> json) =>
      _$RunErrorEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$RunErrorEventToJson(this);
}

@JsonSerializable()
class TextMessageStartEvent extends AgUiEvent {
  final String messageId;
  final String role;

  TextMessageStartEvent({required this.messageId, required this.role})
    : super(type: AgUiEventType.textMessageStart);

  factory TextMessageStartEvent.fromJson(Map<String, dynamic> json) =>
      _$TextMessageStartEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$TextMessageStartEventToJson(this);
}

@JsonSerializable()
class TextMessageContentEvent extends AgUiEvent {
  final String messageId;
  final String delta;

  TextMessageContentEvent({required this.messageId, required this.delta})
    : super(type: AgUiEventType.textMessageContent);

  factory TextMessageContentEvent.fromJson(Map<String, dynamic> json) =>
      _$TextMessageContentEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$TextMessageContentEventToJson(this);
}

@JsonSerializable()
class TextMessageEndEvent extends AgUiEvent {
  final String messageId;

  TextMessageEndEvent({required this.messageId})
    : super(type: AgUiEventType.textMessageEnd);

  factory TextMessageEndEvent.fromJson(Map<String, dynamic> json) =>
      _$TextMessageEndEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$TextMessageEndEventToJson(this);
}

@JsonSerializable()
class ToolCallStartEvent extends AgUiEvent {
  final String toolCallId;
  final String toolCallName;
  final String? parentMessageId;

  ToolCallStartEvent({
    required this.toolCallId,
    required this.toolCallName,
    this.parentMessageId,
  }) : super(type: AgUiEventType.toolCallStart);

  factory ToolCallStartEvent.fromJson(Map<String, dynamic> json) =>
      _$ToolCallStartEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$ToolCallStartEventToJson(this);
}

@JsonSerializable()
class ToolCallArgsEvent extends AgUiEvent {
  final String toolCallId;
  final String delta;

  ToolCallArgsEvent({required this.toolCallId, required this.delta})
    : super(type: AgUiEventType.toolCallArgs);

  factory ToolCallArgsEvent.fromJson(Map<String, dynamic> json) =>
      _$ToolCallArgsEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$ToolCallArgsEventToJson(this);
}

@JsonSerializable()
class ToolCallEndEvent extends AgUiEvent {
  final String toolCallId;

  ToolCallEndEvent({required this.toolCallId})
    : super(type: AgUiEventType.toolCallEnd);

  factory ToolCallEndEvent.fromJson(Map<String, dynamic> json) =>
      _$ToolCallEndEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$ToolCallEndEventToJson(this);
}

@JsonSerializable()
class ToolCallResultEvent extends AgUiEvent {
  final String messageId;
  final String toolCallId;
  final Map<String, dynamic> result;
  final UiCard? ui;

  ToolCallResultEvent({
    required this.messageId,
    required this.toolCallId,
    required this.result,
    this.ui,
  }) : super(type: AgUiEventType.toolCallResult);

  factory ToolCallResultEvent.fromJson(Map<String, dynamic> json) =>
      _$ToolCallResultEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$ToolCallResultEventToJson(this);
}

@JsonSerializable()
class ToolCallErrorEvent extends AgUiEvent {
  final String toolCallId;
  final String error;
  final String? code;

  ToolCallErrorEvent({required this.toolCallId, required this.error, this.code})
    : super(type: AgUiEventType.toolCallError);

  factory ToolCallErrorEvent.fromJson(Map<String, dynamic> json) =>
      _$ToolCallErrorEventFromJson(json);

  @override
  Map<String, dynamic> toJson() => _$ToolCallErrorEventToJson(this);
}