import { useCallback, useEffect, useRef, useState } from "react";
import SocketIO from "socket.io-client";
import { auth } from "../../firebase/firebase.utils";
import { store } from "../../redux/store";
import { addAlerts, setWssStatus } from "../../redux/application/application.actions";
import client from "../../utils/GraphQLClient";
import { useNotification } from "../Notifications/notificationContext.jsx";
import {
  GET_NOTIFICATIONS,
  GET_UNREAD_COUNT,
  MARK_ALL_NOTIFICATIONS_READ,
  MARK_NOTIFICATION_READ,
  UPDATE_NOTIFICATIONS_READ_FRAGMENT
} from "../../graphql/notifications.queries.js";
import { useMutation } from "@apollo/client/react";
import { useTranslation } from "react-i18next";
import { FEATURE_FLAGS_CHANGED_EVENT, useTreatmentsWithConfig } from "../../feature-flags/splitio-react-replacement";
import { INITIAL_NOTIFICATIONS, SocketContext } from "./useSocket.js";

const LIMIT = INITIAL_NOTIFICATIONS;
const TOKEN_SYNC_INTERVAL_MS = 10 * 60 * 1000;

/**
 * Socket Provider - Scenario Notifications / Web Socket related items
 * @param children
 * @param bodyshop
 * @param navigate
 * @param currentUser
 * @returns {JSX.Element}
 * @constructor
 */
