import { streamMessage } from "./api";

const _streams = new Map();
const _listeners = new Map();
let _dispatch = null;

export function setDispatch(dispatch) { _dispatch = dispatch; }

export function getStreamData(chatId) {
  const s = _streams.get(chatId);
  if (!s) return { streaming: false, text: "", thinking: "", isThinking: false };
  return { streaming: true, text: s.text, thinking: s.thinking, isThinking: s.isThinking };
}

export function isStreaming(chatId) { return _streams.has(chatId); }

export function subscribe(chatId, cb) {
  if (!_listeners.has(chatId)) _listeners.set(chatId, new Set());
  _listeners.get(chatId).add(cb);
  return () => { const s = _listeners.get(chatId); if (s) { s.delete(cb); if (!s.size) _listeners.delete(chatId); } };
}

function _notify(id) { const s = _listeners.get(id); if (s) s.forEach((cb) => cb()); }

export function abortStream(chatId) {
  const s = _streams.get(chatId);
  if (s) { s.abortController.abort(); _streams.delete(chatId); _notify(chatId); if (_dispatch) _dispatch({ type: "SET_STREAMING", chatId, streaming: false }); }
}

export function startStream({ token, chatId, body }) {
  if (_streams.has(chatId)) return;
  const ac = new AbortController();
  _streams.set(chatId, { text: "", thinking: "", isThinking: false, abortController: ac });
  if (_dispatch) _dispatch({ type: "SET_STREAMING", chatId, streaming: true });
  _notify(chatId);

  (async () => {
    const s = _streams.get(chatId);
    if (!s) return;
    let usage = {}, msgId = "";
    try {
      for await (const evt of streamMessage(token, chatId, body, ac.signal)) {
        if (ac.signal.aborted || !_streams.has(chatId)) break;
        switch (evt.type) {
          case "thinking_start": s.isThinking = true; _notify(chatId); break;
          case "thinking_delta": s.thinking += evt.content; _notify(chatId); break;
          case "thinking_end": s.isThinking = false; _notify(chatId); break;
          case "text_delta": s.text += evt.content; _notify(chatId); break;
          case "usage": usage = { input_tokens: evt.input_tokens, output_tokens: evt.output_tokens }; break;
          case "title_update": if (_dispatch) _dispatch({ type: "UPDATE_CHAT", chat: { id: chatId, title: evt.title } }); break;
          case "done": msgId = evt.message_id; break;
          case "status": break; // Web search status — ignore in stream
          case "error": s.text += `\n\n**Error:** ${evt.message}`; _notify(chatId); break;
        }
      }
      if (!ac.signal.aborted && _dispatch) {
        _dispatch({ type: "ADD_MESSAGE", chatId, message: {
          id: msgId || `gen-${Date.now()}`, role: "assistant", content: s.text,
          thinking_content: s.thinking || null, input_tokens: usage.input_tokens || 0,
          output_tokens: usage.output_tokens || 0, created_at: new Date().toISOString(), attachments: [],
        }});
      }
    } catch (err) {
      if (!ac.signal.aborted && _dispatch) {
        _dispatch({ type: "ADD_MESSAGE", chatId, message: {
          id: `err-${Date.now()}`, role: "assistant", content: `**Error:** ${err.message}`,
          created_at: new Date().toISOString(), attachments: [],
        }});
      }
    } finally {
      _streams.delete(chatId); _notify(chatId);
      if (_dispatch) _dispatch({ type: "SET_STREAMING", chatId, streaming: false });
    }
  })();
}