export default function transform(file, { j }) {
  const root = j(file.source);

  function isForwardRefCall(node) {
    return (
      j.CallExpression.check(node) &&
      (
        (j.Identifier.check(node.callee) && node.callee.name === 'forwardRef') ||
        (j.MemberExpression.check(node.callee) &&
          node.callee.object.name === 'React' &&
          node.callee.property.name === 'forwardRef')
      )
    );
  }

  // 🛑 Check if forwardRef is imported or React.forwardRef is used
  const hasForwardRefImport = root.find(j.ImportDeclaration, {
    source: { value: 'react' },
  }).find(j.ImportSpecifier, {
    imported: { name: 'forwardRef' },
  }).size() > 0;

  const usesReactForwardRef = root.find(j.CallExpression, {
    callee: {
      type: 'MemberExpression',
      object: { name: 'React' },
      property: { name: 'forwardRef' },
    },
  }).size() > 0;

  // 🚪 Exit early if no relevant usage
  if (!hasForwardRefImport && !usesReactForwardRef) {
    return null;
  }

  // Transform: const X = forwardRef(...)
  root.find(j.VariableDeclarator, {
    init: init => isForwardRefCall(init)
  }).forEach(path => {
    const callExpr = path.node.init;
    const innerFn = callExpr.arguments[0];

    if (!j.FunctionExpression.check(innerFn) && !j.ArrowFunctionExpression.check(innerFn)) return;

    applyTypeFix(j, callExpr, innerFn);
    path.node.init = innerFn;
  });

  // Transform: export default forwardRef(...)
  root.find(j.ExportDefaultDeclaration, {
    declaration: d => isForwardRefCall(d)
  }).forEach(path => {
    const callExpr = path.node.declaration;
    const innerFn = callExpr.arguments[0];

    if (!j.FunctionExpression.check(innerFn) && !j.ArrowFunctionExpression.check(innerFn)) return;

    applyTypeFix(j, callExpr, innerFn);
    path.node.declaration = innerFn;
  });

  // Remove forwardRef import
    root
      .find(j.ImportDeclaration, { source: { value: 'react' } })
      .find(j.ImportSpecifier, { imported: { name: 'forwardRef' }})
      .remove();

  return root.toSource();
}

// Applies the TS type transform to move second type arg to props param
function applyTypeFix(j, callExpr, innerFn) {
  const typeParams = callExpr.typeParameters?.params || [];

  if (typeParams.length === 2) {
    const [_, propsType] = typeParams;
    const propsParam = innerFn.params[0];

    if (j.Identifier.check(propsParam) || j.ObjectPattern.check(propsParam)) {
      propsParam.typeAnnotation = j.tsTypeAnnotation(propsType);
    }
  }

  // Remove type params from call
  callExpr.typeParameters = null;
}