// live.dart
import 'dart:convert';
import 'dart:typed_data';

import 'package:collection/collection.dart';
import 'package:fonnx/models/pyannote/pyannote.dart';
import 'package:freezed_annotation/freezed_annotation.dart';
import 'package:hooks_riverpod/hooks_riverpod.dart';
import 'package:loggy/loggy.dart' hide LogLevel;
import 'package:telosnex/agent/agent_script_select.dart';
import 'package:telosnex/chat/models/tool.dart' hide Tool;
import 'package:telosnex/chat/realtime/audio_ring_buffer.dart';
import 'package:telosnex/chat/realtime/live_history_string.dart';
import 'package:telosnex/chat/realtime/map_has_crosstalk.dart';
import 'package:telosnex/chat/realtime/realtime_models.dart' hide Conversation;
import 'package:telosnex/misc/pcm_to_wav.dart';
import 'package:telosnex/models/companion_providers.dart';
import 'package:telosnex/models/conversation.dart';
import 'package:telosnex/chat/models/role.dart' as tlsnx_role;
import 'package:telosnex/models/script.dart';
import 'package:telosnex/persistence/document_manager.dart';
import 'package:telosnex/persistence/path_providers.dart';
import 'package:telosnex/run/run.dart';
import 'package:telosnex/run/run_status.dart';
import 'package:telosnex/stt/pyannote_path.dart';
import 'package:telosnex/sync/rate_limited_publisher.dart';
import 'package:telosnex/widgets/patrol_keys.dart';

import 'live_architecture.dart';
import 'live_event_handlers.dart';
import 'live_internal_handlers.dart';
import 'live_internal_events.dart';
import 'package:riverpod_annotation/riverpod_annotation.dart';
import 'package:telosnex/loggy/live_loggy.dart';
import 'package:telosnex/chat/realtime/realtime_socket.dart';
import 'package:telosnex/models/conversation_manager.dart';
import 'package:telosnex/models/api_keys.dart';
import 'package:telosnex/run/run_manager.dart';
import 'package:telosnex/models/script_manager.dart';
import 'package:telosnex/models/settings.dart';
import 'package:telosnex/stt/microphone_service.dart';
import 'package:flutter_pcm_sound/flutter_pcm_sound.dart';
import 'dart:async';
import 'package:telosnex/models/loom_step.dart'; // For LiveState initial steps

part 'live.freezed.dart';
part 'live.g.dart';

enum LiveEvent { started, stopped, error }

@freezed
class LiveState with _$LiveState {
  const factory LiveState({
    required bool isRunning,
    required bool isRemoteTalking,
    required bool isSocketConnected,
    required bool didRemoteTalkInFirstSecond,
    required Map<String, LoomStep> itemIdToLoomStep,
    // If *any* audio in the first second is detected, it is assumed that it is
    // echo from the assistant. If thats true, that implies headphones could not
    // be connected. With this information, the crosstalk detection algorithm
    // can make better decisions.
    //
    // Motivated by the assistant reading a poem
    // dramatically, which at second 21, rolled such that there were no segments
    // in the first 1.02 seconds. This was evaluated as a headphone interruption.
    required bool isMicMutedByUser,
    required bool isSpeakerMutedByUser,
    String? conversationId,
    @Default([]) List<LiveEvent> events,
  }) = _LiveState;

  factory LiveState.initial() => LiveState(
        isRunning: false,
        isRemoteTalking: false,
        isSocketConnected: false,
        didRemoteTalkInFirstSecond: false,
        itemIdToLoomStep: {},
        isMicMutedByUser: false,
        isSpeakerMutedByUser: false,
      );
}

const _kLiveDebug = true;

void liveLog(String Function() message) {
  if (_kLiveDebug) {
    logDebug('[LIVE] ${message()}');
  }
}

void liveLogError(String Function() message) {
  logError('[LIVE] ${message()}');
}

@Riverpod(
    dependencies: [ScriptManager, apiKeys, ConversationManager, RunManager])
class LiveNotifier extends _$LiveNotifier with LiveLoggy {
  static const kInjectedStepIdPrefix = 'ClientC_';
  // static const kFpcmLogLevel = LogLevel.none;

  RealtimeSocket? _socket;
  StreamSubscription? _socketEventsSubscription;
  StreamSubscription? _microphoneSubscription;
  MicrophoneService? _microphone;
  AudioRingBuffer? _micDataWhenRemoteTalkingBuffer;
  AudioRingBuffer? _micDataWhenSocketConnectingBuffer;

  final Set<String> _skippedAssistantItemIds = {};
  Pyannote? _pyannote; // Pyannote model
  ThrottledPublisher? _pyannoteThrottler;
  BytesBuilder _lastAudioPlayed = BytesBuilder();
  BytesBuilder _lastAudioFromAssistant = BytesBuilder();

  @override
  FutureOr<LiveState> build() {
    ref.onDispose(() => stop());
    _lastAudioPlayed = BytesBuilder();
    _lastAudioFromAssistant = BytesBuilder();
    ref.listen(companionScriptIdProvider, (previous, next) {
      if (!state.value!.isRunning) return;
      final commands = [RefreshSessionInstructionsCommand(next)];
      _processCommands(commands);
    });
    return LiveState.initial();
  }

