'use client';

import * as React from 'react';
import {
  createContext,
  useCallback,
  useEffect,
  useMemo,
  useReducer,
  useRef,
} from 'react';

import {
  DEFAULT_TREE_LABELS,
  type FlatRow,
  type TreeContextMenuSlot,
  type TreeItemId,
  type TreeLabels,
  type TreeLoadChildren,
  type TreeNode,
  type TreeRootProps,
  type TreeRowSlot,
  type TreeSelectionMode,
} from '../types';
import {
  createChildCache,
  type ChildCache,
  type ChildEntry,
} from '../data/childCache';
import { flattenTree } from '../data/flatten';
import { loadTreeState, saveTreeState } from '../data/persist';
import {
  resolveAppearance,
  type ResolvedAppearance,
  type TreeAppearance,
} from '../data/appearance';

// =====================================================================
// Reducer
// =====================================================================

interface State<T> {
  expanded: Set<TreeItemId>;
  selected: Set<TreeItemId>;
  focused: TreeItemId | null;
  query: string;
  /** Bumped on every cache mutation so memos see a fresh dep. */
  cacheTick: number;
}

type Action<T> =
  | { type: 'expand'; id: TreeItemId }
  | { type: 'collapse'; id: TreeItemId }
  | { type: 'toggle'; id: TreeItemId }
  | { type: 'set-expanded'; ids: TreeItemId[] }
  | { type: 'select'; id: TreeItemId; mode: TreeSelectionMode }
  | { type: 'select-many'; ids: TreeItemId[] }
  | { type: 'clear-selection' }
  | { type: 'focus'; id: TreeItemId | null }
  | { type: 'set-query'; q: string }
  | { type: 'cache-tick' };

const reducer = <T,>(state: State<T>, action: Action<T>): State<T> => {
  switch (action.type) {
    case 'expand': {
      if (state.expanded.has(action.id)) return state;
      const next = new Set(state.expanded);
      next.add(action.id);
      return { ...state, expanded: next };
    }
    case 'collapse': {
      if (!state.expanded.has(action.id)) return state;
      const next = new Set(state.expanded);
      next.delete(action.id);
      return { ...state, expanded: next };
    }
    case 'toggle': {
      const next = new Set(state.expanded);
      if (next.has(action.id)) next.delete(action.id);
      else next.add(action.id);
      return { ...state, expanded: next };
    }
    case 'set-expanded':
      return { ...state, expanded: new Set(action.ids) };
    case 'select': {
      if (action.mode === 'none') return state;
      if (action.mode === 'single') {
        return { ...state, selected: new Set([action.id]), focused: action.id };
      }
      const next = new Set(state.selected);
      if (next.has(action.id)) next.delete(action.id);
      else next.add(action.id);
      return { ...state, selected: next, focused: action.id };
    }
    case 'select-many':
      return { ...state, selected: new Set(action.ids) };
    case 'clear-selection':
      return { ...state, selected: new Set() };
    case 'focus':
      return { ...state, focused: action.id };
    case 'set-query':
      return { ...state, query: action.q };
    case 'cache-tick':
      return { ...state, cacheTick: state.cacheTick + 1 };
    default:
      return state;
  }
};

// =====================================================================
// Context value
// =====================================================================

export interface TreeContextValue<T> {
  // State
  expanded: ReadonlySet<TreeItemId>;
  selected: ReadonlySet<TreeItemId>;
  focused: TreeItemId | null;
  query: string;

  // Flattened render rows (visible items only)
  flatRows: FlatRow<T>[];
  /** Search-matching node ids (subset of all flatRows). */
  matchingIds: ReadonlySet<TreeItemId>;

  // Imperative actions
  expand: (id: TreeItemId) => void;
  collapse: (id: TreeItemId) => void;
  toggle: (id: TreeItemId) => void;
  expandAll: () => void;
  collapseAll: () => void;
  select: (id: TreeItemId) => void;
  setSelectedIds: (ids: TreeItemId[]) => void;
  clearSelection: () => void;
  setFocus: (id: TreeItemId | null) => void;
  setQuery: (q: string) => void;
  refresh: (id: TreeItemId) => Promise<void>;
  refreshAll: () => Promise<void>;
  activate: (node: TreeNode<T>) => void;

