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

import 'dart:convert';

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';
  static const stateSnapshot = 'STATE_SNAPSHOT';
  static const messagesSnapshot = 'MESSAGES_SNAPSHOT';
}

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

// wire 类型到枚举的映射
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,
  AgUiEventTypeWire.stateSnapshot: AgUiEventType.stateSnapshot,
  AgUiEventTypeWire.messagesSnapshot: AgUiEventType.messagesSnapshot,
};

// 枚举到 wire 类型的映射
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.stateSnapshot: AgUiEventTypeWire.stateSnapshot,
  AgUiEventType.messagesSnapshot: AgUiEventTypeWire.messagesSnapshot,
  AgUiEventType.unknown: '',
};

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

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

// 类型到工厂函数的映射，用于简化 fromJson
final _typeToFactory = {
  AgUiEventType.runStarted: RunStartedEvent.fromJson,
  AgUiEventType.runFinished: RunFinishedEvent.fromJson,
  AgUiEventType.runError: RunErrorEvent.fromJson,
  AgUiEventType.textMessageStart: TextMessageStartEvent.fromJson,
  AgUiEventType.textMessageContent: TextMessageContentEvent.fromJson,
  AgUiEventType.textMessageEnd: TextMessageEndEvent.fromJson,
  AgUiEventType.toolCallStart: ToolCallStartEvent.fromJson,
  AgUiEventType.toolCallArgs: ToolCallArgsEvent.fromJson,
  AgUiEventType.toolCallEnd: ToolCallEndEvent.fromJson,
  AgUiEventType.toolCallResult: ToolCallResultEvent.fromJson,
  AgUiEventType.toolCallError: ToolCallErrorEvent.fromJson,
  AgUiEventType.stateSnapshot: StateSnapshotEvent.fromJson,
  AgUiEventType.messagesSnapshot: MessagesSnapshotEvent.fromJson,
  AgUiEventType.unknown: UnknownAgUiEvent.fromJson,
};

@JsonSerializable(createFactory: false)
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);
    return _typeToFactory[type]?.call(json) ?? UnknownAgUiEvent.fromJson(json);
  }

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

@JsonSerializable(createFactory: false, createToJson: false)
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(createFactory: false, createToJson: false)
class ToolCallResultEvent extends AgUiEvent {
  final String messageId;
  final String toolCallId;
  final String content;

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

  Map<String, dynamic> get payload {
    try {
      final decoded = jsonDecode(content);
      if (decoded is Map<String, dynamic>) {
        return decoded;
      }
    } catch (_) {}
    return {'content': content};
  }

  Map<String, dynamic> get result {
    final rawResult = payload['result'];
    if (rawResult is Map<String, dynamic>) {
      return rawResult;
    }
    return payload;
  }

  UiCard? get ui {
    final rawUi = payload['ui'];
    if (rawUi is Map<String, dynamic>) {
      return UiCard.fromJson(rawUi);
    }
    final rawResult = payload['result'];
    if (rawResult is Map<String, dynamic>) {
      final type = rawResult['type'];
      final data = rawResult['data'];
      if (type is String && data is Map<String, dynamic>) {
        return UiCard.fromJson(rawResult);
      }
    }
    return null;
  }

  factory ToolCallResultEvent.fromJson(Map<String, dynamic> json) {
    final rawContent = json['content'];
    final content = rawContent is String ? rawContent : '';
    return ToolCallResultEvent(
      messageId: json['messageId'] as String,
      toolCallId: json['toolCallId'] as String,
      content: content,
    );
  }

  @override
  Map<String, dynamic> toJson() => {
    'type': agUiEventTypeToWire(type),
    'messageId': messageId,
    'toolCallId': toolCallId,
    'content': content,
  };
}

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

@JsonSerializable(createFactory: false, createToJson: false)
class StateSnapshotEvent extends AgUiEvent {
  final Map<String, dynamic> snapshot;

  StateSnapshotEvent({required this.snapshot})
    : super(type: AgUiEventType.stateSnapshot);

  factory StateSnapshotEvent.fromJson(Map<String, dynamic> json) {
    final rawSnapshot = json['snapshot'];
    return StateSnapshotEvent(
      snapshot: rawSnapshot is Map<String, dynamic>
          ? rawSnapshot
          : <String, dynamic>{},
    );
  }

  @override
  Map<String, dynamic> toJson() => {
    'type': agUiEventTypeToWire(type),
    'snapshot': snapshot,
  };
}

@JsonSerializable()
class MessagesSnapshotEvent extends AgUiEvent {
  final List<SnapshotMessage> messages;

  MessagesSnapshotEvent({required this.messages})
    : super(type: AgUiEventType.messagesSnapshot);

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

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

@JsonSerializable()
class SnapshotMessage {
  final String id;
  final String role;
  final String? content;
  final String? toolCallId;
  final UiCard? ui;
  final DateTime? timestamp;

  SnapshotMessage({
    required this.id,
    required this.role,
    this.content,
    this.toolCallId,
    this.ui,
    this.timestamp,
  });

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

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