  void muteMic() {
    final event = MicMuteEvent(isMuted: true);
    final result = handleInternalEvent(state.value!, event);
    state = AsyncData(result.newState);
    _processCommands(result.commands);
  }

  void unmuteMic() {
    final event = MicMuteEvent(isMuted: false);
    final result = handleInternalEvent(state.value!, event);
    state = AsyncData(result.newState);
    _processCommands(result.commands);
  }

  void muteSpeaker() {
    final event = SpeakerMuteEvent(isMuted: true);
    final result = handleInternalEvent(state.value!, event);
    state = AsyncData(result.newState);
    _processCommands(result.commands);
  }

  void unmuteSpeaker() {
    final event = SpeakerMuteEvent(isMuted: false);
    final result = handleInternalEvent(state.value!, event);
    state = AsyncData(result.newState);
    _processCommands(result.commands);
  }

  void skipAssistant() {
    final lastAssistantItemId = state.value!.itemIdToLoomStep.entries
        .lastWhereOrNull((e) =>
            e.value is LoomChatCompletionText &&
            (e.value as LoomChatCompletionText).role ==
                tlsnx_role.Role.assistant)
        ?.key;

    if (lastAssistantItemId == null) {
      liveLogError(() => 'No assistant item found to skip');
      return;
    }
    _skippedAssistantItemIds.add(lastAssistantItemId);
    final event = AssistantSpeechStopRequestedEvent(lastAssistantItemId, 10000);
    final result = handleInternalEvent(state.value!, event);
    state = AsyncData(result.newState);
    _processCommands(result.commands);
  }

  Future<String> start({
    required MicrophoneService microphone,
    required String initialPrompt,
    required String? chatId,
  }) async {
    _lastAudioPlayed = BytesBuilder();
    _lastAudioFromAssistant = BytesBuilder();
    await _microphone?.stop();
    _micDataWhenRemoteTalkingBuffer = AudioRingBuffer(
        (MicrophoneConfig.kBitsPerSample / 8).ceil() * /* bytes per sample */
            MicrophoneConfig.kSampleRate /* samples per second */ *
            8 /* seconds */);
    _micDataWhenSocketConnectingBuffer = AudioRingBuffer(
        (MicrophoneConfig.kBitsPerSample / 8).ceil() * /* bytes per sample */
            MicrophoneConfig.kSampleRate /* samples per second */ *
            8 /* seconds */);
    _microphone = microphone;
    _pyannote ??= Pyannote.load(await getPyannoteModelPath());
    _pyannoteThrottler =
        ThrottledPublisher(interval: const Duration(milliseconds: 300));

    // Trigger SessionStartRequestedEvent
    final event = SessionStartRequestedEvent(
      initialPrompt,
      chatId,
      ref.read(apiKeysProvider),
    );
    final result = handleInternalEvent(state.value!, event);
    state = AsyncData(result.newState);
    _processCommands(result.commands);
    return state.valueOrNull!.conversationId!;
  }

  void stop() {
    final event = LiveStopCalledEvent();
    final result = handleInternalEvent(state.value!, event);
    state = AsyncData(result.newState);
    _processCommands(result.commands);
  }

  void _onSocketEvent(ServerEvent event) {
    final result = handleServerEvent(state.value!, event);
    state = AsyncData(result.newState);
    _processCommands(result.commands);
  }

