import type { StateCreator } from "zustand";
import { create } from "zustand";
import { devtools, persist } from "zustand/middleware";

import type { LastNTimeRangeKey } from "@phoenix/components/datetime/types";
import type { AIQueryModelConfig } from "@phoenix/components/filter/ai/types";
import type { PackageManager, ProgrammingLanguage } from "@phoenix/types/code";
import {
  pythonPackageManagers,
  typescriptPackageManagers,
} from "@phoenix/types/code";
import { getSupportedTimezones } from "@phoenix/utils/timeUtils";

import type { ModelConfig } from "./playground";

/**
 * Per-language package manager preferences.
 */
export type PackageManagerByLanguage = Record<
  ProgrammingLanguage,
  PackageManager
>;

/**
 * Available package managers for each programming language.
 */
export const packageManagersByLanguage: Record<
  ProgrammingLanguage,
  readonly PackageManager[]
> = {
  Python: pythonPackageManagers,
  TypeScript: typescriptPackageManagers,
};

/**
 * The default package manager preference for each language.
 */
const defaultPackageManagerByLanguage: PackageManagerByLanguage = {
  Python: "pip",
  TypeScript: "npm",
};

export type MarkdownDisplayMode = "text" | "markdown";

export const awsBedrockModelPrefixes = [
  "",
  "apac",
  "au",
  "ca",
  "eu",
  "global",
  "il",
  "jp",
  "us",
  "us-gov",
] as const;

export type AwsBedrockModelPrefix = (typeof awsBedrockModelPrefixes)[number];

export type ModelConfigByProvider = Partial<
  Record<
    ModelProvider,
    Omit<ModelConfig, "supportedInvocationParameters" | "customProvider">
  >
>;

export type ProjectViewMode = "table" | "grid";

export type ProjectSortOrder = {
  column: "name" | "endTime";
  direction: "asc" | "desc";
};

export type DisplayTimezone = string;

export interface PreferencesProps {
  /**
   * The display mode of markdown text
   * @default "text"
   */
  markdownDisplayMode: MarkdownDisplayMode;
  /**
   * Whether or not streaming is enabled for a projects traces
   * @default true
   */
  traceStreamingEnabled: boolean;
  /**
   * The last N time range to load data for
   */
  lastNTimeRangeKey: LastNTimeRangeKey;
  /**
   * Whether or not to automatically refresh projects
   * @default true
   */
  projectsAutoRefreshEnabled: boolean;
  /**
   * Whether or not the trace tree shows metrics
   */
  showMetricsInTraceTree: boolean;
  /**
   * Whether the tracing and experiments tables' rows wrap their content or
   * clip it to a single line. A reading preference, so it spans datasets,
   * projects and tabs.
   * @default false
   */
  areTableRowsExpanded: boolean;
  /**
   * {@link ModelConfig|ModelConfigs} for llm providers will be used as the default parameters for the provider
   */
  modelConfigByProvider: ModelConfigByProvider;
  /**
   * Whether or not the playground is in streaming mode
   * Note: this is always false in environments that do not support streaming
   */
  playgroundStreamingEnabled: boolean;
  /**
   * Whether or not the span details are in annotating mode
   */
  isAnnotatingSpans: boolean;
  /**
   * Whether the note-taking bar is shown at the bottom of the span details
   */
  isTakingSpanNotes: boolean;
  /**
   * The view mode for projects
   */
  projectViewMode: ProjectViewMode;
  /**
   * The sort order for projects
   */
  projectSortOrder: ProjectSortOrder;
  /**
   * The last project selected on the dashboards page.
   */
  lastSelectedDashboardProjectId?: string;
  /**
   * Whether the user prefers the side nav to be open or closed. Responsive
   * layout constraints can temporarily override this preference.
   * @default true
   */
  isSideNavExpanded: boolean;
  /**
   * The timezone to display timestamps in
   * @default undefined - use the browser's local timezone
   */
  displayTimezone?: DisplayTimezone;
  /**
   * The preferred programming language for code snippets
   * @default "Python"
   */
  programmingLanguage: ProgrammingLanguage;
  /**
   * The preferred package manager for install commands, per language
   */
  packageManagerByLanguage: PackageManagerByLanguage;
  /**
   * The AWS Bedrock cross-region inference model prefix
   * @default "us"
   * @see https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html
   */
  awsBedrockModelPrefix: AwsBedrockModelPrefix;
  /**
   * Whether or not the assistant agent is enabled
   * @default true
   */
  isAssistantAgentEnabled: boolean;
  /**
   * The user's preferred default model provider for the playground.
   * When undefined, falls back to {@link DEFAULT_MODEL_PROVIDER}.
   */
  defaultModelProvider?: ModelProvider;
  /**
   * The user's preferred default model name for the playground.
   * When undefined, falls back to {@link DEFAULT_MODEL_NAME}.
   */
  defaultModelName?: string;
  /**
   * Whether AI query is enabled on filter condition fields
   * @default true
   */
  isAIQueryEnabled: boolean;
  /**
   * How AI query on filter fields resolves its language model. When
   * undefined, the on-device browser model is used.
   */
  aiQueryModelConfig?: AIQueryModelConfig;
}

