src/services/astParser.ts

/**
 * src/services/astParser.ts
 *
 * Parses JS/JSX/TS/TSX source files using Babel's AST parser to produce
 * a lightweight Skeleton.json describing the app's structure.
 *
 * WHY AST INSTEAD OF SENDING RAW CODE TO AN LLM?
 *   Sending an entire React codebase to an LLM context window is expensive,
 *   slow, and risks hitting token limits. Instead, we extract only what the
 *   architecture-planning LLM needs: component names, state shapes, hooks
 *   used, API endpoints, and routing information.
 *
 * This module runs ENTIRELY locally — zero LLM calls.
 */

import * as babelParser from '@babel/parser';
import _traverse from '@babel/traverse';
import * as t from '@babel/types';
import path from 'path';
import { logger } from '../utils/logger';
import type { ExtractedFile } from './zipService';

// @babel/traverse ships CJS + ESM; handle the default-export dance safely
const traverse = (
  typeof (_traverse as unknown as { default: unknown }).default === 'function'
    ? (_traverse as unknown as { default: typeof _traverse }).default
    : _traverse
) as typeof _traverse;

// ─── Skeleton Types ───────────────────────────────────────────────────────────

export interface StateVariable {
  name: string;
  /** Best-effort guess at the initial value (string representation) */
  initialValue: string;
}

export interface ImportEntry {
  source: string;
  specifiers: string[];
}

export interface ComponentSkeleton {
  name: string;
  filePath: string;
  type: 'functional' | 'class' | 'unknown';
  imports: ImportEntry[];
  stateVariables: StateVariable[];
  propNames: string[];
  hooks: string[];
  /** Top-level JSX tag names used (e.g. div, button, Link) */
  jsxElements: string[];
  hasRouter: boolean;
  hasContext: boolean;
  /** Raw URL strings found in fetch() / axios calls */
  apiEndpoints: string[];
  exports: string[];
  /** Rough token estimate for chunking: content.length / 4 */
  tokenEstimate: number;
}

export interface AppSkeleton {
  appName: string;
  totalFiles: number;
  entryPoint: string | null;
  components: ComponentSkeleton[];
  routes: string[];
  globalStateHints: string[];
  cssClasses: string[];
}

// ─── Main Public API ──────────────────────────────────────────────────────────

/**
 * Parses all JS/TS files in `files` and assembles a complete AppSkeleton.
 * CSS/SCSS files are scanned separately for class names.
 */