  void _processCommands(List<DomainCommand> commands) async {
    if (commands.isEmpty) return;
    if (commands.length == 1 &&
        (commands.first is SendSocketMessageCommand ||
            // ignore: no-empty-block
            commands.first is PersistConversationCommand)) {
      // avoiding log spam from mic uploads temporarily while debugging
    } else {
      liveLog(() =>
          '_processCommands: ${commands.map((e) => e.runtimeType).join(', ')}');
    }
    for (final cmd in commands) {
      switch (cmd) {
        case CacheLastAssistantAudioCommand _:
          _lastAudioFromAssistant = _lastAudioPlayed;
          _lastAudioPlayed = BytesBuilder();
        case RefreshSessionInstructionsCommand _:
          _sendInitialPrompt();
        case ClearRemoteSpeakingMicBufferCommand _:
          _micDataWhenRemoteTalkingBuffer?.clear();
        case PlayAudioCommand c:
          if (_skippedAssistantItemIds.contains(c.itemId)) {
            liveLog(() => 'Skipping audio for item ${c.itemId}');
            return;
          }
          if (state.valueOrNull?.isSpeakerMutedByUser == true) {
            liveLog(() =>
                'Speaker muted by user, not playing audio for item ${c.itemId}');
            return;
          }
          _lastAudioPlayed.add(c.bytes);
          final Int16List samples = c.bytes.buffer.asInt16List();
          // `await` is crucial to avoid a ton of wasted time in debug mode on
          //  overflow logs that aren't shown. https://api.flutter.dev/flutter/dart-ui/ChannelBuffers-class.html
          // It is likely these logs may also be a real issue since they concern the
          // platform channel being backed up with messages and having to overwrite
          // ones that have not been processed.
          await FlutterPcmSound.feed(
              PcmArrayInt16(bytes: samples.buffer.asByteData()));
        case SetConversationCommand c:
          final maybeChatId = c.id;
          final convoMgr = ref.read(conversationManagerProvider.notifier);
          final existing =
              maybeChatId == null ? null : convoMgr.getById(maybeChatId);
          final Conversation chat;
          if (existing != null) {
            chat = existing;
          } else {
            chat = Conversation(
              id: LoomStep.newId(),
              modifiedAt: DateTime.now(),
              syncLevel:
                  ref.read(settingsNotifierProvider).syncLevelForNewObjects,
            );
            unawaited(convoMgr.upsert(chat));
          }
          state = AsyncData(state.value!.copyWith(conversationId: chat.id));
        case PersistConversationCommand c:
          final convoMgr = ref.read(conversationManagerProvider.notifier);
          final convo = convoMgr.getById(c.conversationId);
          if (convo == null) {
            liveLogError(() =>
                'PersistConversationCommand: Conversation not found for id ${c.conversationId}');
            return;
          }
          final List<LoomStep> makeEmptyMessagesInvisible = c.steps
              .map((e) {
                if (e is LoomChatCompletionText &&
                    e.isVisible &&
                    e.text.isEmpty) {
                  return e.copyWith(isVisible: false);
                }
                if (e is LoomChatCompletionText &&
                    e.text.isNotEmpty &&
                    !e.isVisible) {
                  return e.copyWith(isVisible: false);
                }
                if (e is LoomChatDialogue && e.text.isEmpty) {
                  return e.copyWith(isVisible: false);
                }
                if (e is LoomChatDialogue &&
                    e.text.isNotEmpty &&
                    !e.isVisible) {
                  return e.copyWith(isVisible: true);
                }
                return e;
              })
              .toList()
              .cast<LoomStep>();
          final newConvo = convo.mergeSteps(makeEmptyMessagesInvisible);
          unawaited(convoMgr.upsert(newConvo));
        case SendSocketMessageCommand c:
          // liveLog(() => '${c.message.runtimeType} sending socket message');
          _socket?.sendEvent(c.message);
        case TruncateAssistantMessageCommand c:
          _socket?.sendEvent(RealtimeEvent.conversationItemTruncate(
              itemId: c.itemId,
              contentIndex: c.contentIndex,
              audioEndMs: c.audioEndMs));
        case RunScriptCommand c:
          _executeRunScript(c.input, c.serverEvent);
        case ResetAudioSystemCommand _:
          _resetAiAudioPlayback();
        case ConnectSocketCommand c:
          await _connectSocket(c.apiKeys, c.isReconnect);
        case DisconnectSocketCommand _:
          await _disconnectSocket();
        case StartMicrophoneCommand c:
          _startMicrophone(c.config);
        case StopMicrophoneCommand _:
          await _stopMicrophone();
        case LoadPyannoteModelCommand _:
          await _loadPyannoteModel();
        case BufferMicDataWhenRemoteTalkingCommand c:
          _bufferMicDataWhenRemoteTalking(c.bytes);
        case BufferMicDataWhenSocketDisconnectedCommand c:
          _bufferMicDataWhenSocketDisconnected(c.bytes);
      }
    }
  }

  void _bufferMicDataWhenSocketDisconnected(Uint8List bytes) {
    liveLog(() => 'Buffering mic data while socket disconnected.');
    _micDataWhenSocketConnectingBuffer?.write(bytes);
  }