export interface PreferencesState extends PreferencesProps {
  /**
   * Sets the display mode of markdown text
   * @param markdownDisplayMode
   */
  setMarkdownDisplayMode: (markdownDisplayMode: MarkdownDisplayMode) => void;
  /**
   * Setter for enabling/disabling trace streaming
   * @param traceStreamingEnabled
   */
  setTraceStreamingEnabled: (traceStreamingEnabled: boolean) => void;
  /**
   * Setter for the last N time range to load data for
   */
  setLastNTimeRangeKey: (lastNTimeRangeKey: LastNTimeRangeKey) => void;
  /**
   * Setter for enabling/disabling project auto refresh
   * @param projectsAutoRefreshEnabled
   */
  setProjectAutoRefreshEnabled: (projectsAutoRefreshEnabled: boolean) => void;
  /**
   * Setter for enabling/disabling metrics in the trace tree
   * @param showMetricsInTraceTree
   */
  setShowMetricsInTraceTree: (showMetricsInTraceTree: boolean) => void;
  /** Setter for whether tracing tables' rows are expanded */
  setAreTableRowsExpanded: (areTableRowsExpanded: boolean) => void;
  /**
   * Setter for the model configs by provider
   */
  setModelConfigForProvider: ({
    provider,
    modelConfig,
  }: {
    provider: ModelProvider;
    modelConfig: Omit<ModelConfig, "supportedInvocationParameters">;
  }) => void;
  /**
   * Setter for enabling/disabling playground streaming
   */
  setPlaygroundStreamingEnabled: (playgroundStreamingEnabled: boolean) => void;
  /**
   * Setter for enabling/disabling span annotating
   */
  setIsAnnotatingSpans: (isAnnotatingSpans: boolean) => void;
  /**
   * Setter for showing/hiding the span note-taking bar
   */
  setIsTakingSpanNotes: (isTakingSpanNotes: boolean) => void;
  /**
   * Setter for the project view mode
   */
  setProjectViewMode: (projectViewMode: ProjectViewMode) => void;
  /**
   * Setter for the project sort order
   */
  setProjectSortOrder: (projectSortOrder: ProjectSortOrder) => void;
  /**
   * Setter for the last project selected on the dashboards page.
   */
  setLastSelectedDashboardProjectId: (projectId: string) => void;
  /**
   * Setter for the user's preferred side nav open state
   */
  setIsSideNavExpanded: (isSideNavExpanded: boolean) => void;
  /**
   * Setter for the display timezone
   */
  setDisplayTimezone: (displayTimezone: DisplayTimezone | undefined) => void;
  /**
   * Setter for the preferred programming language
   */
  setProgrammingLanguage: (programmingLanguage: ProgrammingLanguage) => void;
  /**
   * Setter for the preferred package manager for a given language
   */
  setPackageManager: (
    language: ProgrammingLanguage,
    packageManager: PackageManager
  ) => void;
  /**
   * Setter for the AWS Bedrock model prefix
   */
  setAwsBedrockModelPrefix: (
    awsBedrockModelPrefix: AwsBedrockModelPrefix
  ) => void;
  /**
   * Setter for enabling/disabling the assistant agent
   */
  setIsAssistantAgentEnabled: (isAssistantAgentEnabled: boolean) => void;
  /**
   * Setter for the user's preferred default model provider for the playground.
   * Pass `undefined` to clear the preference.
   */
  setDefaultModelProvider: (
    defaultModelProvider: ModelProvider | undefined
  ) => void;
  /**
   * Setter for the user's preferred default model name for the playground.
   * Pass `undefined` (or an empty string) to clear the preference.
   */
  setDefaultModelName: (defaultModelName: string | undefined) => void;
  /**
   * Setter for enabling/disabling AI query on filter condition fields
   */
  setIsAIQueryEnabled: (isAIQueryEnabled: boolean) => void;
  /**
   * Setter for how AI query resolves its language model. Pass `undefined`
   * to fall back to the on-device browser model.
   */
  setAIQueryModelConfig: (
    aiQueryModelConfig: AIQueryModelConfig | undefined
  ) => void;
}

