import type { ColumnSizingState, Updater } from "@tanstack/react-table";
import type { StateCreator } from "zustand";
import { create } from "zustand";
import { devtools, persist } from "zustand/middleware";

import type { ProjectTab } from "@phoenix/pages/project/constants";
import { TRACE_ANNOTATIONS_COLUMN_ID } from "@phoenix/pages/project/tableUtils";

type VisibilityState = Record<string, boolean>;
export interface TracingProps {
  /**
   * The project ID for this tracing context
   */
  projectId: string;
  /**
   * Map of the column id to the visibility state
   */
  columnVisibility: VisibilityState;
  /**
   * Map of the annotation column names that are toggled on
   */
  annotationColumnVisibility: VisibilityState;
  /**
   * Map of the trace annotation column names that are toggled on
   */
  traceAnnotationColumnVisibility: VisibilityState;
  /**
   * Map of the column id to the width
   */
  columnSizing: ColumnSizingState;
  /** Order of top-level columns by id. Empty means natural order. */
  columnOrder: string[];
}

export interface TracingState extends TracingProps {
  /**
   * Sets the visibility state of a column
   * @param columnVisibility
   * @returns
   */
  setColumnVisibility: (columnVisibility: VisibilityState) => void;
  /**
   * Sets the visibility of the annotation columns
   */
  setAnnotationColumnVisibility: (
    annotationColumnVisibility: VisibilityState
  ) => void;
  /**
   * Sets the visibility of the trace annotation columns
   */
  setTraceAnnotationColumnVisibility: (
    traceAnnotationColumnVisibility: VisibilityState
  ) => void;
  /**
   * Sets the width of a column
   */
  setColumnSizing: (updater: Updater<ColumnSizingState>) => void;
  /**
   * Sets the order of the top-level columns
   */
  setColumnOrder: (columnOrder: string[]) => void;
}

const makeTracingStoreKey = ({
  projectId,
  tableId,
}: {
  projectId: string;
  tableId: ProjectTab;
}) => `arize-phoenix-tracing-${projectId}-${tableId}`;

export type CreateTracingStoreProps = {
  projectId: string;
  tableId: ProjectTab;
} & Partial<TracingProps>;

export const createTracingStore = (initialProps: CreateTracingStoreProps) => {
  const tracingStore: StateCreator<
    TracingState,
    [["zustand/devtools", unknown]]
  > = (set) => ({
    projectId: initialProps.projectId,
    columnVisibility: {
      metadata: false,
      spanNotes: false,
      traceNotes: false,
      spanId: false,
      traceId: false,
      [TRACE_ANNOTATIONS_COLUMN_ID]: false,
    },
    columnSizing: {
      metadata: 200,
    },
    annotationColumnVisibility: {},
    traceAnnotationColumnVisibility: {},
    columnOrder: [],
    setColumnVisibility: (columnVisibility) => {
      set({ columnVisibility }, false, { type: "setColumnVisibility" });
    },
    setAnnotationColumnVisibility: (annotationColumnVisibility) => {
      set({ annotationColumnVisibility }, false, {
        type: "setAnnotationColumnVisibility",
      });
    },
    setTraceAnnotationColumnVisibility: (traceAnnotationColumnVisibility) => {
      set({ traceAnnotationColumnVisibility }, false, {
        type: "setTraceAnnotationColumnVisibility",
      });
    },
    setColumnOrder: (columnOrder) => {
      set({ columnOrder }, false, { type: "setColumnOrder" });
    },
    setColumnSizing: (columnSizing) => {
      if (typeof columnSizing === "function") {
        set(
          (state) => ({
            columnSizing: columnSizing(state.columnSizing),
          }),
          false,
          { type: "setColumnSizing" }
        );
      } else {
        set({ columnSizing }, false, { type: "setColumnSizing" });
      }
    },
  });
  return create<TracingState>()(
    persist(
      devtools(tracingStore, {
        name: "tracingStore",
      }),
      {
        name: makeTracingStoreKey(initialProps),
      }
    )
  );
};

export type TracingStore = ReturnType<typeof createTracingStore>;