  void _bufferMicDataWhenRemoteTalking(Uint8List bytes) {
    _micDataWhenRemoteTalkingBuffer?.write(bytes);

    // ignore: avoid-passing-async-when-sync-expected
    _pyannoteThrottler!.publish(() async {
      if (state.value?.isRemoteTalking == false) {
        // liveLog(() =>
        // '[_bufferMicDataWhenRemoteTalking] Stream publisher checking *after* assistant stopped talking.');
        _micDataWhenRemoteTalkingBuffer?.clear();
        return;
      }
      final bufferedMicData =
          _micDataWhenRemoteTalkingBuffer!.getCurrentSamples();
      if (bufferedMicData.isEmpty) {
        // Needed because pyannote expects *some* audio, it does sublist
        // operations while calculating windows that ultimately leads to an
        // error of "Failed to process audio in isolate" from FONNX due to
        // a List.sublist() call in FONNX.
        liveLogError(
            () => '[_bufferMicDataWhenRemoteTalking] No data to process');
        return;
      }
      final floats = resampleTo16k(pcmSamplesToFloat32(bufferedMicData));
      final durationInMs = (bufferedMicData.length /
              (MicrophoneConfig.kBitsPerSample / 8) /
              MicrophoneConfig.kSampleRate *
              1000)
          .round();
      final crosstalkResultAsMap = await _pyannote!.process(floats);
      final howLongRemoteTalkingMs = _micDataWhenRemoteTalkingBuffer!
              .bytesWrittenSinceLastClear
              .toDouble() ~/
          (MicrophoneConfig.kBitsPerSample / 8) ~/
          MicrophoneConfig.kSampleRate *
          1000;
      final micDataSizeInMs = _micDataWhenRemoteTalkingBuffer!.capacityInBytes /
          (MicrophoneConfig.kBitsPerSample / 8) /
          MicrophoneConfig.kSampleRate *
          1000;
      if (!state.value!.didRemoteTalkInFirstSecond &&
          howLongRemoteTalkingMs < micDataSizeInMs) {
        final hasRemoteTalkedAtLeast2Seconds = howLongRemoteTalkingMs > 2000;
        final anySegmentInFirstSecond =
            crosstalkResultAsMap.any((s) => s['start'] < 1000);
        if (hasRemoteTalkedAtLeast2Seconds && anySegmentInFirstSecond) {
          // liveLog(() =>
          // 'Remote talked for at least 2 seconds and has segments in the first second. From here on out, we assume there is echo, therefore the user cannot be using headphones.');
          state = AsyncData(
              state.value!.copyWith(didRemoteTalkInFirstSecond: true));
        }
      }
      final crosstalkResult = crosstalkResultAsMap.detectCrosstalk(
        areHeadphonesPossible: !state.value!.didRemoteTalkInFirstSecond,
        isDefinitelyHeadphones: state.value!.isSpeakerMutedByUser,
      );
      // liveLog(() =>
      // 'Crosstalk result: $crosstalkResult. Duration analyzed: $durationInMs ms. As map: $crosstalkResultAsMap');
      // ignore: dead_code
      if (false && durationInMs > 2000) {
        unawaited(writeWavFile(
            '${PathProvider.documentsDirectory}/${crosstalkResult.hasCrosstalk}--${DateTime.now().millisecondsSinceEpoch}.wav',
            floats,
            sampleRate: 16000));
      }
      if (crosstalkResult.hasCrosstalk) {
        if (durationInMs < 2000) {
          // We don't have confidence the models results will be stable &
          // accurate until it has at least 2 seconds of audio.
          liveLog(() => 'Crosstalk detected but duration too short');
          return;
        }

        final audioEndMs = (crosstalkResult.crosstalkStartTime! *
                MicrophoneConfig.kSampleRate /
                MicrophoneConfig.kSampleRateOpenAiRealtime)
            .round();

        final result = handleInternalEvent(
          state.valueOrNull!,
          CrosstalkDetectedEvent(
            audioEndMs: audioEndMs,
            bufferedMicData: bufferedMicData.buffer.asUint8List(),
          ),
        );
        state = AsyncData(result.newState);
        _processCommands(result.commands);
      }
    });
  }

  void _executeRunScript(LoomChatCompletionFunction input,
      ResponseFunctionCallArgumentsDoneEvent serverEvent) async {
    final isScriptSelect =
        input.functionName == AgentScriptSelect.selectScriptToolName;
    final isScriptRun = input.functionName == kUseScriptToolname;
    final Script script;
    if (isScriptRun) {
      final scriptId = ref.read(companionScriptIdProvider);
      final maybeScript = ref
          .read(scriptManagerProvider)
          .valueOrNull
          ?.firstWhereOrNull((e) => e.id == scriptId);
      if (maybeScript == null) {
        liveLogError(() =>
            'Script with id $scriptId not found. Script manager: ${ref.read(scriptManagerProvider).valueOrNull}');
        return;
      }
      script = maybeScript;
    } else if (isScriptSelect) {
      script = AgentScriptSelect.scriptForFunctionArguments(
          input, ref.read(scriptManagerProvider).valueOrNull ?? []);
    } else {
      liveLogError(
          () => 'Cannot determine script for RunScriptCommand: $input');
      return;
    }

    final runManager = ref.read(runManagerProvider.notifier);
    final streamController = await runManager.run(
      imageMessages: [],
      companionFiles: ref.read(companionFilesObjectProvider),
      conversationId: state.value!.conversationId!,
      preprompt: [],
      prompt: runPromptStep(
          input.functionArguments[kScriptSelectPromptFunctionKey] ?? ''),
      script: script,
      history: ref.read(companionContextInputProvider),
      conversationManager: ref.read(conversationManagerProvider.notifier),
      documentManager: ref.read(documentManagerProvider.notifier),
      initialStatus: RunStatusPickingScript(script.label),
      // Getting the rewritten prompt from the tool output helps significantly,
      // the user prompt transcript often is not available yet when Run starts.
      //
      // However, the rewritten prompt is also added to the chat history. Lets
      // make it invisible so there's no duplication.
      isPromptVisible: false,
    );
    if (streamController == null) {
      loggy.error('Run manager did not return a stream controller');
      return;
    }
    final lastRunOutput = await streamController.stream.last;
    final visibleSteps = lastRunOutput.steps
        .whereType<LoomChatCompletionText>()
        .where((e) => e.isVisible && e.text.isNotEmpty)
        .toList();
    final concatenatedOutput =
        visibleSteps.map((e) => '${e.role.label}: ${e.text}').join('\n\n');
    final socket = _socket;
    if (socket == null) {
      loggy.error('Socket is null');
      return;
    }

    final injectedResponseCreatedItemId =
        '$kInjectedStepIdPrefix${DateTime.now().millisecondsSinceEpoch}';
    socket.sendEvent(
      ConversationItemCreateEvent(
        item: Item.functionCallOutput(
          id: injectedResponseCreatedItemId,
          callId: serverEvent.callId,
          output: '# Output\n$concatenatedOutput\n\n\n',
        ),
      ),
    );
    // TODO: Adding this to test if it prevents cases where it apologizes and does the search again
    // after the script has been run and we add these steps.
    await socket.events.firstWhere((e) => e is ConversationItemCreatedEvent);
    final injectedRequestResponseEventId =
        '$kInjectedStepIdPrefix${DateTime.now().millisecondsSinceEpoch}';
    final eventToSend = ResponseCreateEvent(
      eventId: injectedRequestResponseEventId,
      response: Response(
        // id of response has to be null for ResponseCreateEvent, otherwise,
        // OpenAI server errors.
        id: null,
        // Can only be audio & text, or, text. Audio only is not allowed, which
        // is unfortunate, as it would allow us to avoid injecting another text
        // response that copies the existing script output.
        modalities: ['audio', 'text'],
        maxOutputTokens: 4096,
        temperature:
            createDefaultSession(initialPrompt: '', ref: ref).temperature,
        instructions:
            '# Task \nAbove is the output of the script. Please vocally respond to the user.',
      ),
    );
    socket.sendEvent(eventToSend);
  }

