import { useMemo } from 'react';
import { useQueryClient, type QueryClient } from '@tanstack/react-query';
import type { StoreApi, UseBoundStore } from 'zustand';
import type { ResearchComposerState } from '../hooks/useResearchComposer';
import {
  replaceProjectionWorkspace,
  upsertProjectionGhost,
  upsertProjectionNote,
} from '../lib/researchProjection';
import type {
  GhostReviewAction,
  GhostRecommendedAction,
  GhostResolutionStatus,
  NoteType,
  ResearchFilters,
  ResearchStage,
  ResearchWorkspace,
  ResearchFile,
  ThesisStatus,
  UploadResearchFileInput,
  WorkspaceProjection,
  WorkspaceViewKind,
} from '../types/research';
import { researchApi } from './researchApi';
import {
  researchQueryKeys,
  researchWorkspacesQueryOptions,
} from './researchQueries';
import {
  type ResearchUiStoreState,
  useResearchUiStore,
} from './researchUiStore';

const DEFAULT_WORKSPACE_VIEW: WorkspaceViewKind = 'canvas';

const normalizeTicker = (ticker?: string | null) => ticker?.trim().toUpperCase() ?? '';

const upsertWorkspaceRecord = (
  current: ResearchWorkspace[],
  nextWorkspace: ResearchWorkspace,
) => {
  const existingIndex = current.findIndex((workspace) => workspace.id === nextWorkspace.id);
  if (existingIndex === -1) {
    return [nextWorkspace, ...current];
  }

  return current.map((workspace) =>
    workspace.id === nextWorkspace.id ? nextWorkspace : workspace,
  );
};

const mergeJobs = <
  T extends {
    id: string;
  },
>(
  current: T[],
  incoming: T[],
) => {
  const nextById = new Map(current.map((job) => [job.id, job]));
  incoming.forEach((job) => nextById.set(job.id, job));
  return [...nextById.values()].sort((left: any, right: any) =>
    String(right.updatedAt).localeCompare(String(left.updatedAt)),
  );
};

const getCachedWorkspaceDefaultView = (
  queryClient: QueryClient,
  workspaceId: string,
) => {
  const snapshot = queryClient.getQueryData<WorkspaceProjection>(
    researchQueryKeys.snapshot(workspaceId),
  );
  if (snapshot) {
    return snapshot.workspace.defaultView;
  }

  const workspaces =
    queryClient.getQueryData<ResearchWorkspace[]>(researchQueryKeys.workspaces) ?? [];
  return (
    workspaces.find((workspace) => workspace.id === workspaceId)?.defaultView ??
    DEFAULT_WORKSPACE_VIEW
  );
};

const patchSnapshot = (
  queryClient: QueryClient,
  workspaceId: string,
  updater: (snapshot: WorkspaceProjection) => WorkspaceProjection,
) => {
  queryClient.setQueryData<WorkspaceProjection | undefined>(
    researchQueryKeys.snapshot(workspaceId),
    (current) => (current ? updater(current) : current),
  );
};

const upsertWorkspaceCache = (queryClient: QueryClient, workspace: ResearchWorkspace) => {
  queryClient.setQueryData<ResearchWorkspace[]>(
    researchQueryKeys.workspaces,
    (current = []) => upsertWorkspaceRecord(current, workspace),
  );
  patchSnapshot(queryClient, workspace.id, (snapshot) =>
    replaceProjectionWorkspace(snapshot, workspace),
  );
};

const clearWorkspaceCache = (queryClient: QueryClient, workspaceId: string) => {
  queryClient.removeQueries({
    queryKey: researchQueryKeys.snapshot(workspaceId),
    exact: true,
  });
  queryClient.removeQueries({
    queryKey: researchQueryKeys.jobs(workspaceId),
    exact: true,
  });
  queryClient.removeQueries({
    queryKey: researchQueryKeys.searchIdsPrefix(workspaceId),
  });
};

const resolveWorkspaceList = async (queryClient: QueryClient) =>
  queryClient.ensureQueryData(researchWorkspacesQueryOptions());

