import { isAbsolute, relative, resolve } from 'node:path';
import ts from 'typescript';
import type { Plugin } from 'vite';

const escapeAttribute = (value: string): string => value.replace(/&/g, '&amp;').replace(/"/g, '&quot;').replace(/</g, '&lt;').replace(/>/g, '&gt;');

const enclosingComponent = (node: ts.Node): string | undefined => {
  for (let parent = node.parent; parent; parent = parent.parent) {
    if ((ts.isFunctionDeclaration(parent) || ts.isFunctionExpression(parent)) && parent.name && /^[A-Z]/.test(parent.name.text)) {
      return parent.name.text;
    }
    if (ts.isVariableDeclaration(parent) && ts.isIdentifier(parent.name) && /^[A-Z]/.test(parent.name.text)) {
      return parent.name.text;
    }
    if (ts.isClassDeclaration(parent) && parent.name && /^[A-Z]/.test(parent.name.text)) return parent.name.text;
  }
  return undefined;
};

/** Adds only repository-relative location and component identity; never reads JSX values. */
export const annotateUiSource = (code: string, id: string, uiRoot: string): string | null => {
  if (id.includes('?') || id.includes('\0') || !isAbsolute(id)) return null;
  const path = relative(resolve(uiRoot, 'src'), id).replace(/\\/g, '/');
  if (path.startsWith('../') || isAbsolute(path) || !path.endsWith('.tsx')) return null;
  if (/(?:^|\/)(?:node_modules|vendor|__tests__|tests?|__mocks__)(?:\/|$)/.test(path) || /\.(?:test|spec)\.tsx$/.test(path)) return null;
  const source = ts.createSourceFile(id, code, ts.ScriptTarget.Latest, true, ts.ScriptKind.TSX);
  const insertions: { at: number; text: string }[] = [];
  const visit = (node: ts.Node): void => {
    if ((ts.isJsxOpeningElement(node) || ts.isJsxSelfClosingElement(node)) && ts.isIdentifier(node.tagName) && /^[a-z]/.test(node.tagName.text)) {
      const existing = new Set(node.attributes.properties.filter(ts.isJsxAttribute).map(attribute => attribute.name.getText(source)));
      const { line, character } = source.getLineAndCharacterOfPosition(node.getStart(source));
      let text = existing.has('data-ui-source') ? '' : ` data-ui-source="${escapeAttribute(`ui/src/${path}:${line + 1}:${character + 1}`)}"`;
      const component = enclosingComponent(node);
      if (component && !existing.has('data-ui-component')) text += ` data-ui-component="${escapeAttribute(component)}"`;
      // Raw HTML descendants may carry arbitrary data-* attributes. Mark the
      // trusted JSX host so inspection can resolve to its source instead.
      if (existing.has('dangerouslySetInnerHTML') && !existing.has('data-ui-content-boundary')) text += ' data-ui-content-boundary="true"';
      if (text) insertions.push({ at: node.tagName.end, text });
    }
    ts.forEachChild(node, visit);
  };
  visit(source);
  // Reverse insertions preserve every original character and all original source coordinates.
  let output = code;
  for (const insertion of insertions.sort((a, b) => b.at - a.at)) output = output.slice(0, insertion.at) + insertion.text + output.slice(insertion.at);
  return output;
};

export const uiSourcePlugin = (uiRoot: string): Plugin => ({
  name: 'ui-source',
  enforce: 'pre',
  transform(code, id) {
    const annotated = annotateUiSource(code, id, uiRoot);
    if (annotated === null || annotated === code) return null;
    // No mapping library is a direct dependency. Source attributes retain original
    // coordinates; this insertion-only transform does not claim a generated sourcemap.
    return { code: annotated, map: null };
  },
});