Core Atoms

/**
 * Create a primitive atom with an initial value
 * @param initialValue - The initial value for the atom
 * @returns PrimitiveAtom that can be read and written
 */
function atom<Value>(initialValue: Value): PrimitiveAtom<Value>;

/**
 * Create a primitive atom without an initial value
 * @returns PrimitiveAtom with undefined initial value
 */
function atom<Value>(): PrimitiveAtom<Value | undefined>;

/**
 * Create a read-only derived atom
 * @param read - Function to compute the atom's value from other atoms
 * @returns Read-only Atom
 */
function atom<Value>(read: Read<Value>): Atom<Value>;

/**
 * Create a writable derived atom with custom read and write logic
 * @param read - Function to compute the atom's value from other atoms
 * @param write - Function to handle write operations
 * @returns WritableAtom with custom behavior
 */
function atom<Value, Args extends unknown[], Result>(
  read: Read<Value, SetAtom<Args, Result>>,
  write: Write<Args, Result>
): WritableAtom<Value, Args, Result>;

/**
 * Create a primitive atom with custom write behavior
 * @param initialValue - The initial value for the atom
 * @param write - Custom write function
 * @returns WritableAtom with custom write behavior
 */
function atom<Value, Args extends unknown[], Result>(
  initialValue: Value,
  write: Write<Args, Result>
): WritableAtom<Value, Args, Result>;

import { atom } from "jotai";

// Primitive atom
const countAtom = atom(0);

// Derived read-only atom
const doubleCountAtom = atom((get) => get(countAtom) * 2);

// Write-only atom (action)
const incrementAtom = atom(null, (get, set) => {
  set(countAtom, get(countAtom) + 1);
});

// Writable derived atom
const uppercaseAtom = atom(
  (get) => get(textAtom).toUpperCase(),
  (get, set, newText: string) => {
    set(textAtom, newText.toLowerCase());
  }
);

// Async atom
const userAtom = atom(async () => {
  const response = await fetch("/api/user");
  return response.json();
});

interface Atom<Value> {
  /** Returns a string representation of the atom */
  toString: () => string;
  /** Function to read the atom's value */
  read: Read<Value>;
  /** Optional function to compare atoms for equality */
  unstable_is?(atom: Atom<unknown>): boolean;
  /** Optional debug label for development */
  debugLabel?: string;
  /** Internal flag for marking atoms as private */
  debugPrivate?: boolean;
  /** Optional initialization callback when atom is first referenced */
  unstable_onInit?(store: Store): void;
}

interface WritableAtom<Value, Args extends unknown[], Result> extends Atom<Value> {
  /** Function to read the atom's value with setSelf capability */
  read: Read<Value, SetAtom<Args, Result>>;
  /** Function to handle write operations */
  write: Write<Args, Result>;
  /** Optional mount callback for side effects */
  onMount?: OnMount<Args, Result>;
}

type PrimitiveAtom<Value> = WritableAtom<Value, [SetStateAction<Value>], void>;

type Read<Value, SetSelf = never> = (
  get: Getter,
  options: { 
    readonly signal: AbortSignal; 
    readonly setSelf: SetSelf 
  }
) => Value;

type Write<Args extends unknown[], Result> = (
  get: Getter,
  set: Setter,
  ...args: Args
) => Result;

type OnMount<Args extends unknown[], Result> = <S extends SetAtom<Args, Result>>(
  setAtom: S
) => OnUnmount | void;

type OnUnmount = () => void;

type SetAtom<Args extends unknown[], Result> = <A extends Args>(
  ...args: A
) => Result;

type Getter = <Value>(atom: Atom<Value>) => Value;

type Setter = <Value, Args extends unknown[], Result>(
  atom: WritableAtom<Value, Args, Result>,
  ...args: Args
) => Result;

type SetStateAction<Value> = Value | ((prev: Value) => Value);

interface Store {
  /** Get the current value of an atom */
  get: <Value>(atom: Atom<Value>) => Value;
  /** Set the value of a writable atom */
  set: <Value, Args extends unknown[], Result>(
    atom: WritableAtom<Value, Args, Result>,
    ...args: Args
  ) => Result;
  /** Subscribe to atom changes */
  sub: (atom: AnyAtom, listener: Listener) => Unsubscribe;
}

type AnyAtom = Atom<unknown>;
type Listener = () => void;
type Unsubscribe = () => void;

/**
 * Create a new isolated atom store
 * @returns New Store instance
 */
function createStore(): Store;

/**
 * Get the default global store, creating one if it doesn't exist
 * @returns The default Store instance
 */
function getDefaultStore(): Store;

/**
 * Internal function to override the default store creation behavior
 * @param fn - Function that returns a new createStore function
 */
function INTERNAL_overrideCreateStore(
  fn: (prev: typeof createStore | undefined) => typeof createStore
): void;

import { createStore, getDefaultStore } from "jotai";

// Create a custom store
const myStore = createStore();

// Get the default store
const defaultStore = getDefaultStore();

// Use custom store
const value = myStore.get(countAtom);
myStore.set(countAtom, 10);

/**
 * Extract the value type from an atom type
 */
type ExtractAtomValue<AtomType> = AtomType extends Atom<infer Value> ? Value : never;

/**
 * Extract the argument types from a writable atom type
 */
type ExtractAtomArgs<AtomType> = AtomType extends WritableAtom<any, infer Args, any> ? Args : never;

/**
 * Extract the result type from a writable atom type
 */
type ExtractAtomResult<AtomType> = AtomType extends WritableAtom<any, any, infer Result> ? Result : never;

import type { ExtractAtomValue, ExtractAtomArgs } from "jotai";

const numberAtom = atom(42);
const incrementAtom = atom(null, (get, set, amount: number) => {
  set(numberAtom, get(numberAtom) + amount);
});

// Extract types from atoms
type NumberValue = ExtractAtomValue<typeof numberAtom>; // number
type IncrementArgs = ExtractAtomArgs<typeof incrementAtom>; // [number]