const runMutation = async <T,>(
  store: UseBoundStore<StoreApi<ResearchUiStoreState>>,
  key: string,
  task: () => Promise<T>,
): Promise<T> => {
  store.getState().setMutationPending(key);
  try {
    const result = await task();
    store.getState().setMutationSucceeded(key);
    return result;
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    store.getState().setMutationFailed(key, message);
    throw error;
  }
};

export const researchActionKeys = {
  createWorkspace: 'workspace:create',
  updateWorkspace: (workspaceId: string) => `workspace:update:${workspaceId}`,
  archiveWorkspace: (workspaceId: string) => `workspace:archive:${workspaceId}`,
  captureNote: 'note:capture',
  updateNote: (noteId: string) => `note:update:${noteId}`,
  archiveNote: (noteId: string) => `note:archive:${noteId}`,
  promoteNote: (noteId: string) => `note:promote:${noteId}`,
  reviewGhost: (ghostId: string, action: GhostReviewAction) =>
    `ghost:${action}:${ghostId}`,
  resolveGhost: (ghostId: string, status: GhostResolutionStatus) =>
    `ghost:resolve:${status}:${ghostId}`,
  retryJobs: (workspaceId: string) => `jobs:retry:${workspaceId}`,
  uploadFiles: (workspaceId: string) => `files:upload:${workspaceId}`,
};

