app/src/admin/hooks/usePlanWork.ts

import { useState, useCallback, useMemo } from "react";
import {
  useQuery,
  useQueryClient,
  useMutation,
  keepPreviousData,
} from "@tanstack/react-query";
import apiFetch from "../utils/api";
import { planKeys, GridData, ResolvedCell } from "../lib/queries/plan";

export type ViewMode = "week" | "month";
export type ModalMode =
  | { kind: "closed" }
  | { kind: "create"; userId: number; date: string }
  | { kind: "create-override"; userId: number; date: string }
  | { kind: "edit-range"; entryId: number; userId: number; date: string }
  | { kind: "edit-override"; overrideId: number; userId: number; date: string }
  | { kind: "day-in-range"; userId: number; date: string; entryId: number }
  | { kind: "day"; userId: number; date: string }
  | { kind: "view"; userId: number; date: string };

function isoDate(d: Date): string {
  return d.toISOString().slice(0, 10);
}

function startOfWeek(d: Date): Date {
  const day = d.getUTCDay(); // 0 = Sun
  const diff = day === 0 ? -6 : 1 - day; // Monday-start
  const r = new Date(d);
  r.setUTCDate(d.getUTCDate() + diff);
  return r;
}

function addDays(d: Date, n: number): Date {
  const r = new Date(d);
  r.setUTCDate(d.getUTCDate() + n);
  return r;
}

async function apiCall(
  url: string,
  options: RequestInit = {},
): Promise<unknown> {
  const res = await apiFetch(url, options);
  if (!res.ok) {
    let message = "Chyba serveru";
    try {
      const body = await res.json();
      if (body && typeof body.error === "string") message = body.error;
    } catch {
      // body wasn't JSON; use default
    }
    throw new Error(message);
  }
  try {
    return await res.json();
  } catch {
    return null;
  }
}

// Request-body shapes for the plan mutations. The server validates the full
// schema; these capture the fields the optimistic patches read.
export interface PlanEntryBody {
  user_id: number;
  date_from: string;
  date_to: string;
  project_id?: number | null;
  category: string;
  note: string;
}

export interface PlanEntryPatchBody {
  date_from?: string;
  date_to?: string;
  project_id?: number | null;
  category?: string;
  note?: string;
}

export interface PlanOverrideBody {
  user_id: number;
  shift_date: string;
  project_id?: number | null;
  category: string;
  note: string;
}

export interface PlanOverridePatchBody {
  project_id?: number | null;
  category?: string;
  note?: string;
}

export interface BulkCreateBody {
  user_ids: number[];
  date_from: string;
  date_to: string;
  include_weekends: boolean;
  project_id: number | null;
  category: string;
  note: string;
}

export interface UsePlanWorkArgs {
  initialDate?: Date;
  canEdit: boolean;
}