export const createPreferencesStore = (
  initialProps?: Partial<PreferencesProps>
) => {
  const preferencesStore: StateCreator<
    PreferencesState,
    [["zustand/devtools", unknown]]
  > = (set) => ({
    markdownDisplayMode: "text",
    setMarkdownDisplayMode: (markdownDisplayMode) => {
      set({ markdownDisplayMode }, false, { type: "setMarkdownDisplayMode" });
    },
    traceStreamingEnabled: true,
    setTraceStreamingEnabled: (traceStreamingEnabled) => {
      set({ traceStreamingEnabled }, false, {
        type: "setTraceStreamingEnabled",
      });
    },
    lastNTimeRangeKey: "7d",
    setLastNTimeRangeKey: (lastNTimeRangeKey) => {
      set({ lastNTimeRangeKey });
    },
    projectsAutoRefreshEnabled: true,
    setProjectAutoRefreshEnabled: (projectsAutoRefreshEnabled) => {
      set({ projectsAutoRefreshEnabled }, false, {
        type: "setProjectAutoRefreshEnabled",
      });
    },
    showMetricsInTraceTree: true,
    setShowMetricsInTraceTree: (showMetricsInTraceTree) => {
      set({ showMetricsInTraceTree }, false, {
        type: "setShowMetricsInTraceTree",
      });
    },
    // Clipped to start: an even row height is what makes a table scan as a grid
    areTableRowsExpanded: false,
    setAreTableRowsExpanded: (areTableRowsExpanded) => {
      set({ areTableRowsExpanded }, false, {
        type: "setAreTableRowsExpanded",
      });
    },
    modelConfigByProvider: {},
    setModelConfigForProvider: ({ provider, modelConfig }) => {
      set(
        (state) => {
          return {
            modelConfigByProvider: {
              ...state.modelConfigByProvider,
              [provider]: modelConfig,
            },
          };
        },
        false,
        { type: "setModelConfigForProvider" }
      );
    },
    playgroundStreamingEnabled: true,
    setPlaygroundStreamingEnabled: (playgroundStreamingEnabled) => {
      set({ playgroundStreamingEnabled }, false, {
        type: "setPlaygroundStreamingEnabled",
      });
    },
    isAnnotatingSpans: false,
    setIsAnnotatingSpans: (isAnnotatingSpans) => {
      set({ isAnnotatingSpans }, false, { type: "setIsAnnotatingSpans" });
    },
    isTakingSpanNotes: false,
    setIsTakingSpanNotes: (isTakingSpanNotes) => {
      set({ isTakingSpanNotes }, false, { type: "setIsTakingSpanNotes" });
    },
    projectViewMode: "grid",
    setProjectViewMode: (projectViewMode) => {
      set({ projectViewMode }, false, { type: "setProjectViewMode" });
    },
    projectSortOrder: {
      column: "endTime",
      direction: "desc",
    },
    setProjectSortOrder: (projectSortOrder) => {
      set({ projectSortOrder }, false, { type: "setProjectSortOrder" });
    },
    lastSelectedDashboardProjectId: undefined,
    setLastSelectedDashboardProjectId: (projectId) => {
      set({ lastSelectedDashboardProjectId: projectId }, false, {
        type: "setLastSelectedDashboardProjectId",
      });
    },
    isSideNavExpanded: true,
    setIsSideNavExpanded: (isSideNavExpanded) => {
      set({ isSideNavExpanded }, false, { type: "setIsSideNavExpanded" });
    },
    setDisplayTimezone: (displayTimezone) => {
      // Just to be extra safe of what we store in local storage.
      if (
        displayTimezone &&
        !getSupportedTimezones().includes(displayTimezone)
      ) {
        throw new Error(`Invalid timezone: ${displayTimezone}`);
      }
      set({ displayTimezone }, false, { type: "setDisplayTimezone" });
    },
    programmingLanguage: "Python",
    setProgrammingLanguage: (programmingLanguage) => {
      set({ programmingLanguage }, false, { type: "setProgrammingLanguage" });
    },
    packageManagerByLanguage: { ...defaultPackageManagerByLanguage },
    setPackageManager: (language, packageManager) => {
      set(
        (state) => ({
          packageManagerByLanguage: {
            ...state.packageManagerByLanguage,
            [language]: packageManager,
          },
        }),
        false,
        { type: "setPackageManager" }
      );
    },
    awsBedrockModelPrefix: "us",
    setAwsBedrockModelPrefix: (awsBedrockModelPrefix) => {
      set({ awsBedrockModelPrefix }, false, {
        type: "setAwsBedrockModelPrefix",
      });
    },
    isAssistantAgentEnabled: true,
    setIsAssistantAgentEnabled: (isAssistantAgentEnabled) => {
      set({ isAssistantAgentEnabled }, false, {
        type: "setIsAssistantAgentEnabled",
      });
    },
    defaultModelProvider: undefined,
    setDefaultModelProvider: (defaultModelProvider) => {
      set({ defaultModelProvider }, false, { type: "setDefaultModelProvider" });
    },
    defaultModelName: undefined,
    setDefaultModelName: (defaultModelName) => {
      const trimmed = defaultModelName?.trim();
      set({ defaultModelName: trimmed ? trimmed : undefined }, false, {
        type: "setDefaultModelName",
      });
    },
    isAIQueryEnabled: true,
    setIsAIQueryEnabled: (isAIQueryEnabled) => {
      set({ isAIQueryEnabled }, false, { type: "setIsAIQueryEnabled" });
    },
    aiQueryModelConfig: undefined,
    setAIQueryModelConfig: (aiQueryModelConfig) => {
      set({ aiQueryModelConfig }, false, { type: "setAIQueryModelConfig" });
    },
    ...initialProps,
  });
  return create<PreferencesState>()(
    persist(devtools(preferencesStore, { name: "preferencesStore" }), {
      name: "arize-phoenix-preferences",
    })
  );
};

export type PreferencesStore = ReturnType<typeof createPreferencesStore>;
