import * as React from "react";
import { useSearchParams } from "react-router";
import { useColumnsContext } from "../contexts/ColumnsContext";

export function toUrlSafeBase64(input: string): string {
  const utf8Bytes = new TextEncoder().encode(input);
  let base64 = btoa(String.fromCharCode(...utf8Bytes));
  base64 = base64.replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/g, "");
  return base64;
}

export function fromUrlSafeBase64(urlSafeBase64: string): string {
  let base64 = urlSafeBase64.replace(/-/g, "+").replace(/_/g, "/");
  while (base64.length % 4 !== 0) {
    base64 += "=";
  }
  const binaryString = atob(base64);
  const bytes = new Uint8Array(binaryString.length);
  for (let i = 0; i < binaryString.length; i++) {
    bytes[i] = binaryString.charCodeAt(i);
  }
  return new TextDecoder().decode(bytes);
}

export function brevityReplacer(_key: string, value: any): any {
  if (value?.__type === "column" && value.key) {
    return { __type: "columnRef", key: value.key };
  }
  return value;
}

export function parseState<T extends Record<string, any>>(
  search: string,
  reviver: (key: string, value: any) => any,
): T {
  const params = new URLSearchParams(search);
  return Array.from(params.entries()).reduce(
    (memo, [key, value]) => {
      try {
        const decoded = fromUrlSafeBase64(value);
        const parsed = JSON.parse(decoded, reviver);
        memo[key] = parsed;
        return memo;
      } catch (error) {
        memo[key] = value;
        return memo;
      }
    },
    {} as Record<string, any>,
  ) as T;
}

export function useServerState<
  T extends Record<string, any> | null | undefined,
>(key: string, initialValue: T) {
  const [searchParams, setSearchParams] = useSearchParams();
  const $$columns = useColumnsContext();

  // Read the state from the URL on first render
  const readFromUrl = React.useCallback((): T => {
    const encoded = searchParams.get(key);
    if (encoded) {
      try {
        const decoded = fromUrlSafeBase64(encoded);
        return JSON.parse(decoded, (_key, value) => {
          if (value?.__type === "columnRef" && value.key) {
            return $$columns?.[value.key];
          }
          return value;
        }) as T;
      } catch (error) {
        // If decoding or parsing fails, fallback to initialValue
        return initialValue;
      }
    }
    return initialValue;
  }, [key, searchParams, initialValue]);

  const [state, setState] = React.useState<T>(readFromUrl);

  // Whenever state changes, encode & sync it to the URL
  React.useEffect(() => {
    try {
      const newSearchParams = new URLSearchParams(searchParams);
      if (!initialValue && (state === null || state === undefined)) {
        newSearchParams.delete(key);
      } else {
        const encoded = toUrlSafeBase64(JSON.stringify(state, brevityReplacer));
        // Copy current search params so we don't lose other params
        if (encoded) {
          newSearchParams.set(key, encoded);
        } else if (newSearchParams.has(key)) {
          newSearchParams.delete(key);
        }
      }
      setSearchParams(newSearchParams, {
        replace: true,
        preventScrollReset: true,
      });
    } catch (err) {
      // Handle error in encoding if necessary
      console.error("Failed to encode URL state:", err);
    }
  }, [key, searchParams, setSearchParams, state, initialValue]);

  return [state, setState] as const;
}