export function usePlanWork({ initialDate, canEdit }: UsePlanWorkArgs) {
  const qc = useQueryClient();
  const [view, setView] = useState<ViewMode>("week");
  const [anchor, setAnchor] = useState<Date>(initialDate ?? new Date());
  const [filterActive, setFilterActive] = useState(true);
  const [modal, setModal] = useState<ModalMode>({ kind: "closed" });

  const range = useMemo(() => {
    if (view === "week") {
      const from = startOfWeek(anchor);
      const to = addDays(from, 6);
      return { from, to };
    }
    const from = new Date(
      Date.UTC(anchor.getUTCFullYear(), anchor.getUTCMonth(), 1),
    );
    const to = new Date(
      Date.UTC(anchor.getUTCFullYear(), anchor.getUTCMonth() + 1, 0),
    );
    return { from, to };
  }, [view, anchor]);

  const gridQuery = useQuery({
    queryKey: planKeys.grid(isoDate(range.from), isoDate(range.to), view),
    queryFn: async (): Promise<GridData> => {
      const res = await apiFetch(
        `/api/admin/plan/grid?date_from=${isoDate(range.from)}&date_to=${isoDate(range.to)}&view=${view}`,
      );
      if (!res.ok) {
        let message = "Chyba serveru";
        try {
          const errBody = await res.json();
          if (errBody && typeof errBody.error === "string")
            message = errBody.error;
        } catch {
          // body wasn't JSON; use default
        }
        throw new Error(message);
      }
      const body = (await res.json()) as { data: GridData };
      return body.data;
    },
    // keepPreviousData holds the old range's data visible while the new
    // range's fetch is in flight. Without this, `gridQuery.data` is
    // immediately undefined on key change and the grid flashes empty.
    // `isPlaceholderData` is true while the displayed data is stale.
    placeholderData: keepPreviousData,
  });

  const invalidate = useCallback(() => {
    qc.invalidateQueries({ queryKey: planKeys.all });
    // Plan mutations write audit-log rows, so they touch the audit domain
    // too. Invalidate the dashboard activity feed and the audit-log page so
    // they reflect the change without a manual F5 (prefix-matches all
    // ["audit-log", {...}] filter variants). The dashboard query is active
    // when visible and refetches immediately despite its staleTime.
    qc.invalidateQueries({ queryKey: ["dashboard"] });
    qc.invalidateQueries({ queryKey: ["audit-log"] });
  }, [qc]);

  // --- Optimistic cache patches ---------------------------------------------
  // After a mutation succeeds, we patch the visible grid query's data
  // *before* the refetch resolves. This is what stops the full-grid
  // re-animation: the visible data updates in place (and the new cell
  // gets a one-shot pulse via PlanWork's lastMutated state), so the grid
  // never unmounts. The subsequent invalidate+refetch confirms the patch
  // from the server (which is the source of truth) and is a no-op if the
  // server's data matches our patch.
  //
  // If the mutation fails, the mutation's onError handler in PlanWork
  // rolls the cache back to the snapshot we took before the patch.

  const currentGridKey = planKeys.grid(
    isoDate(range.from),
    isoDate(range.to),
    view,
  );

  // Walk the days in [from, to] inclusive (YYYY-MM-DD strings). Used to
  // figure out which cell keys to patch for a range entry.
  function eachDay(from: string, to: string): string[] {
    const out: string[] = [];
    const a = new Date(from + "T00:00:00.000Z");
    const b = new Date(to + "T00:00:00.000Z");
    for (let d = new Date(a); d <= b; d.setUTCDate(d.getUTCDate() + 1)) {
      out.push(d.toISOString().slice(0, 10));
    }
    return out;
  }

  // Snapshot the currently-cached grid for a given key. Returns
  // `undefined` if there's nothing cached (we don't patch empty data).
  function snapshotGrid(key: readonly unknown[]): GridData | undefined {
    return qc.getQueryData<GridData>(key);
  }

  // Apply a per-date cell patch to the cached grid at `key`. Returns the
  // previous cells[userId] map so the caller can roll back on error.
  function patchCells(
    key: readonly unknown[],
    userId: number,
    dates: string[],
    mutator: (prev: ResolvedCell[]) => ResolvedCell[],
  ): Record<string, ResolvedCell[]> | null {
    const prev = snapshotGrid(key);
    if (!prev) return null;
    const userPrev = prev.cells[userId] ?? {};
    const rolled: Record<string, ResolvedCell[]> = {};
    const userNext: Record<string, ResolvedCell[]> = { ...userPrev };
    for (const date of dates) {
      rolled[date] = userPrev[date] ?? [];
      userNext[date] = mutator(rolled[date]).slice(0, 3);
    }
    qc.setQueryData(key, {
      ...prev,
      cells: { ...prev.cells, [userId]: userNext },
    });
    return rolled;
  }

  // Build a ResolvedCell for a freshly-created entry. The server returns
  // a row with the same fields, so we mirror that shape. We don't have
  // entryId/overrideId from the request body — but the createEntry
  // response *does* include id; we could plumb it through. For now we
  // leave entryId as a placeholder that the refetch will replace.
  // (PlanGrid only uses source/entryId/overrideId to decide edit flow;
  // a fresh cell from an entry is "entry" source — the refetch fills
  // the actual id within ~200ms.)
  function makeEntryCell(args: {
    userId: number;
    date: string;
    projectId: number | null;
    category: string;
    note: string;
    rangeFrom: string;
    rangeTo: string;
    entryId: number | null;
  }): ResolvedCell {
    return {
      source: "entry",
      entryId: args.entryId,
      overrideId: null,
      user_id: args.userId,
      shift_date: args.date,
      project_id: args.projectId,
      // Optimistic placeholder — the invalidate+refetch fills the real
      // project label (resolveGrid embeds it) within ~200ms.
      project_number: null,
      project_name: null,
      category: args.category,
      note: args.note,
      rangeFrom: args.rangeFrom,
      rangeTo: args.rangeTo,
    };
  }

  function makeOverrideCell(args: {
    userId: number;
    date: string;
    projectId: number | null;
    category: string;
    note: string;
    overrideId: number | null;
  }): ResolvedCell {
    return {
      source: "override",
      entryId: null,
      overrideId: args.overrideId,
      user_id: args.userId,
      shift_date: args.date,
      project_id: args.projectId,
      // Optimistic placeholder — refetch fills the real project label.
      project_number: null,
      project_name: null,
      category: args.category,
      note: args.note,
      rangeFrom: null,
      rangeTo: null,
    };
  }

  // --- Mutations ---
  //
  // Optimistic-update pattern (idiomatic TanStack Query):
  //   onMutate  → apply the optimistic grid patch and RETURN the rollback
  //               snapshot as this invocation's context (concurrency-safe:
  //               each in-flight mutation owns its own context object, so two
  //               concurrent mutations can't clobber each other's snapshot —
  //               the previous `(mutation as any)._rolled = …` stash could).
  //   onError   → restore the cells captured in onMutate's context.
  //   onSettled → invalidate the domain so the server (source of truth)
  //               reconciles the patch on both success AND error.
  // The visible UX is unchanged: the new/updated cell still appears
  // immediately and the broad invalidate still refetches the grid.

  type CellRollback = {
    key: readonly unknown[];
    userId: number;
    rolled: Record<string, ResolvedCell[]> | null;
  } | null;

  // Restore the cells a patch overwrote, from an onMutate-returned context.
  // `ctx` may be `undefined` (e.g. if onMutate itself threw) — guarded below.
  function rollbackCells(ctx: CellRollback | undefined) {
    if (!ctx || !ctx.rolled) return;
    const prev = qc.getQueryData<GridData>(ctx.key);
    if (!prev) return;
    const userPrev = prev.cells[ctx.userId] ?? {};
    const userNext = { ...userPrev };
    for (const [date, cells] of Object.entries(ctx.rolled)) {
      userNext[date] = cells;
    }
    qc.setQueryData(ctx.key, {
      ...prev,
      cells: { ...prev.cells, [ctx.userId]: userNext },
    });
  }

  const createEntry = useMutation<unknown, Error, PlanEntryBody, CellRollback>({
    mutationFn: (body: PlanEntryBody) =>
      apiCall("/api/admin/plan/entries", {
        method: "POST",
        body: JSON.stringify(body),
        headers: { "Content-Type": "application/json" },
      }),
    onMutate: (body): CellRollback => {
      // Patch the visible grid with the new entry. We don't have the
      // server-assigned id yet (entryId stays null) — the refetch fills it
      // within ~200ms; PlanGrid only uses entryId for the edit flow.
      const days = eachDay(body.date_from, body.date_to);
      const rolled = patchCells(currentGridKey, body.user_id, days, (prev) => [
        makeEntryCell({
          userId: body.user_id,
          date: days[0],
          projectId: body.project_id ?? null,
          category: body.category,
          note: body.note,
          rangeFrom: body.date_from,
          rangeTo: body.date_to,
          entryId: null,
        }),
        ...prev,
      ]);
      return { key: currentGridKey, userId: body.user_id, rolled };
    },
    onError: (_err, _body, ctx) => rollbackCells(ctx),
    onSettled: () => invalidate(),
  });

  const updateEntry = useMutation<
    unknown,
    Error,
    { id: number; body: PlanEntryPatchBody; force?: boolean },
    CellRollback
  >({
    mutationFn: ({ id, body, force }) =>
      apiCall(`/api/admin/plan/entries/${id}${force ? "?force=1" : ""}`, {
        method: "PATCH",
        body: JSON.stringify(body),
        headers: { "Content-Type": "application/json" },
      }),
    onMutate: ({ id, body }): CellRollback => {
      // We don't know the new full range without the response, but we
      // do have the body's date_from/date_to. If those are present,
      // patch the new range. Old cells outside the new range are NOT
      // cleared here — they'll either still be valid (date_from/to
      // were partial updates) or the refetch will fix them.
      if (!body.date_from || !body.date_to) return null;
      const days = eachDay(body.date_from, body.date_to);
      // Find the user that owns this entry in the current grid by
      // looking for any cell with entryId === id (we already know
      // the id from vars; it doesn't change across the update).
      const grid = qc.getQueryData<GridData>(currentGridKey);
      let ownerUserId: number | null = null;
      if (grid) {
        for (const [uidStr, byDate] of Object.entries(grid.cells)) {
          for (const cells of Object.values(byDate)) {
            if (cells.some((c) => c.entryId === id)) {
              ownerUserId = Number(uidStr);
              break;
            }
          }
          if (ownerUserId !== null) break;
        }
      }
      if (ownerUserId === null) return null;
      const rolled = patchCells(currentGridKey, ownerUserId, days, (prev) => {
        const existing = prev.find((c) => c.entryId === id) ?? null;
        const updated = makeEntryCell({
          userId: ownerUserId!,
          date: days[0],
          projectId:
            body.project_id === undefined
              ? (existing?.project_id ?? null)
              : body.project_id,
          category: body.category ?? existing?.category ?? "work",
          note: body.note ?? existing?.note ?? "",
          rangeFrom: body.date_from!,
          rangeTo: body.date_to!,
          entryId: id,
        });
        return [updated, ...prev.filter((c) => c.entryId !== id)];
      });
      return { key: currentGridKey, userId: ownerUserId, rolled };
    },
    onError: (_err, _vars, ctx) => rollbackCells(ctx),
    onSettled: () => invalidate(),
  });

  const deleteEntry = useMutation<
    unknown,
    Error,
    { id: number; force?: boolean },
    CellRollback
  >({
    mutationFn: ({ id, force }) =>
      apiCall(`/api/admin/plan/entries/${id}${force ? "?force=1" : ""}`, {
        method: "DELETE",
      }),
    onMutate: (vars): CellRollback => {
      // For delete we need to know the entry's user_id and full range.
      // Look it up from the current grid: find the user that has a cell
      // with entryId === id, and read rangeFrom/rangeTo from that cell.
      const grid = qc.getQueryData<GridData>(currentGridKey);
      if (!grid) return null;
      for (const [uidStr, byDate] of Object.entries(grid.cells)) {
        let entryRange: { from: string; to: string } | null = null;
        for (const cells of Object.values(byDate)) {
          const hit = cells.find(
            (c) => c.entryId === vars.id && c.rangeFrom && c.rangeTo,
          );
          if (hit) {
            entryRange = { from: hit.rangeFrom!, to: hit.rangeTo! };
            break;
          }
        }
        if (entryRange) {
          const days = eachDay(entryRange.from, entryRange.to);
          const rolled = patchCells(
            currentGridKey,
            Number(uidStr),
            days,
            (prev) => prev.filter((c) => c.entryId !== vars.id),
          );
          return { key: currentGridKey, userId: Number(uidStr), rolled };
        }
      }
      return null;
    },
    onError: (_err, _vars, ctx) => rollbackCells(ctx),
    onSettled: () => invalidate(),
  });

  const createOverride = useMutation<
    unknown,
    Error,
    PlanOverrideBody,
    CellRollback
  >({
    mutationFn: (body: PlanOverrideBody) =>
      apiCall("/api/admin/plan/overrides", {
        method: "POST",
        body: JSON.stringify(body),
        headers: { "Content-Type": "application/json" },
      }),
    onMutate: (vars): CellRollback => {
      // overrideId stays null optimistically; the refetch fills the real id.
      const rolled = patchCells(
        currentGridKey,
        vars.user_id,
        [vars.shift_date],
        (prev) => [
          makeOverrideCell({
            userId: vars.user_id,
            date: vars.shift_date,
            projectId: vars.project_id ?? null,
            category: vars.category,
            note: vars.note,
            overrideId: null,
          }),
          ...prev,
        ],
      );
      return { key: currentGridKey, userId: vars.user_id, rolled };
    },
    onError: (_err, _vars, ctx) => rollbackCells(ctx),
    onSettled: () => invalidate(),
  });

  const updateOverride = useMutation<
    unknown,
    Error,
    { id: number; body: PlanOverridePatchBody; force?: boolean },
    CellRollback
  >({
    mutationFn: ({ id, body, force }) =>
      apiCall(`/api/admin/plan/overrides/${id}${force ? "?force=1" : ""}`, {
        method: "PATCH",
        body: JSON.stringify(body),
        headers: { "Content-Type": "application/json" },
      }),
    onMutate: (vars): CellRollback => {
      // Find the user/date for this overrideId in the current grid, then
      // patch that single cell with the new values.
      const grid = qc.getQueryData<GridData>(currentGridKey);
      if (!grid) return null;
      for (const [uidStr, byDate] of Object.entries(grid.cells)) {
        for (const [date, cells] of Object.entries(byDate)) {
          if (cells.some((c) => c.overrideId === vars.id)) {
            const rolled = patchCells(
              currentGridKey,
              Number(uidStr),
              [date],
              (prev) => {
                const existing =
                  prev.find((c) => c.overrideId === vars.id) ?? null;
                const updated = makeOverrideCell({
                  userId: Number(uidStr),
                  date,
                  projectId:
                    vars.body.project_id ?? existing?.project_id ?? null,
                  category: vars.body.category ?? existing?.category ?? "work",
                  note: vars.body.note ?? existing?.note ?? "",
                  overrideId: vars.id,
                });
                return [
                  updated,
                  ...prev.filter((c) => c.overrideId !== vars.id),
                ];
              },
            );
            return { key: currentGridKey, userId: Number(uidStr), rolled };
          }
        }
      }
      return null;
    },
    onError: (_err, _vars, ctx) => rollbackCells(ctx),
    onSettled: () => invalidate(),
  });

  const deleteOverride = useMutation<
    unknown,
    Error,
    { id: number; force?: boolean },
    CellRollback
  >({
    mutationFn: ({ id, force }) =>
      apiCall(`/api/admin/plan/overrides/${id}${force ? "?force=1" : ""}`, {
        method: "DELETE",
      }),
    onMutate: (vars): CellRollback => {
      const grid = qc.getQueryData<GridData>(currentGridKey);
      if (!grid) return null;
      for (const [uidStr, byDate] of Object.entries(grid.cells)) {
        for (const [date, cells] of Object.entries(byDate)) {
          if (cells.some((c) => c.overrideId === vars.id)) {
            const rolled = patchCells(
              currentGridKey,
              Number(uidStr),
              [date],
              (prev) => prev.filter((c) => c.overrideId !== vars.id),
            );
            return { key: currentGridKey, userId: Number(uidStr), rolled };
          }
        }
      }
      return null;
    },
    onError: (_err, _vars, ctx) => rollbackCells(ctx),
    onSettled: () => invalidate(),
  });

  const bulkCreate = useMutation<unknown, Error, BulkCreateBody>({
    mutationFn: (body: BulkCreateBody) =>
      apiCall("/api/admin/plan/entries/bulk", {
        method: "POST",
        body: JSON.stringify(body),
        headers: { "Content-Type": "application/json" },
      }),
    onSuccess: () => {
      // No optimistic patch — bulk can touch many cells; just invalidate and
      // let the grid refetch the authoritative state.
      invalidate();
    },
  });

  // Rollback is now handled idiomatically inside each mutation's `onError`
  // from its `onMutate`-returned context (concurrency-safe per-invocation
  // snapshot). This is kept as a no-op for the existing call sites in
  // PlanWork.tsx so they continue to compile unchanged; by the time a
  // `mutateAsync` promise rejects, the mutation's own `onError` has already
  // restored the patched cells, so there is nothing left to do here.

  function rollbackMutation(_mutation: unknown, _userId: number) {
    /* no-op — see onMutate/onError above */
  }

  return {
    view,
    setView,
    anchor,
    setAnchor,
    filterActive,
    setFilterActive,
    range,
    grid: gridQuery.data,
    gridLoading: gridQuery.isLoading,
    gridFetching: gridQuery.isFetching,
    gridIsPlaceholder: gridQuery.isPlaceholderData,
    modal,
    setModal,
    canEdit,
    createEntry,
    updateEntry,
    deleteEntry,
    createOverride,
    updateOverride,
    deleteOverride,
    bulkCreate,
    invalidate,
    // Exposed so PlanWork can roll back an optimistic patch when a
    // mutation throws. Pass the failed mutation and the userId it
    // targeted.
    rollbackMutation,
  };
}

export function getCells(
  grid: GridData | undefined,
  userId: number,
  date: string,
): ResolvedCell[] {
  return grid?.cells?.[userId]?.[date] ?? [];
}

export function getCell(
  grid: GridData | undefined,
  userId: number,
  date: string,
): ResolvedCell | null {
  return getCells(grid, userId, date)[0] ?? null;
}