import { parse } from '@babel/parser';
import fs from 'node:fs';
import path from 'node:path';

/**
 * Vite plugin — "component docs" virtual module.
 *
 * Reads JSDoc blocks tagged `@component` from the source files listed in
 * SOURCES, parses them into a registry, and exposes it to the dev gallery as
 * `import { componentDocs } from 'virtual:component-docs'`.
 *
 * The JSDoc format (kept next to the implementation):
 *
 *   Block comment: "/**" ... "*​/" (see parseDoc for tag handling)
 *
 *   @component PanelTile        ← component name used in the gallery
 *   @import dashgui/panels      ← sub-path export for the import snippet
 *   @defaultExport              ← optional: `import X from '…'` (default export)
 *
 *   @prop {type} [name=default] - Description   ← optional props use [name=default]
 *   @prop {type} requiredProp - Description     ← props without brackets are required
 *   @prop {type} [optionalProp] - Description   ← optional without a default
 *
 *   @behaviour A behaviour note.
 *   @behaviour Another behaviour note.
 */
const VIRTUAL_ID = 'virtual:component-docs';
const RESOLVED_VIRTUAL_ID = '\0' + VIRTUAL_ID;

const SOURCES = [
  { file: 'src/panels/PanelTile.jsx', importPath: 'dashgui/panels' },
  { file: 'src/panels/EditablePanelTile.jsx', importPath: 'dashgui/panels' },
  { file: 'src/ElevatedPill.jsx', importPath: 'dashgui/ElevatedPill' },
  { file: 'src/SingleLineTextEdit.jsx', importPath: 'dashgui/SingleLineTextEdit' },
  { file: 'src/tiles/Tile.jsx', importPath: 'dashgui/tiles' },
  { file: 'src/lists/RecordList.jsx', importPath: 'dashgui/lists' },
];

/** Parse a single `@prop {type} [name=default] - description` line. */
function parseProp(text) {
  const typeMatch = text.match(/^\{(.*?)\}\s+(.*)$/s);
  const type = typeMatch ? typeMatch[1].trim() : '';
  const rest = typeMatch ? typeMatch[2].trim() : text.trim();

  let name;
  let defaultValue;
  let required = true;
  let description = '';

  const optionalMatch = rest.match(/^\[([^\]=]*)(?:=([^\]]*))?\]\s*(?:-\s*)?(.*)$/s);
  if (optionalMatch) {
    required = false;
    name = optionalMatch[1].trim();
    defaultValue = optionalMatch[2] !== undefined ? optionalMatch[2].trim() : undefined;
    description = optionalMatch[3].trim();
  } else {
    const reqMatch = rest.match(/^([^\s-]+)(?:\s*-\s*)?(.*)$/s);
    if (reqMatch) {
      name = reqMatch[1].trim();
      description = reqMatch[2].trim();
    }
  }

  return { name, type, required, default: defaultValue, description };
}

/** Parse a `/** … *​/` comment value (as provided by @babel/parser). */
function parseDoc(commentValue, fallback) {
  const lines = commentValue
    .split('\n')
    .map((l) => l.replace(/^\s*\*\s?/, '').trim())
    .filter((l, idx, arr) => !(l === '' && (idx === 0 || idx === arr.length - 1 || arr[idx - 1] === '')));

  const descLines = [];
  const props = [];
  const behaviour = [];
  let name = fallback.name || '';
  let importPath = fallback.importPath || '';
  let defaultExport = false;

  let i = 0;
  while (i < lines.length && !lines[i].startsWith('@')) {
    if (lines[i] !== '') descLines.push(lines[i]);
    i++;
  }

  for (; i < lines.length; i++) {
    const m = lines[i].match(/^@(\w+)\s*(.*)$/);
    if (!m) continue;
    const [, tag, rest] = m;
    switch (tag) {
      case 'component':
        name = rest.trim();
        break;
      case 'import':
        importPath = rest.trim();
        break;
      case 'defaultExport':
        defaultExport = true;
        break;
      case 'prop':
        props.push(parseProp(rest.trim()));
        break;
      case 'behaviour':
        behaviour.push(rest.trim());
        break;
      default:
        break;
    }
  }

  return { name, importPath, defaultExport, description: descLines.join(' '), props, behaviour };
}

/** Walk a source file and collect every `@component`-tagged JSDoc block. */
function extractFromFile(root, source) {
  const fullPath = path.join(root, source.file);
  let code;
  try {
    code = fs.readFileSync(fullPath, 'utf8');
  } catch (err) {
    console.warn(`[component-docs] cannot read ${source.file}: ${err.message}`);
    return {};
  }

  let ast;
  try {
    ast = parse(code, {
      sourceType: 'module',
      plugins: ['jsx', 'classProperties'],
      errorRecovery: true,
    });
  } catch (err) {
    console.warn(`[component-docs] cannot parse ${source.file}: ${err.message}`);
    return {};
  }

  const found = {};
  for (const stmt of ast.program.body) {
    const comments = stmt.leadingComments || [];
    const doc = comments[comments.length - 1];
    if (!doc || doc.type !== 'CommentBlock' || !doc.value.includes('@component')) continue;

    const parsed = parseDoc(doc.value, { name: '', importPath: source.importPath });
    if (!parsed.name) continue;
    found[parsed.name] = parsed;
  }
  return found;
}

function buildRegistry(root) {
  const registry = {};
  for (const source of SOURCES) {
    Object.assign(registry, extractFromFile(root, source));
  }
  return registry;
}

export function componentDocsPlugin() {
  const documentedFiles = SOURCES.map((s) => s.file);

  return {
    name: 'component-docs',
    enforce: 'pre',

    resolveId(id) {
      if (id === VIRTUAL_ID) return RESOLVED_VIRTUAL_ID;
      return null;
    },

    load(id) {
      if (id !== RESOLVED_VIRTUAL_ID) return null;
      const registry = buildRegistry(process.cwd());
      return `export const componentDocs = ${JSON.stringify(registry)};`;
    },

    handleHotUpdate(ctx) {
      const { file, server } = ctx;
      if (!documentedFiles.some((f) => file.endsWith(f))) return;
      const mod = server.moduleGraph.getModuleById(RESOLVED_VIRTUAL_ID);
      if (mod) server.moduleGraph.invalidateModule(mod);
      server.ws.send({ type: 'full-reload' });
    },
  };
}