const SocketProvider = ({ children, bodyshop, navigate, currentUser }) => {
  const socketRef = useRef(null);
  const tokenSyncIntervalRef = useRef(null);
  const [clientId, setClientId] = useState(null);
  const [isConnected, setIsConnected] = useState(false);
  const notification = useNotification();
  const userAssociationId = bodyshop?.associations?.[0]?.id;
  const { t } = useTranslation();

  const {
    treatments: { Realtime_Notifications_UI }
  } = useTreatmentsWithConfig({
    attributes: {},
    names: ["Realtime_Notifications_UI"],
    splitKey: bodyshop?.imexshopid
  });

  const [markNotificationRead] = useMutation(MARK_NOTIFICATION_READ, {
    update: (cache, { data: { update_notifications } }) => {
      const timestamp = new Date().toISOString();
      const updatedNotification = update_notifications.returning[0];

      cache.modify({
        fields: {
          notifications(existing = [], { readField }) {
            return existing.map((notif) =>
              readField("id", notif) === updatedNotification.id
                ? {
                    ...notif,
                    read: timestamp
                  }
                : notif
            );
          }
        }
      });

      const unreadCountQuery = cache.readQuery({
        query: GET_UNREAD_COUNT,
        variables: { associationid: userAssociationId }
      });

      if (unreadCountQuery?.notifications_aggregate?.aggregate?.count > 0) {
        cache.writeQuery({
          query: GET_UNREAD_COUNT,
          variables: { associationid: userAssociationId },
          data: {
            notifications_aggregate: {
              ...unreadCountQuery.notifications_aggregate,
              aggregate: {
                ...unreadCountQuery.notifications_aggregate.aggregate,
                count: unreadCountQuery.notifications_aggregate.aggregate.count - 1
              }
            }
          }
        });
      }

      if (socketRef.current && isConnected) {
        socketRef.current.emit("sync-notification-read", {
          email: currentUser?.email,
          bodyshopId: bodyshop.id,
          notificationId: updatedNotification.id
        });
      }
    },
    onError: (err) =>
      console.error("MARK_NOTIFICATION_READ error:", {
        message: err?.message,
        stack: err?.stack
      })
  });

  const [markAllNotificationsRead] = useMutation(MARK_ALL_NOTIFICATIONS_READ, {
    variables: { associationid: userAssociationId },
    update: (cache) => {
      const timestamp = new Date().toISOString();
      cache.modify({
        fields: {
          notifications(existing = [], { readField }) {
            return existing.map((notif) =>
              readField("read", notif) === null && readField("associationid", notif) === userAssociationId
                ? { ...notif, read: timestamp }
                : notif
            );
          },
          notifications_aggregate() {
            return { aggregate: { count: 0, __typename: "notifications_aggregate_fields" } };
          }
        }
      });

      const baseWhereClause = { associationid: { _eq: userAssociationId } };
      const cachedNotifications = cache.readQuery({
        query: GET_NOTIFICATIONS,
        variables: { limit: INITIAL_NOTIFICATIONS, offset: 0, where: baseWhereClause }
      });

      if (cachedNotifications?.notifications) {
        cache.writeQuery({
          query: GET_NOTIFICATIONS,
          variables: { limit: INITIAL_NOTIFICATIONS, offset: 0, where: baseWhereClause },
          data: {
            notifications: cachedNotifications.notifications.map((notif) =>
              notif.read === null ? { ...notif, read: timestamp } : notif
            )
          }
        });
      }

      if (socketRef.current && isConnected) {
        socketRef.current.emit("sync-all-notifications-read", {
          email: currentUser?.email,
          bodyshopId: bodyshop.id
        });
      }
    },
    onError: (err) => console.error("MARK_ALL_NOTIFICATIONS_READ error:", err)
  });

  const reconnectSocket = useCallback(
    async ({ forceRefreshToken = true } = {}) => {
      const socketInstance = socketRef.current;
      if (!socketInstance || !auth.currentUser || !bodyshop?.id) return false;

      try {
        const token = await auth.currentUser.getIdToken(forceRefreshToken);
        socketInstance.auth = { token, bodyshopId: bodyshop.id };

        if (socketInstance.connected) {
          socketInstance.emit("update-token", { token, bodyshopId: bodyshop.id });
        }

        socketInstance.disconnect();
        socketInstance.connect();
        return true;
      } catch (error) {
        console.error("Socket reconnect failed:", error?.message || error);
        return false;
      }
    },
    [bodyshop?.id]
  );

  useEffect(() => {
    const initializeSocket = async (token) => {
      if (!bodyshop?.id || socketRef.current) return;

      const endpoint = import.meta.env.PROD ? import.meta.env.VITE_APP_AXIOS_BASE_API_URL : "";
      const socketInstance = SocketIO(endpoint, {
        path: "/wss",
        withCredentials: true,
        auth: { token, bodyshopId: bodyshop.id },
        reconnectionAttempts: Infinity,
        reconnectionDelay: 2000,
        reconnectionDelayMax: 60000
        // randomizationFactor: 0.5,
        // transports: ["websocket", "polling"], // Add this to prefer WebSocket with polling fallback
        // rememberUpgrade: true
      });

      socketRef.current = socketInstance;

      const handleBodyshopMessage = (message) => {
        if (!message?.type) return;
        switch (message.type) {
          case "alert-update":
            store.dispatch(addAlerts(message.payload));
            break;
          case "task-created":
          case "task-updated":
          case "task-deleted": {
            const payload = message.payload;
            const assignedToId = bodyshop?.employees?.find((e) => e.user_email === currentUser?.email)?.id;
            if (!assignedToId || payload.assigned_to !== assignedToId) return;

            const dueVars = {
              bodyshop: bodyshop?.id,
              assigned_to: assignedToId,
              order: [{ due_date: "asc" }, { created_at: "desc" }]
            };
            const noDueVars = {
              bodyshop: bodyshop?.id,
              assigned_to: assignedToId,
              order: [{ created_at: "desc" }],
              limit: LIMIT,
              offset: 0
            };

            const whereBase = {
              bodyshopid: { _eq: bodyshop?.id },
              assigned_to: { _eq: assignedToId },
              deleted: { _eq: false },
              completed: { _eq: false }
            };
            const whereDue = { ...whereBase, due_date: { _is_null: false } };
            const whereNoDue = { ...whereBase, due_date: { _is_null: true } };

            // Helper to invalidate a cache entry
            const invalidateCache = (fieldName, args) => {
              try {
                client.cache.evict({
                  id: "ROOT_QUERY",
                  fieldName,
                  args
                });
              } catch (error) {
                console.error("Error invalidating cache:", error);
              }
            };

            // Invalidate lists and aggregates based on event type
            if (message.type === "task-deleted" || message.type === "task-updated") {
              // Invalidate both lists and no due aggregate for deletes and updates
              invalidateCache("tasks", { where: whereDue, order_by: dueVars.order });
              invalidateCache("tasks", {
                where: whereNoDue,
                order_by: noDueVars.order,
                limit: noDueVars.limit,
                offset: noDueVars.offset
              });
              invalidateCache("tasks_aggregate", { where: whereNoDue });
            } else if (message.type === "task-created") {
              // For creates, invalidate the target list and no due aggregate if applicable
              const hasDue = !!payload.due_date;
              if (hasDue) {
                invalidateCache("tasks", { where: whereDue, order_by: dueVars.order });
              } else {
                invalidateCache("tasks", {
                  where: whereNoDue,
                  order_by: noDueVars.order,
                  limit: noDueVars.limit,
                  offset: noDueVars.offset
                });
                invalidateCache("tasks_aggregate", { where: whereNoDue });
              }
            }

            // Always invalidate the total count for all events (handles creates, deletes, updates including completions)
            invalidateCache("tasks_aggregate", { where: whereBase });

            // Garbage collect after evictions
            client.cache.gc();

            break;
          }
          default:
            break;
        }
      };

      const handleFeatureFlagsChanged = (payload) => {
        window.dispatchEvent(new CustomEvent(FEATURE_FLAGS_CHANGED_EVENT, { detail: payload }));
      };

      const syncCurrentTokenToSocket = async () => {
        try {
          if (!auth.currentUser || !bodyshop?.id) return;
          const token = await auth.currentUser.getIdToken();
          socketInstance.auth = { token, bodyshopId: bodyshop.id };
          socketInstance.emit("update-token", { token, bodyshopId: bodyshop.id });
        } catch (error) {
          console.error("Failed to sync token to socket:", error?.message || error);
        }
      };

      const forceRefreshAndSyncToken = async () => {
        try {
          if (!auth.currentUser || !bodyshop?.id) return;
          const token = await auth.currentUser.getIdToken(true);
          socketInstance.auth = { token, bodyshopId: bodyshop.id };
          socketInstance.emit("update-token", { token, bodyshopId: bodyshop.id });
        } catch (error) {
          console.error("Failed to force-refresh token for socket:", error?.message || error);
        }
      };

      const handleConnect = () => {
        socketInstance.emit("join-bodyshop-room", bodyshop.id);
        syncCurrentTokenToSocket();
        setClientId(socketInstance.id);
        setIsConnected(true);
        store.dispatch(setWssStatus("connected"));
      };

      const handleReconnect = () => {
        forceRefreshAndSyncToken();
        setIsConnected(true);
        store.dispatch(setWssStatus("connected"));
      };

      const handleTokenUpdated = ({ success, error }) => {
        if (success) return;
        const err = String(error || "");
        if (/stale token|id-token-expired/i.test(err)) {
          forceRefreshAndSyncToken();
        }
      };

      const handleConnectionError = (err) => {
        console.error("Socket connection error:", err);
        setIsConnected(false);
        if (err?.message?.includes("auth/id-token-expired")) {
          console.warn("Token expired, refreshing...");
          auth.currentUser?.getIdToken(true).then((newToken) => {
            socketInstance.auth = { token: newToken, bodyshopId: bodyshop.id };
            if (socketInstance.connected) {
              socketInstance.emit("update-token", { token: newToken, bodyshopId: bodyshop.id });
            }
            socketInstance.connect();
          });
        } else {
          store.dispatch(setWssStatus("error"));
        }
      };

      const handleDisconnect = (reason) => {
        console.warn("Socket disconnected:", reason);
        setIsConnected(false);
        store.dispatch(setWssStatus("disconnected"));
        if (!socketInstance.connected && reason !== "io server disconnect") {
          setTimeout(() => {
            if (socketInstance.disconnected) {
              console.log("Manually triggering reconnection...");
              socketInstance.connect();
            }
          }, 2000);
        }
      };

      const handleNotification = (data) => {
        if (Realtime_Notifications_UI?.treatment !== "on") {
          return;
        }

        const { jobId, jobRoNumber, notificationId, associationId, notifications } = data;
        if (associationId !== userAssociationId) return;

        const newNotification = {
          __typename: "notifications",
          id: notificationId,
          jobid: jobId,
          associationid: associationId,
          scenario_text: JSON.stringify(notifications.map((notif) => notif.body)),
          fcm_text: notifications.map((notif) => notif.body).join(". ") + ".",
          scenario_meta: JSON.stringify(notifications.map((notif) => notif.variables || {})),
          created_at: new Date(notifications[0].timestamp).toISOString(),
          read: null,
          job: {
            ro_number: jobRoNumber,
            id: jobId,
            __typename: "jobs"
          }
        };

        const baseVariables = {
          limit: INITIAL_NOTIFICATIONS,
          offset: 0,
          where: { associationid: { _eq: userAssociationId } }
        };

        try {
          const existingNotifications =
            client.cache.readQuery({
              query: GET_NOTIFICATIONS,
              variables: baseVariables
            })?.notifications || [];
          if (!existingNotifications.some((n) => n.id === newNotification.id)) {
            client.cache.writeQuery({
              query: GET_NOTIFICATIONS,
              variables: baseVariables,
              data: {
                notifications: [newNotification, ...existingNotifications].sort(
                  (a, b) => new Date(b.created_at) - new Date(a.created_at)
                )
              },
              broadcast: true
            });

            const unreadVariables = {
              ...baseVariables,
              where: { ...baseVariables.where, read: { _is_null: true } }
            };
            const unreadNotifications =
              client.cache.readQuery({
                query: GET_NOTIFICATIONS,
                variables: unreadVariables
              })?.notifications || [];
            if (newNotification.read === null && !unreadNotifications.some((n) => n.id === newNotification.id)) {
              client.cache.writeQuery({
                query: GET_NOTIFICATIONS,
                variables: unreadVariables,
                data: {
                  notifications: [newNotification, ...unreadNotifications].sort(
                    (a, b) => new Date(b.created_at) - new Date(a.created_at)
                  )
                },
                broadcast: true
              });
            }

            client.cache.modify({
              id: "ROOT_QUERY",
              fields: {
                notifications_aggregate(existing = { aggregate: { count: 0 } }) {
                  return {
                    ...existing,
                    aggregate: {
                      ...existing.aggregate,
                      count: existing.aggregate.count + (newNotification.read === null ? 1 : 0)
                    }
                  };
                }
              }
            });

            notification.info({
              title: (
                <div
                  onClick={() => {
                    markNotificationRead({ variables: { id: notificationId } })
                      .then(() => navigate(`/manage/jobs/${jobId}`))
                      .catch((e) => console.error(`Error marking notification read: ${e?.message || ""}`));
                  }}
                >
                  {t("notifications.labels.notification-popup-title", {
                    ro_number: jobRoNumber || t("general.labels.na")
                  })}
                </div>
              ),
              description: (
                <ul
                  className="notification-alert-unordered-list"
                  onClick={() => {
                    markNotificationRead({ variables: { id: notificationId } })
                      .then(() => navigate(`/manage/jobs/${jobId}`))
                      .catch((e) => console.error(`Error marking notification read: ${e?.message || ""}`));
                  }}
                >
                  {notifications.map((notif, index) => (
                    <li className="notification-alert-unordered-list-item" key={index}>
                      {notif.body}
                    </li>
                  ))}
                </ul>
              )
            });
          }
        } catch (error) {
          console.error(`Error handling new notification: ${error?.message || ""}`);
        }
      };

      const handleSyncNotificationRead = ({ notificationId, timestamp }) => {
        if (Realtime_Notifications_UI?.treatment !== "on") {
          return;
        }

        try {
          const notificationRef = client.cache.identify({
            __typename: "notifications",
            id: notificationId
          });
          client.cache.writeFragment({
            id: notificationRef,
            fragment: UPDATE_NOTIFICATIONS_READ_FRAGMENT,
            data: { read: timestamp }
          });

          const unreadCountData = client.cache.readQuery({
            query: GET_UNREAD_COUNT,
            variables: { associationid: userAssociationId }
          });
          if (unreadCountData?.notifications_aggregate?.aggregate?.count > 0) {
            const newCount = Math.max(unreadCountData.notifications_aggregate.aggregate.count - 1, 0);
            client.cache.writeQuery({
              query: GET_UNREAD_COUNT,
              variables: { associationid: userAssociationId },
              data: {
                notifications_aggregate: {
                  __typename: "notifications_aggregate",
                  aggregate: {
                    __typename: "notifications_aggregate_fields",
                    count: newCount
                  }
                }
              }
            });
          }
        } catch (error) {
          console.error("Error in handleSyncNotificationRead:", error);
        }
      };

      const handleSyncAllNotificationsRead = ({ timestamp }) => {
        if (Realtime_Notifications_UI?.treatment !== "on") {
          return;
        }

        try {
          const queryVars = {
            limit: INITIAL_NOTIFICATIONS,
            offset: 0,
            where: { associationid: { _eq: userAssociationId } }
          };
          const cachedData = client.cache.readQuery({
            query: GET_NOTIFICATIONS,
            variables: queryVars
          });

          if (cachedData?.notifications) {
            cachedData.notifications.forEach((notif) => {
              if (!notif.read) {
                const notifRef = client.cache.identify({ __typename: "notifications", id: notif.id });
                client.cache.writeFragment({
                  id: notifRef,
                  fragment: UPDATE_NOTIFICATIONS_READ_FRAGMENT,
                  data: { read: timestamp }
                });
              }
            });
          }

          client.cache.writeQuery({
            query: GET_UNREAD_COUNT,
            variables: { associationid: userAssociationId },
            data: {
              notifications_aggregate: {
                __typename: "notifications_aggregate",
                aggregate: {
                  __typename: "notifications_aggregate_fields",
                  count: 0
                }
              }
            }
          });
        } catch (error) {
          console.error(`Error In HandleSyncAllNotificationsRead: ${error?.message || ""}`);
        }
      };

      socketInstance.on("connect", handleConnect);
      socketInstance.on("reconnect", handleReconnect);
      socketInstance.on("connect_error", handleConnectionError);
      socketInstance.on("disconnect", handleDisconnect);
      socketInstance.on("bodyshop-message", handleBodyshopMessage);
      socketInstance.on("notification", handleNotification);
      socketInstance.on("sync-notification-read", handleSyncNotificationRead);
      socketInstance.on("sync-all-notifications-read", handleSyncAllNotificationsRead);
      socketInstance.on(FEATURE_FLAGS_CHANGED_EVENT, handleFeatureFlagsChanged);
      socketInstance.on("token-updated", handleTokenUpdated);

      if (tokenSyncIntervalRef.current) {
        clearInterval(tokenSyncIntervalRef.current);
      }
      tokenSyncIntervalRef.current = setInterval(() => {
        if (!socketInstance.connected) return;
        syncCurrentTokenToSocket();
      }, TOKEN_SYNC_INTERVAL_MS);
    };

    const unsubscribe = auth.onIdTokenChanged(async (user) => {
      if (!user) {
        if (tokenSyncIntervalRef.current) {
          clearInterval(tokenSyncIntervalRef.current);
          tokenSyncIntervalRef.current = null;
        }
        socketRef.current?.disconnect();
        socketRef.current = null;
        setIsConnected(false);
        return;
      }

      const token = await user.getIdToken();
      if (socketRef.current) {
        socketRef.current.auth = { token, bodyshopId: bodyshop.id };
        if (socketRef.current.connected) {
          socketRef.current.emit("update-token", { token, bodyshopId: bodyshop.id });
        }
      } else {
        initializeSocket(token).catch((err) =>
          console.error("Something went wrong Initializing Sockets:", err?.message || "")
        );
      }
    });

    return () => {
      unsubscribe();
      if (tokenSyncIntervalRef.current) {
        clearInterval(tokenSyncIntervalRef.current);
        tokenSyncIntervalRef.current = null;
      }
      if (socketRef.current) {
        socketRef.current.disconnect();
        socketRef.current = null;
        setIsConnected(false);
      }
    };
  }, [bodyshop?.id, userAssociationId, Realtime_Notifications_UI?.treatment]);

  return (
    <SocketContext.Provider
      value={{
        socket: socketRef.current,
        clientId,
        isConnected,
        reconnectSocket,
        markNotificationRead,
        markAllNotificationsRead,
        scenarioNotificationsOn: Realtime_Notifications_UI?.treatment === "on"
      }}
    >
      {children}
    </SocketContext.Provider>
  );
};

export default SocketProvider;