export function buildSkeleton(files: ExtractedFile[], appName: string): AppSkeleton {
  const components: ComponentSkeleton[] = [];
  const cssClasses = new Set<string>();
  let entryPoint: string | null = null;

  const jsFiles = files.filter(f =>
    ['.js', '.jsx', '.ts', '.tsx'].includes(f.extension)
  );

  logger.info(`AstParser: scanning ${jsFiles.length} JS/TS file(s)`);

  for (const file of jsFiles) {
    try {
      const skeleton = parseOneFile(file);
      if (!skeleton) continue;

      components.push(skeleton);

      // Detect entry point by conventional filename
      const base = path.basename(file.relativePath);
      if (
        ['index.js', 'index.jsx', 'index.ts', 'index.tsx',
         'App.js', 'App.jsx', 'App.ts', 'App.tsx',
         'main.js', 'main.jsx', 'main.ts', 'main.tsx'].includes(base)
      ) {
        entryPoint = file.relativePath;
      }
    } catch (err) {
      logger.warn(`AstParser: skipping "${file.relativePath}" — parse error: ${String(err)}`);
    }
  }

  // Scan CSS/SCSS for class names (useful for mapping to Flutter styles)
  for (const file of files.filter(f => ['.css', '.scss', '.sass'].includes(f.extension))) {
    for (const match of file.content.matchAll(/\.([a-zA-Z][\w-]+)\s*[{,]/g)) {
      cssClasses.add(match[1]);
    }
  }

  // Derive global state hints from import patterns
  const globalStateHints = new Set<string>();
  for (const comp of components) {
    for (const imp of comp.imports) {
      if (imp.source.includes('redux') || imp.source.includes('/store')) {
        globalStateHints.add('redux');
      }
      if (imp.source.includes('zustand'))   globalStateHints.add('zustand');
      if (imp.source.includes('recoil'))    globalStateHints.add('recoil');
      if (imp.source.includes('jotai'))     globalStateHints.add('jotai');
      if (imp.source.includes('mobx'))      globalStateHints.add('mobx');
    }
    if (comp.hasContext) globalStateHints.add('react-context');
  }

  // Collect all route paths found in react-router-aware components
  const routes = [
    ...new Set(
      components
        .filter(c => c.hasRouter)
        .flatMap(c => c.apiEndpoints)
    ),
  ];

  return {
    appName,
    totalFiles: files.length,
    entryPoint,
    components,
    routes,
    globalStateHints: [...globalStateHints],
    cssClasses: [...cssClasses].slice(0, 150), // cap to keep Skeleton.json small
  };
}

// ─── Single-File Parser ───────────────────────────────────────────────────────

function parseOneFile(file: ExtractedFile): ComponentSkeleton | null {
  const skeleton: ComponentSkeleton = {
    name: deriveComponentName(file.relativePath),
    filePath: file.relativePath,
    type: 'unknown',
    imports: [],
    stateVariables: [],
    propNames: [],
    hooks: [],
    jsxElements: [],
    hasRouter: false,
    hasContext: false,
    apiEndpoints: [],
    exports: [],
    tokenEstimate: Math.ceil(file.content.length / 4),
  };

  const ast = tryParse(file.content);
  if (!ast) return null;

  // Track seen hook names to avoid duplicates
  const seenHooks = new Set<string>();
  const seenJsx = new Set<string>();

  traverse(ast, {
    // ── Import declarations ───────────────────────────────────────────────
    ImportDeclaration(p) {
      const source = p.node.source.value;
      const specifiers = p.node.specifiers.map(s => {
        if (t.isImportDefaultSpecifier(s)) return `default:${s.local.name}`;
        if (t.isImportNamespaceSpecifier(s)) return `*:${s.local.name}`;
        if (t.isImportSpecifier(s)) {
          return t.isIdentifier(s.imported) ? s.imported.name : s.imported.value;
        }
        return '';
      }).filter(Boolean);

      skeleton.imports.push({ source, specifiers });

      if (source.includes('react-router')) skeleton.hasRouter = true;
      if (specifiers.some(s => s.toLowerCase().includes('context') || s === 'useContext')) {
        skeleton.hasContext = true;
      }
    },

    // ── Functional component declarations (function Foo() {}) ─────────────
    FunctionDeclaration(p) {
      const name = p.node.id?.name ?? '';
      if (isPascalCase(name)) {
        skeleton.type = 'functional';
        if (skeleton.name === 'Unknown') skeleton.name = name;
      }
    },

    // ── Arrow function components: const Foo = () => ... ──────────────────
    VariableDeclarator(p) {
      if (
        t.isIdentifier(p.node.id) &&
        isPascalCase(p.node.id.name) &&
        (t.isArrowFunctionExpression(p.node.init) || t.isFunctionExpression(p.node.init))
      ) {
        skeleton.type = 'functional';
        if (skeleton.name === 'Unknown') skeleton.name = p.node.id.name;
      }
    },

    // ── Class components: class Foo extends (React.)Component ─────────────
    ClassDeclaration(p) {
      const sup = p.node.superClass;
      const isReactComponent =
        (t.isIdentifier(sup) && sup.name === 'Component') ||
        (t.isMemberExpression(sup) &&
          t.isIdentifier(sup.object) && sup.object.name === 'React' &&
          t.isIdentifier(sup.property) && sup.property.name === 'Component');

      if (isReactComponent) {
        skeleton.type = 'class';
        if (p.node.id?.name) skeleton.name = p.node.id.name;
      }
    },

    // ── Hook calls & API calls ────────────────────────────────────────────
    CallExpression(p) {
      const callee = p.node.callee;

      // React hooks
      if (t.isIdentifier(callee)) {
        const hookNames = [
          'useState', 'useEffect', 'useCallback', 'useMemo',
          'useRef', 'useReducer', 'useContext', 'useLayoutEffect',
          'useNavigate', 'useParams', 'useLocation', 'useSearchParams',
        ];

        if (hookNames.includes(callee.name) && !seenHooks.has(callee.name)) {
          seenHooks.add(callee.name);
          skeleton.hooks.push(callee.name);
        }

        // useState — capture state var name + initial value
        if (callee.name === 'useState') {
          const parent = p.parent;
          if (
            t.isVariableDeclarator(parent) &&
            t.isArrayPattern(parent.id) &&
            parent.id.elements[0] &&
            t.isIdentifier(parent.id.elements[0])
          ) {
            const initArg = p.node.arguments[0];
            skeleton.stateVariables.push({
              name: parent.id.elements[0].name,
              initialValue: (initArg && !t.isArgumentPlaceholder(initArg))
                ? nodeToValueString(initArg)
                : 'null',
            });
          }
        }

        // fetch() calls
        if (callee.name === 'fetch') {
          const firstArg = p.node.arguments[0];
          if (firstArg && !t.isArgumentPlaceholder(firstArg)) {
            skeleton.apiEndpoints.push(extractEndpoint(firstArg));
          }
        }
      }

      // axios.get/post/put/delete/patch calls
      if (
        t.isMemberExpression(callee) &&
        t.isIdentifier(callee.object) &&
        callee.object.name === 'axios' &&
        t.isIdentifier(callee.property)
      ) {
        const firstArg = p.node.arguments[0];
        if (firstArg && !t.isArgumentPlaceholder(firstArg)) {
          skeleton.apiEndpoints.push(extractEndpoint(firstArg));
        }
      }
    },

    // ── JSX opening elements ──────────────────────────────────────────────
    JSXOpeningElement(p) {
      if (t.isJSXIdentifier(p.node.name)) {
        const tag = p.node.name.name;
        if (!seenJsx.has(tag)) {
          seenJsx.add(tag);
          skeleton.jsxElements.push(tag);
        }
      }
    },

    // ── PropTypes.shape or function parameters → extract prop names ────────
    AssignmentExpression(p) {
      // Foo.propTypes = { ... }
      if (
        t.isMemberExpression(p.node.left) &&
        t.isIdentifier(p.node.left.property) &&
        p.node.left.property.name === 'propTypes' &&
        t.isObjectExpression(p.node.right)
      ) {
        for (const prop of p.node.right.properties) {
          if (t.isObjectProperty(prop) && t.isIdentifier(prop.key)) {
            skeleton.propNames.push(prop.key.name);
          }
        }
      }
    },

    // ── Exports ───────────────────────────────────────────────────────────
    ExportDefaultDeclaration(p) {
      const decl = p.node.declaration;
      if (t.isIdentifier(decl)) skeleton.exports.push(`default:${decl.name}`);
      else skeleton.exports.push('default:anonymous');
    },
    ExportNamedDeclaration(p) {
      for (const spec of p.node.specifiers) {
        if (t.isExportSpecifier(spec) && t.isIdentifier(spec.exported)) {
          skeleton.exports.push(spec.exported.name);
        }
      }
    },
  });

  return skeleton;
}

// ─── File Chunker ─────────────────────────────────────────────────────────────

/**
 * Groups extracted JS/TS files into chunks sized below `maxTokensPerChunk`.
 * This keeps individual LLM requests within context-window limits.
 *
 * Rough heuristic: 1 token ≈ 4 characters (conservative for code).
 * Default: 8 000 tokens ≈ ~32 KB of source code per chunk.
 */
export function chunkFiles(
  files: ExtractedFile[],
  maxTokensPerChunk = 8_000
): Array<ExtractedFile[]> {
  const chunks: Array<ExtractedFile[]> = [];
  let current: ExtractedFile[] = [];
  let currentTokens = 0;

  for (const file of files) {
    const tokens = Math.ceil(file.content.length / 4);

    if (currentTokens + tokens > maxTokensPerChunk && current.length > 0) {
      chunks.push(current);
      current = [];
      currentTokens = 0;
    }

    // Very large single file — put it alone in its own chunk
    if (tokens > maxTokensPerChunk) {
      if (current.length > 0) { chunks.push(current); current = []; currentTokens = 0; }
      chunks.push([file]);
      continue;
    }

    current.push(file);
    currentTokens += tokens;
  }

  if (current.length > 0) chunks.push(current);

  logger.info(
    `AstParser: ${files.length} file(s) → ${chunks.length} chunk(s) ` +
    `(max ${maxTokensPerChunk} tokens/chunk)`
  );
  return chunks;
}

// ─── Private Helpers ──────────────────────────────────────────────────────────

/** Tries multiple Babel plugin combinations; returns null on complete failure. */
function tryParse(code: string): ReturnType<typeof babelParser.parse> | null {
  const commonPlugins: babelParser.ParserPlugin[] = [
    'jsx',
    'classProperties',
    'classStaticBlock',
    'optionalChaining',
    'nullishCoalescingOperator',
    'decorators-legacy',
    'exportDefaultFrom',
  ];

  // Attempt 1: full TypeScript + JSX
  try {
    return babelParser.parse(code, {
      sourceType: 'module',
      plugins: [...commonPlugins, 'typescript'],
      errorRecovery: true,
    });
  } catch { /* fall through */ }

  // Attempt 2: plain JS + JSX (handles .js files with JSX syntax)
  try {
    return babelParser.parse(code, {
      sourceType: 'module',
      plugins: commonPlugins,
      errorRecovery: true,
    });
  } catch { /* fall through */ }

  // Attempt 3: script mode (e.g. CommonJS require() files)
  try {
    return babelParser.parse(code, {
      sourceType: 'script',
      plugins: ['jsx'],
      errorRecovery: true,
    });
  } catch { return null; }
}

function isPascalCase(name: string): boolean {
  return /^[A-Z][a-zA-Z0-9]*$/.test(name);
}

function deriveComponentName(filePath: string): string {
  const base = path.basename(filePath, path.extname(filePath));
  return base.charAt(0).toUpperCase() + base.slice(1);
}

function extractEndpoint(node: t.Expression | t.SpreadElement | t.JSXNamespacedName | undefined): string {
  if (!node) return '<unknown>';
  if (t.isStringLiteral(node)) return node.value;
  if (t.isTemplateLiteral(node)) return '<dynamic-url>';
  return '<expression-url>';
}

function nodeToValueString(
  node: t.Expression | t.SpreadElement | t.JSXNamespacedName | undefined
): string {
  if (!node) return 'null';
  if (t.isStringLiteral(node)) return `"${node.value}"`;
  if (t.isNumericLiteral(node)) return String(node.value);
  if (t.isBooleanLiteral(node)) return String(node.value);
  if (t.isNullLiteral(node)) return 'null';
  if (t.isArrayExpression(node)) return '[]';
  if (t.isObjectExpression(node)) return '{}';
  if (t.isIdentifier(node) && node.name === 'undefined') return 'null';
  return 'unknown';
}