  void _resetAiAudioPlayback() {
    if (P.isTesting) {
      // Avoid MissingPluginException(No implementation found for method setLogLevel on channel flutter_pcm_sound/methods)
      // affects ~30 of 1200 tests.
      return;
    }
    if (state.valueOrNull?.isRemoteTalking == true) {
      liveLog(() => '[LIVE] _resetAudioSystem isRemoteTalking = false');
      state = AsyncData(state.value!.copyWith(isRemoteTalking: false));
    }
    liveLog(() => 'Resetting audio system');
    FlutterPcmSound.release();
    FlutterPcmSound.setup(
      sampleRate: 24000,
      channelCount: 1,
      iosAudioCategory: IosAudioCategory.playAndRecord,
    );
    // This is delicate - the feed callback is how we decide when the assistant
    // stops talking. If we set a higher feed threshold, say, 8000 frames,
    // there's significant slippage between the actual speech stop and when we
    // expect it to stop. We don't have a back buffer to feed audio anyway,
    // we give it as soon as we get it, so might as well use feed threshold 0.
    FlutterPcmSound.setFeedThreshold(0);
    FlutterPcmSound.setLogLevel(LogLevel.none);
    FlutterPcmSound.setFeedCallback((index) {
      if (state.value!.isRemoteTalking) {
        liveLog(() => 'Remote talking will end in 450 ms.');

        Future.delayed(const Duration(milliseconds: 450), () {
          state = AsyncData(state.value!.copyWith(isRemoteTalking: false));
          liveLog(() => 'Remote talking ended');
        });
      }
    });
  }

  Future<void> _connectSocket(ApiKeys apiKeys, bool isReconnect) async {
    try {
      state = AsyncData(state.value!.copyWith(isSocketConnected: false));
      await _socket?.disconnect();
      _socket = RealtimeSocket();
      _socketEventsSubscription = _socket!.events.listen(_onSocketEvent,
          onDone: _onSocketDone, onError: _onSocketError);
      await _socket?.connect(apiKeys);
      await _socket?.events.firstWhere((e) => e is SessionCreatedEvent);
      _sendInitialPrompt();
      await _socket?.events.firstWhere((e) => e is SessionUpdatedEvent);
      if (isReconnect && _lastAudioFromAssistant.isNotEmpty) {
        liveLog(() =>
            'Reconnecting, sending last assistant audio to condition voice to match.');
        _socket?.sendEvent(
          ConversationItemCreateEvent(
            item: MessageItem(
              id: LiveNotifier.kInjectedStepIdPrefix +
                  DateTime.now().millisecondsSinceEpoch.toString(),
              role: Role.user,
              content: [
                Content(
                  type: ContentType.inputAudio,
                  audio: base64Encode(_lastAudioFromAssistant.takeBytes()),
                ),
                Content(
                    type: ContentType.inputText,
                    text:
                        'That audio is the last clip from the assistant before we disconnected, the regional dialect/accent should match that clip. Don\'t mention the reconnection, or re-greet the user. Stay silent, if you have to say something, just note there was a hiccup there.'),
              ],
            ),
          ),
        );
        await _socket?.events
            .firstWhere((e) => e is ConversationItemCreatedEvent);
      }
      liveLog(() => 'Socket connected.');
      final backAudio = _micDataWhenSocketConnectingBuffer
          ?.getCurrentSamples()
          .buffer
          .asUint8List();
      if (backAudio != null && backAudio.isNotEmpty) {
        liveLog(() =>
            'Sending mic audio while connecting. Duration: ${(backAudio.length / 2.0 / MicrophoneConfig.kSampleRateOpenAiRealtime).toStringAsFixed(3)} s');
        final command = MicDataEvent(backAudio);
        final result = handleInternalEvent(state.value!, command);
        state = AsyncData(result.newState);
        _processCommands(result.commands);
      } else {
        liveLog(() =>
            'No mic audio to send while connecting. _micDataWhenSocketConnectingBuffer: $_micDataWhenSocketConnectingBuffer.');
      }
      liveLog(() => 'Updating state to indictate socket connected.');
      state = AsyncData(state.value!.copyWith(isSocketConnected: true));
    } catch (e, s) {
      liveLogError(() => 'Error connecting socket: $e\n$s');
      state = AsyncData(state.value!.copyWith(isSocketConnected: false));
      stop();
    }
  }

