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,
}

AgUiEventType agUiEventTypeFromWire(String wire) {
  switch (wire) {
    case AgUiEventTypeWire.runStarted:
      return AgUiEventType.runStarted;
    case AgUiEventTypeWire.runFinished:
      return AgUiEventType.runFinished;
    case AgUiEventTypeWire.runError:
      return AgUiEventType.runError;
    case AgUiEventTypeWire.textMessageStart:
      return AgUiEventType.textMessageStart;
    case AgUiEventTypeWire.textMessageContent:
      return AgUiEventType.textMessageContent;
    case AgUiEventTypeWire.textMessageEnd:
      return AgUiEventType.textMessageEnd;
    case AgUiEventTypeWire.toolCallStart:
      return AgUiEventType.toolCallStart;
    case AgUiEventTypeWire.toolCallArgs:
      return AgUiEventType.toolCallArgs;
    case AgUiEventTypeWire.toolCallEnd:
      return AgUiEventType.toolCallEnd;
    case AgUiEventTypeWire.toolCallResult:
      return AgUiEventType.toolCallResult;
    case AgUiEventTypeWire.toolCallError:
      return AgUiEventType.toolCallError;
    default:
      return AgUiEventType.unknown;
  }
}

String agUiEventTypeToWire(AgUiEventType type) {
  switch (type) {
    case AgUiEventType.runStarted:
      return AgUiEventTypeWire.runStarted;
    case AgUiEventType.runFinished:
      return AgUiEventTypeWire.runFinished;
    case AgUiEventType.runError:
      return AgUiEventTypeWire.runError;
    case AgUiEventType.textMessageStart:
      return AgUiEventTypeWire.textMessageStart;
    case AgUiEventType.textMessageContent:
      return AgUiEventTypeWire.textMessageContent;
    case AgUiEventType.textMessageEnd:
      return AgUiEventTypeWire.textMessageEnd;
    case AgUiEventType.toolCallStart:
      return AgUiEventTypeWire.toolCallStart;
    case AgUiEventType.toolCallArgs:
      return AgUiEventTypeWire.toolCallArgs;
    case AgUiEventType.toolCallEnd:
      return AgUiEventTypeWire.toolCallEnd;
    case AgUiEventType.toolCallResult:
      return AgUiEventTypeWire.toolCallResult;
    case AgUiEventType.toolCallError:
      return AgUiEventTypeWire.toolCallError;
    case AgUiEventType.unknown:
      return '';
  }
}

@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);
}