export const createResearchActions = (
  queryClient: QueryClient,
  store: UseBoundStore<StoreApi<ResearchUiStoreState>> = useResearchUiStore,
) => {
  const selectWorkspace = (workspaceId: string) => {
    const defaultView = getCachedWorkspaceDefaultView(queryClient, workspaceId);
    const state = store.getState();
    state.initializeWorkspaceState(workspaceId, defaultView);
    state.setActiveWorkspaceId(workspaceId);
    state.clearSelection(workspaceId);
    state.setMobilePane('none');
  };

  const ensureWorkspace = async (request: {
    ticker?: string;
    workspaceId?: string | null;
  }) => {
    const workspaces = await resolveWorkspaceList(queryClient);

    if (request.workspaceId) {
      const existing =
        workspaces.find((workspace) => workspace.id === request.workspaceId) ?? null;
      if (existing) {
        selectWorkspace(existing.id);
      }
      return existing;
    }

    const normalizedTicker = normalizeTicker(request.ticker);
    if (normalizedTicker) {
      const existing =
        workspaces.find(
          (workspace) => normalizeTicker(workspace.primaryTicker) === normalizedTicker,
        ) ?? null;
      if (existing) {
        selectWorkspace(existing.id);
        return existing;
      }

      return runMutation(store, researchActionKeys.createWorkspace, async () => {
        const created = await researchApi.createWorkspace({
          name: `${normalizedTicker} Research`,
          primaryTicker: normalizedTicker,
        });
        upsertWorkspaceCache(queryClient, created);
        store.getState().initializeWorkspaceState(created.id, created.defaultView);
        selectWorkspace(created.id);
        await queryClient.invalidateQueries({
          queryKey: researchQueryKeys.workspaces,
        });
        return created;
      });
    }

    const firstWorkspace = workspaces[0] ?? null;
    if (firstWorkspace) {
      selectWorkspace(firstWorkspace.id);
    }
    return firstWorkspace;
  };

  return {
    ensureWorkspace,
    createWorkspace: async (ticker: string) =>
      runMutation(store, researchActionKeys.createWorkspace, async () => {
        const normalizedTicker = normalizeTicker(ticker);
        const created = await researchApi.createWorkspace({
          name: `${normalizedTicker} Research`,
          primaryTicker: normalizedTicker,
        });
        upsertWorkspaceCache(queryClient, created);
        store.getState().initializeWorkspaceState(created.id, created.defaultView);
        selectWorkspace(created.id);
        store.getState().closeCreateWorkspaceModal();
        await queryClient.invalidateQueries({
          queryKey: researchQueryKeys.workspaces,
        });
        return created;
      }),
    selectWorkspace,
    updateWorkspace: async (request: {
      workspaceId: string;
      name?: string;
      stage?: ResearchStage;
      defaultView?: WorkspaceViewKind;
      pinnedNoteIds?: string[];
    }) =>
      runMutation(
        store,
        researchActionKeys.updateWorkspace(request.workspaceId),
        async () => {
          const updated = await researchApi.updateWorkspace(request);
          upsertWorkspaceCache(queryClient, updated);
          store
            .getState()
            .initializeWorkspaceState(updated.id, updated.defaultView);
          await queryClient.invalidateQueries({
            queryKey: researchQueryKeys.workspaces,
          });
          return updated;
        },
      ),
    archiveWorkspace: async (workspaceId: string) =>
      runMutation(
        store,
        researchActionKeys.archiveWorkspace(workspaceId),
        async () => {
          const archived = await researchApi.archiveWorkspace({
            workspaceId,
            archived: true,
          });
          queryClient.setQueryData<ResearchWorkspace[]>(
            researchQueryKeys.workspaces,
            (current = []) => current.filter((workspace) => workspace.id !== archived.id),
          );
          clearWorkspaceCache(queryClient, archived.id);
          const currentActiveWorkspaceId = store.getState().activeWorkspaceId;
          if (currentActiveWorkspaceId === archived.id) {
            const remainingWorkspaces = await resolveWorkspaceList(queryClient);
            const nextWorkspace =
              remainingWorkspaces.find((workspace) => workspace.id !== archived.id) ?? null;
            if (nextWorkspace) {
              selectWorkspace(nextWorkspace.id);
            } else {
              store.getState().setActiveWorkspaceId(null);
            }
          }
          await queryClient.invalidateQueries({
            queryKey: researchQueryKeys.workspaces,
          });
          return archived;
        },
      ),
    captureNote: async (args: {
      draft: ResearchComposerState;
      fallbackTicker?: string;
      explicitWorkspaceId?: string | null;
      autoCreateFromTicker?: boolean;
    }) =>
      runMutation(store, researchActionKeys.captureNote, async () => {
        const normalizedText = args.draft.rawText.trim();
        if (!normalizedText) {
          throw new Error('Note text is required.');
        }

        const resolvedTicker =
          normalizeTicker(args.draft.ticker) || normalizeTicker(args.fallbackTicker) || undefined;
        const resolvedWorkspace = await ensureWorkspace({
          workspaceId: args.explicitWorkspaceId || args.draft.workspaceId || undefined,
          ticker: args.autoCreateFromTicker ? resolvedTicker : undefined,
        });
        if (!resolvedWorkspace) {
          throw new Error(
            'Select or create a research workspace before saving this note.',
          );
        }

        const result = await researchApi.captureNote({
          workspaceId: resolvedWorkspace.id,
          rawText: normalizedText,
          bodyHtml: args.draft.bodyHtml || undefined,
          embeds: args.draft.embeds,
          fileIds: args.draft.fileIds,
          ticker: resolvedTicker,
          userNoteTypeOverride: args.draft.userNoteTypeOverride,
          sourceRef:
            args.draft.sourceUrl || args.draft.sourceTitle
              ? {
                  kind: args.draft.sourceKind,
                  title: args.draft.sourceTitle || undefined,
                  url: args.draft.sourceUrl || undefined,
                }
              : undefined,
        });

        patchSnapshot(queryClient, resolvedWorkspace.id, (snapshot) =>
          upsertProjectionNote(snapshot, result.note),
        );
        queryClient.setQueryData(
          researchQueryKeys.jobs(resolvedWorkspace.id),
          (current: any[] = []) => mergeJobs(current, result.queuedJobs),
        );
        store.getState().initializeWorkspaceState(
          resolvedWorkspace.id,
          resolvedWorkspace.defaultView,
        );
        selectWorkspace(resolvedWorkspace.id);
        store.getState().selectNote(resolvedWorkspace.id, result.note.id);
        store.getState().closeCaptureModal();
        await queryClient.invalidateQueries({
          queryKey: researchQueryKeys.searchIdsPrefix(resolvedWorkspace.id),
        });
        await queryClient.invalidateQueries({
          queryKey: researchQueryKeys.jobs(resolvedWorkspace.id),
        });
        return result;
      }),
    updateNote: async (request: {
      noteId: string;
      workspaceId: string;
      rawText?: string;
      bodyHtml?: string;
      embeds?: import('../types/research').NoteEmbedInput[];
      fileIds?: string[];
      title?: string;
      noteType?: NoteType;
      pinned?: boolean;
      thesisStatus?: ThesisStatus;
    }) =>
      runMutation(store, researchActionKeys.updateNote(request.noteId), async () => {
        const updated = await researchApi.updateNote({
          noteId: request.noteId,
          rawText: request.rawText,
          bodyHtml: request.bodyHtml,
          embeds: request.embeds,
          fileIds: request.fileIds,
          title: request.title,
          noteType: request.noteType,
          pinned: request.pinned,
          thesisStatus: request.thesisStatus,
        });
        patchSnapshot(queryClient, updated.workspaceId, (snapshot) =>
          upsertProjectionNote(snapshot, updated),
        );
        await Promise.all([
          queryClient.invalidateQueries({
            queryKey: researchQueryKeys.searchIdsPrefix(updated.workspaceId),
          }),
          queryClient.invalidateQueries({
            queryKey: researchQueryKeys.auditTrail(updated.id),
          }),
          queryClient.invalidateQueries({
            queryKey: researchQueryKeys.jobs(updated.workspaceId),
          }),
        ]);
        return updated;
      }),
    archiveNote: async (noteId: string, workspaceId: string, archived: boolean) =>
      runMutation(store, researchActionKeys.archiveNote(noteId), async () => {
        const updated = await researchApi.archiveNote({ noteId, archived });
        patchSnapshot(queryClient, workspaceId, (snapshot) =>
          upsertProjectionNote(snapshot, updated),
        );
        await Promise.all([
          queryClient.invalidateQueries({
            queryKey: researchQueryKeys.searchIdsPrefix(workspaceId),
          }),
          queryClient.invalidateQueries({
            queryKey: researchQueryKeys.auditTrail(noteId),
          }),
          queryClient.invalidateQueries({
            queryKey: researchQueryKeys.jobs(workspaceId),
          }),
        ]);
        return updated;
      }),
    promoteNote: async (
      noteId: string,
      workspaceId: string,
      thesisStatus?: ThesisStatus,
      noteType?: NoteType,
    ) =>
      runMutation(store, researchActionKeys.promoteNote(noteId), async () => {
        const promoted = await researchApi.promoteNote({
          noteId,
          thesisStatus,
          noteType,
        });
        patchSnapshot(queryClient, workspaceId, (snapshot) =>
          upsertProjectionNote(snapshot, promoted),
        );
        await Promise.all([
          queryClient.invalidateQueries({
            queryKey: researchQueryKeys.searchIdsPrefix(workspaceId),
          }),
          queryClient.invalidateQueries({
            queryKey: researchQueryKeys.auditTrail(noteId),
          }),
          queryClient.invalidateQueries({
            queryKey: researchQueryKeys.jobs(workspaceId),
          }),
        ]);
        return promoted;
      }),
    reviewGhost: async (
      ghostId: string,
      workspaceId: string,
      action: GhostReviewAction,
    ) =>
      runMutation(
        store,
        researchActionKeys.reviewGhost(ghostId, action),
        async () => {
          const reviewed = await researchApi.reviewGhost({
            ghostNoteId: ghostId,
            action,
          });
          patchSnapshot(queryClient, workspaceId, (snapshot) =>
            upsertProjectionGhost(snapshot, reviewed),
          );
          await Promise.all([
            ...reviewed.supportingNoteIds.map((noteId) =>
              queryClient.invalidateQueries({
                queryKey: researchQueryKeys.auditTrail(noteId),
              }),
            ),
            ...reviewed.contradictingNoteIds.map((noteId) =>
              queryClient.invalidateQueries({
                queryKey: researchQueryKeys.auditTrail(noteId),
              }),
            ),
          ]);
          return reviewed;
        },
      ),
    resolveGhostIntervention: async (request: {
      ghostId: string;
      workspaceId: string;
      resolutionStatus: GhostResolutionStatus;
      actionTaken?: GhostRecommendedAction;
      notes?: string;
      linkedNoteIds?: string[];
    }) =>
      runMutation(
        store,
        researchActionKeys.resolveGhost(request.ghostId, request.resolutionStatus),
        async () => {
          const resolved = await researchApi.resolveGhostIntervention({
            ghostNoteId: request.ghostId,
            workspaceId: request.workspaceId,
            resolutionStatus: request.resolutionStatus,
            actionTaken: request.actionTaken,
            notes: request.notes,
            linkedNoteIds: request.linkedNoteIds,
          });
          patchSnapshot(queryClient, request.workspaceId, (snapshot) =>
            upsertProjectionGhost(snapshot, resolved),
          );
          await Promise.all([
            ...resolved.supportingNoteIds.map((noteId) =>
              queryClient.invalidateQueries({
                queryKey: researchQueryKeys.auditTrail(noteId),
              }),
            ),
            ...resolved.contradictingNoteIds.map((noteId) =>
              queryClient.invalidateQueries({
                queryKey: researchQueryKeys.auditTrail(noteId),
              }),
            ),
          ]);
          return resolved;
        },
      ),
    retryJobs: async (workspaceId: string) =>
      runMutation(store, researchActionKeys.retryJobs(workspaceId), async () => {
        const jobs = await researchApi.retryJobs({ workspaceId });
        queryClient.setQueryData(
          researchQueryKeys.jobs(workspaceId),
          (current: any[] = []) => mergeJobs(current, jobs),
        );
        await queryClient.invalidateQueries({
          queryKey: researchQueryKeys.jobs(workspaceId),
        });
        return jobs;
      }),
    uploadFiles: async (
      workspaceId: string,
      files: UploadResearchFileInput[],
    ): Promise<ResearchFile[]> =>
      runMutation(store, researchActionKeys.uploadFiles(workspaceId), async () => {
        const uploaded = await researchApi.uploadFiles({ workspaceId, files });
        queryClient.setQueryData<ResearchFile[]>(
          researchQueryKeys.files(workspaceId),
          (current = []) => {
            const byId = new Map(current.map((file) => [file.id, file]));
            uploaded.forEach((file) => byId.set(file.id, file));
            return [...byId.values()].sort((left, right) =>
              String(right.createdAt).localeCompare(String(left.createdAt)),
            );
          },
        );
        await queryClient.invalidateQueries({
          queryKey: researchQueryKeys.files(workspaceId),
        });
        return uploaded;
      }),
    setView: (workspaceId: string, view: WorkspaceViewKind) => {
      store.getState().initializeWorkspaceState(workspaceId, view);
      store.getState().setWorkspaceView(workspaceId, view);
    },
    setFilters: (workspaceId: string, patch: Partial<ResearchFilters>) =>
      store.getState().setWorkspaceFilters(workspaceId, patch),
    resetFilters: (workspaceId: string) =>
      store.getState().resetWorkspaceFilters(workspaceId),
    selectNote: (workspaceId: string, noteId: string) =>
      store.getState().selectNote(workspaceId, noteId),
    selectGhost: (workspaceId: string, ghostId: string) =>
      store.getState().selectGhost(workspaceId, ghostId),
    clearSelection: (workspaceId: string) =>
      store.getState().clearSelection(workspaceId),
    setMemoDraft: (workspaceId: string, value: string) =>
      store.getState().setMemoDraft(workspaceId, value),
  };
};

export const useResearchActions = () => {
  const queryClient = useQueryClient();
  return useMemo(() => createResearchActions(queryClient), [queryClient]);
};