  void _onSocketDone() {
    liveLog(() => 'Socket connection closed by the server or completed');
    // If the session is still running, attempt to reconnect.
    if (state.value?.isRunning == true) {
      liveLog(() => 'Session is running, attempting to reconnect...');
      final apiKeys = ref.read(apiKeysProvider);
      final result =
          handleInternalEvent(state.value!, ReconnectRequestedEvent(apiKeys));
      state = AsyncData(result.newState);
      _processCommands(result.commands);
    } else {
      // Not running, just mark socket as disconnected.
      state = AsyncData(state.value!.copyWith(isSocketConnected: false));
    }
  }

  void _onSocketError(dynamic error, StackTrace stack) {
    liveLogError(() => 'Socket error encountered: $error\n$stack');

    // If we are running, treat this as a lost connection and attempt to reconnect.
    if (state.value?.isRunning == true) {
      liveLogError(() => 'Socket error while running. Attempting reconnect...');
      final apiKeys = ref.read(apiKeysProvider);
      final result =
          handleInternalEvent(state.value!, ReconnectRequestedEvent(apiKeys));
      state = AsyncData(result.newState);
      _processCommands(result.commands);
    } else {
      // Not running, just mark socket as disconnected and log the error.
      state = AsyncData(state.value!.copyWith(isSocketConnected: false));
    }
  }

  Future<void> _disconnectSocket() async {
    await _socket?.disconnect();
    await _socketEventsSubscription?.cancel();
    _socket = null;
    _socketEventsSubscription = null;
  }

  void _startMicrophone(MicrophoneConfig config) {
    final micStream = _microphone!.getMicrophoneStream(config);
    _microphoneSubscription = micStream.listen((bytes) {
      final command = MicDataEvent(bytes);
      final result = handleInternalEvent(state.value!, command);
      state = AsyncData(result.newState);
      _processCommands(result.commands);
    }, onError: (e, s) {
      liveLogError(() => 'Error in microphone stream: $e\n$s');
    }, onDone: () {
      liveLog(() => 'Microphone stream done');
    });
  }

  Future<void> _stopMicrophone() async {
    await _microphone?.stop();
    await _microphoneSubscription?.cancel();
    _microphone = null;
    _microphoneSubscription = null;
  }

  Future<void> _loadPyannoteModel() async {
    final modelPath = await getPyannoteModelPath();
    _pyannote = Pyannote.load(modelPath);
  }

  void _sendInitialPrompt() async {
    Future<String> getInitialPromptWithChatHistory() async {
      final initialPrompt = StringBuffer();
      initialPrompt
          .write(ref.read(settingsNotifierProvider).agent.liveInstructions);
      final chatId = state.value?.conversationId;
      final String historyString;
      if (chatId == null) {
        loggy.error('Cannot include history; no conversation ID found');
        historyString = '';
      } else {
        final chat =
            ref.read(conversationManagerProvider.notifier).getById(chatId);
        if (chat == null) {
          loggy.error('Cannot include history; conversation not found');
          historyString = '';
        } else {
          final contextInput = ref.read(companionContextInputProvider);
          historyString =
              await liveBuildHistoryString(chat, contextInput.tokenCount);
        }
      }

      liveLog(() => 'Got chat history string: $historyString');

      if (historyString.trim().isNotEmpty) {
        initialPrompt.write(historyString);
      }
      return initialPrompt.toString();
    }

    final initialPrompt = await getInitialPromptWithChatHistory();
    final updateEvent = SessionUpdateEvent(
      session: createDefaultSession(
        initialPrompt: initialPrompt,
        ref: ref,
      ),
    );
    _socket?.sendEvent(updateEvent);
  }

  static Session createDefaultSession(
      {required String initialPrompt, required Ref ref}) {
    final scriptId = ref.read(companionScriptIdProvider);
    final isAuto = scriptId == null;
    final ToolFunction? tool;
    if (isAuto) {
      tool = selectScriptTool(
          '', ref.read(scriptManagerProvider).valueOrNull ?? []);
    } else {
      final script = ref.read(scriptManagerProvider.notifier).getById(scriptId);
      if (script == null) {
        liveLogError(() =>
            'script not found for ID: $scriptId. Cannot use tool, will not update session.');
        tool = null;
      } else {
        tool = useScriptTool(script);
      }
    }

    final session = Session(
      voice: Voice.coral,
      instructions: initialPrompt,
      temperature: 0.8,
      tools: tool == null
          ? []
          : [tool]
              .map(
                (t) => Tool(
                  type: ToolType.function,
                  name: t.name,
                  description: t.description,
                  parameters: jsonDecode(t.parametersAsString),
                ),
              )
              .toList(),
      toolChoice: 'auto',
      maxResponseOutputTokens: 'inf',
      inputAudioTranscription: AudioTranscriptionConfig(model: 'whisper-1'),
      modalities: ['text', 'audio'],
      inputAudioFormat: AudioFormat.pcm16,
      outputAudioFormat: AudioFormat.pcm16,
      turnDetection: TurnDetectionConfig(
        turnType: 'server_vad',
        threshold: 0.5,
        prefixPaddingMs: 300,
        silenceDurationMs: 1000,
      ),
    );
    return session;
  }
}


