import { 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";
import { useTranslation } from "react-i18next";
import { useSplitTreatments } from "@splitsoftware/splitio-react";
import { INITIAL_NOTIFICATIONS, SocketContext } from "./useSocket.js";
import {
  QUERY_MY_TASKS_COUNT,
  QUERY_TASKS_NO_DUE_DATE_PAGINATED,
  QUERY_TASKS_WITH_DUE_DATES
} from "../../graphql/tasks.queries";

const LIMIT = INITIAL_NOTIFICATIONS;

/**
 * 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 [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 }
  } = useSplitTreatments({
    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)
  });

  useEffect(() => {
    const initializeSocket = async (token) => {
      if (!bodyshop || !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: 10000
      });

      socketRef.current = socketInstance;

      const handleBodyshopMessage = (message) => {
        if (!message || !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 countVars = {
              assigned_to: assignedToId,
              bodyshopid: bodyshop?.id
            };

            // Helper to update a list in cache locally
            const updateListCache = (query, vars, action) => {
              // action: 'add' or 'remove'
              try {
                const current = client.readQuery({ query, variables: vars });
                if (!current) return; // List not loaded yet, skip

                let updatedTasks = [...current.tasks];
                let delta = 0;

                if (action === "remove") {
                  const prevLength = updatedTasks.length;
                  updatedTasks = updatedTasks.filter((task) => task.id !== payload.id);
                  if (updatedTasks.length < prevLength) delta -= 1;
                }

                if (action === "add") {
                  const exists = updatedTasks.some((task) => task.id === payload.id);
                  if (!exists) {
                    const newTask = { ...payload, __typename: "tasks" };
                    updatedTasks.push(newTask);
                    delta += 1;
                  }
                }

                // Sort the array based on the query type if delta !== 0
                if (delta !== 0) {
                  if (query === QUERY_TASKS_WITH_DUE_DATES) {
                    updatedTasks.sort((a, b) => {
                      const da = new Date(a.due_date);
                      const db = new Date(b.due_date);
                      if (da < db) return -1;
                      if (da > db) return 1;
                      const ca = new Date(a.created_at);
                      const cb = new Date(b.created_at);
                      return cb - ca; // desc
                    });
                  } else if (query === QUERY_TASKS_NO_DUE_DATE_PAGINATED) {
                    updatedTasks.sort((a, b) => new Date(b.created_at) - new Date(a.created_at));
                  }
                }

                let data = { ...current, tasks: updatedTasks };

                // Adjust aggregate count for no due date query
                if (query === QUERY_TASKS_NO_DUE_DATE_PAGINATED && delta !== 0 && current.tasks_aggregate) {
                  const newCount = Math.max(0, current.tasks_aggregate.aggregate.count + delta);
                  data = {
                    ...data,
                    tasks_aggregate: {
                      ...current.tasks_aggregate,
                      aggregate: {
                        ...current.tasks_aggregate.aggregate,
                        count: newCount
                      }
                    }
                  };
                }

                client.writeQuery({ query, variables: vars, data });
              } catch (error) {
                console.error("Error updating task list cache:", error);
              }
            };

            // Update lists based on event type
            if (message.type === "task-deleted") {
              updateListCache(QUERY_TASKS_WITH_DUE_DATES, dueVars, "remove");
              updateListCache(QUERY_TASKS_NO_DUE_DATE_PAGINATED, noDueVars, "remove");
            } else {
              const targetQuery = payload.due_date ? QUERY_TASKS_WITH_DUE_DATES : QUERY_TASKS_NO_DUE_DATE_PAGINATED;
              const targetVars = payload.due_date ? dueVars : noDueVars;

              if (message.type === "task-updated") {
                // Remove from both in case of due_date change
                updateListCache(QUERY_TASKS_WITH_DUE_DATES, dueVars, "remove");
                updateListCache(QUERY_TASKS_NO_DUE_DATE_PAGINATED, noDueVars, "remove");
              }

              // Add to the target list
              updateListCache(targetQuery, targetVars, "add");
            }

            // Locally update the total count (no network call)
            try {
              const currentCountQuery = client.readQuery({ query: QUERY_MY_TASKS_COUNT, variables: countVars });
              if (currentCountQuery) {
                let countDelta = 0;
                if (message.type === "task-created") countDelta = 1;
                else if (message.type === "task-deleted") countDelta = -1;
                // For "task-updated", assume no change to count unless specific logic (e.g., completion)

                if (countDelta !== 0) {
                  const newCount = Math.max(0, currentCountQuery.tasks_aggregate.aggregate.count + countDelta);
                  client.writeQuery({
                    query: QUERY_MY_TASKS_COUNT,
                    variables: countVars,
                    data: {
                      tasks_aggregate: {
                        __typename: "tasks_aggregate",
                        aggregate: {
                          __typename: "tasks_aggregate_fields",
                          count: newCount
                        }
                      }
                    }
                  });
                }
              }
            } catch (error) {
              console.error("Error updating task count cache:", error);
            }

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

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

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

        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({
              message: (
                <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);
    };

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

    return () => {
      unsubscribe();
      if (socketRef.current) {
        socketRef.current.disconnect();
        socketRef.current = null;
        setIsConnected(false);
      }
    };
  }, [
    bodyshop,
    notification,
    userAssociationId,
    markNotificationRead,
    markAllNotificationsRead,
    navigate,
    currentUser,
    Realtime_Notifications_UI,
    t
  ]);

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

export default SocketProvider;