import {
  isValidElement,
  useMemo,
  type ComponentType,
  type ReactElement,
} from "react";

type UhuruCodeTarget =
  | ComponentType<any>
  | ReactElement
  | string;

type UhuruCodeOptions = {
  id?: string;
};

const componentCodeRegistry = new WeakMap<object, string>();
const componentIdRegistry = new WeakMap<object, string>();
const namedCodeRegistry = new Map<string, string>();

function resolveComponentTarget(target: UhuruCodeTarget) {
  if (typeof target === "string") {
    return { component: null, id: target };
  }

  if (isValidElement(target)) {
    const elementType = target.type;
    return {
      component: typeof elementType === "function" || typeof elementType === "object" ? elementType : null,
      id: typeof elementType === "string" ? elementType : null,
    };
  }

  return { component: target, id: null };
}

export function registerUhuruCode(
  target: ComponentType<any> | string,
  code: string,
  options: UhuruCodeOptions = {},
) {
  if (typeof target === "string") {
    namedCodeRegistry.set(options.id ?? target, code);
    return target;
  }

  componentCodeRegistry.set(target, code);

  const registryId = options.id ?? target.displayName ?? target.name;
  if (registryId) {
    componentIdRegistry.set(target, registryId);
    namedCodeRegistry.set(registryId, code);
  }

  return target;
}

export function showCode(target: UhuruCodeTarget) {
  const resolved = resolveComponentTarget(target);

  if (resolved.component && componentCodeRegistry.has(resolved.component)) {
    return componentCodeRegistry.get(resolved.component) ?? "";
  }

  if (resolved.component) {
    const registryId = componentIdRegistry.get(resolved.component);
    if (registryId) {
      return namedCodeRegistry.get(registryId) ?? "";
    }
  }

  if (resolved.id) {
    return namedCodeRegistry.get(resolved.id) ?? "";
  }

  return "";
}

export function useUhuruCode(target: UhuruCodeTarget) {
  return useMemo(() => showCode(target), [target]);
}

export function getUhuruCodeById(id: string) {
  return namedCodeRegistry.get(id) ?? "";
}