// live_internal_handlers.dart
import 'dart:convert';

import 'package:telosnex/chat/models/role.dart';
import 'package:telosnex/chat/realtime/realtime_models.dart' hide Role;
import 'package:telosnex/models/loom_step.dart';
import 'package:telosnex/stt/microphone_service.dart';

import 'live_architecture.dart';
import 'live_internal_events.dart';
import 'live.dart'; // LiveState defined in live.dart
import 'package:collection/collection.dart';

HandlerResult handleInternalEvent(LiveState currentState, InternalEvent event) {
  if (event is! MicDataEvent && event is! MicDataWhenRemoteTalkingEvent) {
    liveLog(() => 'handleInternalEvent: $event');
  }
  return switch (event) {
    CrosstalkDetectedEvent e => _handleCrosstalkDetected(currentState, e),
    RemoteTalkingEndedEvent e => _handleRemoteTalkingEnded(currentState, e),
    ScriptChangedEvent e => _handleScriptChanged(currentState, e),
    AssistantSpeechStopRequestedEvent e =>
      _handleAssistantSpeechStopRequested(currentState, e),
    SessionStartRequestedEvent e =>
      _handleSessionStartRequested(currentState, e),
    LiveStopCalledEvent e => _handleSessionStopRequested(currentState, e),
    ReconnectRequestedEvent e => _handleReconnectRequested(currentState, e),
    MicDataEvent e => _handleMicDataEvent(currentState, e),
    MicDataWhenRemoteTalkingEvent e =>
      _handleMicDataWhenRemoteTalking(currentState, e),
    MicDataWhenSocketDisconnectedEvent e =>
      _handleMicDataWhenSocketDisconnected(currentState, e),
    MicMuteEvent e => HandlerResult(currentState.copyWith(isMicMutedByUser: e.isMuted), []),
    SpeakerMuteEvent e => HandlerResult(currentState.copyWith(isSpeakerMutedByUser: e.isMuted), []),
  };
}

HandlerResult _handleMicDataEvent(LiveState state, MicDataEvent event) {
  if (!state.isSocketConnected) {
    return handleInternalEvent(state, MicDataWhenSocketDisconnectedEvent(event.data));
  }

  if (state.isMicMutedByUser) {
    return HandlerResult(state);
  }

  if (state.isRemoteTalking) {
    return handleInternalEvent(
        state, MicDataWhenRemoteTalkingEvent(event.data));
  }
  final commands = <DomainCommand>[
    SendSocketMessageCommand(
        InputAudioBufferAppendEvent(audio: base64Encode(event.data)))
  ];
  return HandlerResult(state, commands);
}

HandlerResult _handleMicDataWhenRemoteTalking(
    LiveState state, MicDataWhenRemoteTalkingEvent event) {
  if (!state.isRemoteTalking) {
    liveLogError(() =>
        'Called _handleMicDataWhenRemoteTalking when not remote talking. Passing to _handleMicDataEvent');
    return handleInternalEvent(state, MicDataEvent(event.data));
  }
  final commands = <DomainCommand>[
    BufferMicDataWhenRemoteTalkingCommand(event.data)
  ];
  return HandlerResult(state, commands);
}

HandlerResult _handleMicDataWhenSocketDisconnected(
    LiveState state, MicDataWhenSocketDisconnectedEvent event) {
  if (state.isSocketConnected) {
    liveLogError(() =>
        'Called _handleMicDataWhenSocketDisconnected when socket is connected. Passing to _handleMicDataEvent');
    return handleInternalEvent(state, MicDataEvent(event.data));
  }
  final commands = <DomainCommand>[
    BufferMicDataWhenSocketDisconnectedCommand(event.data)
  ];
  liveLog(() => 'LiveInternalHandlers._handleMicDataWhenSocketDisconnected: '
      'Buffering ${event.data.length} bytes');
  return HandlerResult(state, commands);
}

HandlerResult _handleCrosstalkDetected(
    LiveState state, CrosstalkDetectedEvent event) {
  final lastAssistantItemId = state.itemIdToLoomStep.entries
      .firstWhereOrNull((e) =>
          e.value is LoomChatCompletionText &&
          (e.value as LoomChatCompletionText).role == Role.assistant)
      ?.key;

  final commands = <DomainCommand>[];
  commands.add(ResetAudioSystemCommand());

  if (lastAssistantItemId != null) {
    commands.add(TruncateAssistantMessageCommand(
        lastAssistantItemId, 0, event.audioEndMs));
  } else {
    liveLogError(
        () => '[_handleCrosstalkDetected] No assistant item found to truncate');
  }

  commands.add(ClearRemoteSpeakingMicBufferCommand());

  if (event.bufferedMicData.isNotEmpty) {
    final data = base64Encode(event.bufferedMicData);
    commands.add(
        SendSocketMessageCommand(InputAudioBufferAppendEvent(audio: data)));
  }
  final newState = state.copyWith(isRemoteTalking: false);
  return HandlerResult(newState, commands);
}

