import React, { createContext, useContext, useReducer } from "react";
import { setDispatch } from "./streamManager";

const AppContext = createContext();

const initialState = {
  token: localStorage.getItem("token"),
  user: JSON.parse(localStorage.getItem("user") || "null"),
  chats: [],
  activeChatId: null,
  chatMessages: {},
  activeStreams: {},
};

function reducer(state, action) {
  switch (action.type) {
    case "LOGIN":
      localStorage.setItem("token", action.token);
      localStorage.setItem("user", JSON.stringify(action.user));
      return { ...state, token: action.token, user: action.user };
    case "LOGOUT":
      localStorage.removeItem("token");
      localStorage.removeItem("user");
      return { ...initialState, token: null, user: null };
    case "SET_CHATS":
      return { ...state, chats: action.chats };
    case "SET_ACTIVE_CHAT":
      return { ...state, activeChatId: action.chatId };
    case "ADD_CHAT":
      return { ...state, chats: [action.chat, ...state.chats] };
    case "UPDATE_CHAT": {
      const chats = state.chats.map((c) =>
        c.id === action.chat.id ? { ...c, ...action.chat } : c
      );
      return { ...state, chats };
    }
    case "REMOVE_CHAT": {
      const chats = state.chats.filter((c) => c.id !== action.chatId);
      const msgs = { ...state.chatMessages };
      delete msgs[action.chatId];
      return {
        ...state, chats, chatMessages: msgs,
        activeChatId: state.activeChatId === action.chatId ? null : state.activeChatId,
      };
    }
    case "SET_MESSAGES":
      return { ...state, chatMessages: { ...state.chatMessages, [action.chatId]: action.messages } };
    case "ADD_MESSAGE": {
      const existing = state.chatMessages[action.chatId] || [];
      return { ...state, chatMessages: { ...state.chatMessages, [action.chatId]: [...existing, action.message] } };
    }
    case "SET_STREAMING":
      return { ...state, activeStreams: { ...state.activeStreams, [action.chatId]: action.streaming } };
    default:
      return state;
  }
}

export function AppProvider({ children }) {
  const [state, dispatch] = useReducer(reducer, initialState);
  setDispatch(dispatch);
  return (
    <AppContext.Provider value={{ state, dispatch }}>
      {children}
    </AppContext.Provider>
  );
}

export function useApp() {
  return useContext(AppContext);
}