  // Config / slots
  labels: TreeLabels;
  /** Resolved cosmetic config — never null. */
  appearance: ResolvedAppearance;
  /** Convenience alias for `appearance.indent`. */
  indent: number;
  selectionMode: TreeSelectionMode;
  enableSearch: boolean;
  showIndentGuides: boolean;
  getItemName: (node: TreeNode<T>) => string;

  renderIcon?: TreeRowSlot<T>;
  renderLabel?: TreeRowSlot<T>;
  renderActions?: TreeRowSlot<T>;
  renderContextMenu?: TreeContextMenuSlot<T>;
}

const TreeContext = createContext<TreeContextValue<unknown> | null>(null);

export function useTreeContext<T>(): TreeContextValue<T> {
  const ctx = React.useContext(TreeContext);
  if (!ctx) {
    throw new Error('useTreeContext must be used inside <TreeProvider>');
  }
  return ctx as TreeContextValue<T>;
}

// =====================================================================
// Provider
// =====================================================================

export interface TreeProviderProps<T>
  extends Pick<
    TreeRootProps<T>,
    | 'data'
    | 'getItemName'
    | 'loadChildren'
    | 'selectionMode'
    | 'initialExpandedIds'
    | 'initialSelectedIds'
    | 'indent'
    | 'appearance'
    | 'onSelectionChange'
    | 'onExpansionChange'
    | 'onActivate'
    | 'enableSearch'
    | 'showIndentGuides'
    | 'renderIcon'
    | 'renderLabel'
    | 'renderActions'
    | 'renderContextMenu'
    | 'labels'
    | 'persistKey'
    | 'persistSelection'
  > {
  children: React.ReactNode;
}

const setEqualsArr = (set: ReadonlySet<string>, arr: readonly string[]) => {
  if (set.size !== arr.length) return false;
  for (const id of arr) if (!set.has(id)) return false;
  return true;
};

const collectAllIds = <T,>(
  roots: TreeNode<T>[],
  cache: ChildCache<T>,
  out: TreeItemId[],
) => {
  for (const node of roots) {
    if (Array.isArray(node.children)) {
      out.push(node.id);
      collectAllIds(node.children, cache, out);
    } else if (node.isFolder) {
      out.push(node.id);
      const entry = cache.get(node.id);
      if (entry?.children) collectAllIds(entry.children, cache, out);
    }
  }
};