HandlerResult _handleRemoteTalkingEnded(
    LiveState state, RemoteTalkingEndedEvent _) {
  final commands = <DomainCommand>[
    ClearRemoteSpeakingMicBufferCommand(),
  ];
  return HandlerResult(state.copyWith(isRemoteTalking: false), commands);
}

HandlerResult _handleScriptChanged(LiveState state, ScriptChangedEvent event) {
  return HandlerResult(
      state, [RefreshSessionInstructionsCommand(event.scriptId)]);
}

HandlerResult _handleAssistantSpeechStopRequested(
    LiveState state, AssistantSpeechStopRequestedEvent event) {
  final commands = <DomainCommand>[
    TruncateAssistantMessageCommand(event.assistantItemId, 0, event.audioEndMs),
    ResetAudioSystemCommand(),
    ClearRemoteSpeakingMicBufferCommand(),
  ];
  return HandlerResult(state.copyWith(isRemoteTalking: false), commands);
}

HandlerResult _handleSessionStartRequested(
    LiveState state, SessionStartRequestedEvent event) {
  // Assume apiKeys and pyannotePath are stored in state or known
  final apiKeys = event.apiKeys;

  final commands = <DomainCommand>[
    // set conversation *must* come before any commands that may run
    // asynchronously. start() calls through here, and expects that
    // state.conversationId is set when it runs.
    SetConversationCommand(event.chatId),
    StartMicrophoneCommand(MicrophoneConfig.openAiRealtime()),
    ConnectSocketCommand(apiKeys, false),
    LoadPyannoteModelCommand(),
    ResetAudioSystemCommand(),
  ];
  final newState =
      state.copyWith(isRunning: true, conversationId: event.chatId);
  return HandlerResult(newState, commands);
}

HandlerResult _handleSessionStopRequested(
    LiveState _, LiveStopCalledEvent __) {
  final commands = [
    DisconnectSocketCommand(DisconnectReason.remoteDisconnectWithEvent),
    StopMicrophoneCommand(),
    ResetAudioSystemCommand()
  ];

  return HandlerResult(
      LiveState.initial(),
      commands);
}

HandlerResult _handleReconnectRequested(
    LiveState state, ReconnectRequestedEvent event) {
  final commands = [
    ConnectSocketCommand(event.apiKeys, true),
  ];
  return HandlerResult(state, commands);
}

// live_internal_events.dart
// Internal events for non-server conditions

import 'dart:typed_data';

import 'package:telosnex/models/api_keys.dart';
import 'package:telosnex/stt/microphone_service.dart';

sealed class InternalEvent {}

class CrosstalkDetectedEvent extends InternalEvent {
  final int audioEndMs;
  final Uint8List bufferedMicData;
  CrosstalkDetectedEvent(
      {required this.audioEndMs, required this.bufferedMicData});
}

class MicDataEvent extends InternalEvent {
  final Uint8List data;
  MicDataEvent(this.data);

  @override
  String toString() =>
      'MicDataEvent(${(data.length / 2.0 / MicrophoneConfig.kSampleRateOpenAiRealtime).toStringAsFixed(3)} seconds, ${data.length} bytes)';
}

class MicDataWhenRemoteTalkingEvent extends InternalEvent {
  final Uint8List data;
  MicDataWhenRemoteTalkingEvent(this.data);
}

class MicDataWhenSocketDisconnectedEvent extends InternalEvent {
  final Uint8List data;
  MicDataWhenSocketDisconnectedEvent(this.data);
}

class RemoteTalkingEndedEvent extends InternalEvent {}

class ScriptChangedEvent extends InternalEvent {
  final String? scriptId;
  ScriptChangedEvent(this.scriptId);
}

class AssistantSpeechStopRequestedEvent extends InternalEvent {
  final String assistantItemId;
  final int audioEndMs;
  AssistantSpeechStopRequestedEvent(this.assistantItemId, this.audioEndMs);
}

// Lifecycle events
class SessionStartRequestedEvent extends InternalEvent {
  final String initialPrompt;
  final String? chatId;
  final ApiKeys apiKeys;
  SessionStartRequestedEvent(this.initialPrompt, this.chatId, this.apiKeys);
}

class LiveStopCalledEvent extends InternalEvent {}

class ReconnectRequestedEvent extends InternalEvent {
  final ApiKeys apiKeys;
  ReconnectRequestedEvent(this.apiKeys);
}

class SpeakerMuteEvent extends InternalEvent {
  final bool isMuted;
  SpeakerMuteEvent({required this.isMuted});
}

class MicMuteEvent extends InternalEvent {
  final bool isMuted;
  MicMuteEvent({required this.isMuted});
}