export function TreeProvider<T>(props: TreeProviderProps<T>) {
  const {
    data,
    getItemName,
    loadChildren,
    selectionMode = 'single',
    initialExpandedIds,
    initialSelectedIds,
    indent,
    appearance,
    onSelectionChange,
    onExpansionChange,
    onActivate,
    enableSearch = false,
    showIndentGuides = false,
    renderIcon,
    renderLabel,
    renderActions,
    renderContextMenu,
    labels: labelsOverride,
    persistKey,
    persistSelection = false,
    children,
  } = props;

  const labels = useMemo(
    () => ({ ...DEFAULT_TREE_LABELS, ...labelsOverride }),
    [labelsOverride],
  );

  const resolvedAppearance = useMemo(
    () => resolveAppearance(appearance, indent),
    [appearance, indent],
  );

  // Persisted state, loaded once.
  const persisted = useMemo(
    () => (persistKey ? loadTreeState(persistKey) : null),
    [persistKey],
  );

  const [state, dispatch] = useReducer(reducer<T>, undefined, () => ({
    expanded: new Set(persisted?.expandedItems ?? initialExpandedIds ?? []),
    selected: new Set(
      (persistSelection ? persisted?.selectedItems : undefined) ??
        initialSelectedIds ??
        [],
    ),
    focused: null,
    query: '',
    cacheTick: 0,
  }));

  // Async cache survives provider re-renders, lives in a ref.
  const cacheRef = useRef<ChildCache<T>>(createChildCache<T>());
  const inflightRef = useRef<Map<TreeItemId, Promise<void>>>(new Map());

  // Trigger one fetch per (folder id) — concurrent expansions are deduped.
  const fetchChildren = useCallback(
    async (node: TreeNode<T>) => {
      if (!loadChildren) return;
      if (Array.isArray(node.children)) return;
      const existing = cacheRef.current.get(node.id);
      if (existing?.status === 'loaded' || existing?.status === 'loading') return;
      const inflight = inflightRef.current.get(node.id);
      if (inflight) return inflight;

      cacheRef.current.set(node.id, { status: 'loading', children: [] });
      dispatch({ type: 'cache-tick' });

      const promise = (async () => {
        try {
          const children = await loadChildren(node);
          cacheRef.current.set(node.id, { status: 'loaded', children });
        } catch (err) {
          cacheRef.current.set(node.id, {
            status: 'error',
            children: [],
            error: err instanceof Error ? err.message : String(err),
          });
        } finally {
          inflightRef.current.delete(node.id);
          dispatch({ type: 'cache-tick' });
        }
      })();

      inflightRef.current.set(node.id, promise);
      return promise;
    },
    [loadChildren],
  );

  // Build a quick id → node map for any visible node (incl. cached children).
  const nodeById = useMemo(() => {
    const map = new Map<TreeItemId, TreeNode<T>>();
    const walk = (nodes: TreeNode<T>[]) => {
      for (const n of nodes) {
        map.set(n.id, n);
        if (Array.isArray(n.children)) walk(n.children);
        else {
          const entry = cacheRef.current.get(n.id);
          if (entry?.children) walk(entry.children);
        }
      }
    };
    walk(data);
    return map;
  }, [data, state.cacheTick]);

  // On expand, kick async fetch (no-op if already cached or has inline children).
  useEffect(() => {
    if (!loadChildren) return;
    for (const id of state.expanded) {
      const node = nodeById.get(id);
      if (!node) continue;
      void fetchChildren(node);
    }
  }, [loadChildren, state.expanded, state.cacheTick, nodeById, fetchChildren]);

  // Flatten on relevant changes.
  const flatRows = useMemo(
    () =>
      flattenTree<T>({
        roots: data,
        expandedIds: state.expanded,
        cache: cacheRef.current,
      }),
    [data, state.expanded, state.cacheTick],
  );

  // Search matches (case-insensitive substring on getItemName).
  const matchingIds = useMemo(() => {
    const set = new Set<TreeItemId>();
    if (!enableSearch || state.query.trim() === '') return set;
    const q = state.query.trim().toLowerCase();
    for (const row of flatRows) {
      if (getItemName(row.node).toLowerCase().includes(q)) {
        set.add(row.node.id);
      }
    }
    return set;
  }, [enableSearch, state.query, flatRows, getItemName]);

  // External callbacks via stable refs.
  const onSelectionChangeRef = useRef(onSelectionChange);
  const onExpansionChangeRef = useRef(onExpansionChange);
  const onActivateRef = useRef(onActivate);
  onSelectionChangeRef.current = onSelectionChange;
  onExpansionChangeRef.current = onExpansionChange;
  onActivateRef.current = onActivate;

  // Notify on changes + persist.
  const lastSelectedArrRef = useRef<TreeItemId[]>([...state.selected]);
  const lastExpandedArrRef = useRef<TreeItemId[]>([...state.expanded]);

  useEffect(() => {
    const arr = [...state.expanded];
    if (!setEqualsArr(state.expanded, lastExpandedArrRef.current)) {
      lastExpandedArrRef.current = arr;
      onExpansionChangeRef.current?.(arr);
      if (persistKey) {
        saveTreeState(persistKey, {
          expandedItems: arr,
          selectedItems: persistSelection ? [...state.selected] : [],
        });
      }
    }
  }, [state.expanded, persistKey, persistSelection, state.selected]);

  useEffect(() => {
    const arr = [...state.selected];
    if (!setEqualsArr(state.selected, lastSelectedArrRef.current)) {
      lastSelectedArrRef.current = arr;
      onSelectionChangeRef.current?.(arr);
      if (persistKey && persistSelection) {
        saveTreeState(persistKey, {
          expandedItems: [...state.expanded],
          selectedItems: arr,
        });
      }
    }
  }, [state.selected, persistKey, persistSelection, state.expanded]);

  // Imperative actions.
  const expand = useCallback((id: TreeItemId) => dispatch({ type: 'expand', id }), []);
  const collapse = useCallback((id: TreeItemId) => dispatch({ type: 'collapse', id }), []);
  const toggle = useCallback((id: TreeItemId) => dispatch({ type: 'toggle', id }), []);

  const expandAll = useCallback(() => {
    const ids: TreeItemId[] = [];
    collectAllIds(data, cacheRef.current, ids);
    dispatch({ type: 'set-expanded', ids });
  }, [data]);

  const collapseAll = useCallback(
    () => dispatch({ type: 'set-expanded', ids: [] }),
    [],
  );

  const select = useCallback(
    (id: TreeItemId) => dispatch({ type: 'select', id, mode: selectionMode }),
    [selectionMode],
  );
  const setSelectedIds = useCallback(
    (ids: TreeItemId[]) => dispatch({ type: 'select-many', ids }),
    [],
  );
  const clearSelection = useCallback(() => dispatch({ type: 'clear-selection' }), []);
  const setFocus = useCallback(
    (id: TreeItemId | null) => dispatch({ type: 'focus', id }),
    [],
  );
  const setQuery = useCallback((q: string) => dispatch({ type: 'set-query', q }), []);

  const refresh = useCallback(
    async (id: TreeItemId) => {
      const node = nodeById.get(id);
      if (!node || !loadChildren) return;
      cacheRef.current.delete(id);
      dispatch({ type: 'cache-tick' });
      await fetchChildren(node);
    },
    [nodeById, loadChildren, fetchChildren],
  );

  const refreshAll = useCallback(async () => {
    cacheRef.current.clear();
    dispatch({ type: 'cache-tick' });
    if (!loadChildren) return;
    await Promise.all(
      [...state.expanded].map((id) => {
        const node = nodeById.get(id);
        return node ? fetchChildren(node) : undefined;
      }),
    );
  }, [loadChildren, state.expanded, nodeById, fetchChildren]);

  const activate = useCallback(
    (node: TreeNode<T>) => onActivateRef.current?.(node),
    [],
  );

  const value = useMemo<TreeContextValue<T>>(
    () => ({
      expanded: state.expanded,
      selected: state.selected,
      focused: state.focused,
      query: state.query,
      flatRows,
      matchingIds,
      expand,
      collapse,
      toggle,
      expandAll,
      collapseAll,
      select,
      setSelectedIds,
      clearSelection,
      setFocus,
      setQuery,
      refresh,
      refreshAll,
      activate,
      labels,
      appearance: resolvedAppearance,
      indent: resolvedAppearance.indent,
      selectionMode,
      enableSearch,
      showIndentGuides,
      getItemName,
      renderIcon,
      renderLabel,
      renderActions,
      renderContextMenu,
    }),
    [
      state.expanded,
      state.selected,
      state.focused,
      state.query,
      flatRows,
      matchingIds,
      expand,
      collapse,
      toggle,
      expandAll,
      collapseAll,
      select,
      setSelectedIds,
      clearSelection,
      setFocus,
      setQuery,
      refresh,
      refreshAll,
      activate,
      labels,
      resolvedAppearance,
      selectionMode,
      enableSearch,
      showIndentGuides,
      getItemName,
      renderIcon,
      renderLabel,
      renderActions,
      renderContextMenu,
    ],
  );

  return (
    <TreeContext.Provider value={value as TreeContextValue<unknown>}>
      {children}
    </TreeContext.Provider>
  );
}

// Re-export internal types referenced by hook consumers.
export type { ChildCache, ChildEntry };
