import fs from 'node:fs';
import path from 'node:path';
import { fileURLToPath } from 'node:url';
import { parse } from 'acorn';

const scriptDir = path.dirname(fileURLToPath(import.meta.url));
const rootDir = path.resolve(scriptDir, '..');
const outputDir = path.join(rootDir, 'output');
const excludedDirs = new Set(['.git', 'data', 'node_modules', 'output']);
const functions = [];

function collectJavaScriptFiles(dir) {
    const files = [];
    for (const entry of fs.readdirSync(dir, { withFileTypes: true })) {
        if (entry.isDirectory() && excludedDirs.has(entry.name)) continue;

        const fullPath = path.join(dir, entry.name);
        if (entry.isDirectory()) {
            files.push(...collectJavaScriptFiles(fullPath));
        } else if (entry.isFile() && entry.name.endsWith('.js')) {
            files.push(fullPath);
        }
    }
    return files;
}

function inferredName(node, parent) {
    if (node.id?.name) return node.id.name;
    if (parent?.type === 'VariableDeclarator' && parent.id.type === 'Identifier') {
        return parent.id.name;
    }
    if (parent?.type === 'AssignmentExpression') {
        return parent.left.name ?? parent.left.property?.name ?? 'assigned';
    }
    if (parent?.type === 'Property' || parent?.type === 'MethodDefinition') {
        return parent.key.name ?? parent.key.value ?? 'method';
    }
    if (parent?.type === 'ExportDefaultDeclaration') return 'default';
    return 'anonymous';
}

function visit(node, parent, source, relativePath) {
    if (!node || typeof node !== 'object') return;

    if (node.type === 'FunctionDeclaration'
        || node.type === 'FunctionExpression'
        || node.type === 'ArrowFunctionExpression') {
        const method = (parent?.type === 'MethodDefinition'
            || (parent?.type === 'Property' && parent.method))
            ? parent
            : node;
        const start = method.start;
        const end = method.end;
        const startLine = source.slice(0, start).split('\n').length;
        const endLine = source.slice(0, end).split('\n').length;
        const name = inferredName(node, parent);
        let signature = source.slice(start, node.body.start).replace(/\s+/g, ' ').trim();
        if (method === node
            && name !== 'anonymous'
            && node.type !== 'FunctionDeclaration'
            && parent?.type !== 'ExportDefaultDeclaration') {
            signature = `${name} = ${signature}`;
        }
        functions.push({
            source: relativePath,
            start,
            end,
            startLine,
            endLine,
            name,
            signature,
            code: source.slice(start, end),
        });
    }

    for (const [key, value] of Object.entries(node)) {
        if (key === 'start' || key === 'end' || key === 'loc') continue;
        if (Array.isArray(value)) {
            for (const child of value) visit(child, node, source, relativePath);
        } else if (value && typeof value === 'object') {
            visit(value, node, source, relativePath);
        }
    }
}

const sourceFiles = collectJavaScriptFiles(rootDir).sort();
for (const filePath of sourceFiles) {
    const source = fs.readFileSync(filePath, 'utf8');
    const relativePath = path.relative(rootDir, filePath).split(path.sep).join('/');
    let ast;
    try {
        ast = parse(source, {
            ecmaVersion: 'latest',
            sourceType: 'module',
            allowHashBang: true,
        });
    } catch (error) {
        throw new Error(`Could not parse ${relativePath}: ${error.message}`, { cause: error });
    }
    visit(ast, null, source, relativePath);
}

fs.mkdirSync(outputDir, { recursive: true });
for (const entry of fs.readdirSync(outputDir, { withFileTypes: true })) {
    if (entry.isFile() && /^function-.*\.js$/.test(entry.name)) {
        fs.unlinkSync(path.join(outputDir, entry.name));
    }
}

functions.sort((a, b) => a.source.localeCompare(b.source) || a.start - b.start);
const nameCounts = new Map();
for (const fn of functions) {
    const sourceSlug = fn.source.replace(/\.js$/, '').replace(/[^a-zA-Z0-9_-]+/g, '-');
    const nameSlug = fn.name.replace(/[^a-zA-Z0-9_-]+/g, '-');
    const baseName = `function-${sourceSlug}-L${fn.startLine}-${nameSlug}`;
    const sequence = (nameCounts.get(baseName) ?? 0) + 1;
    nameCounts.set(baseName, sequence);

    const outputName = `${baseName}-${sequence}.js`;
    const contents = [
        `// Source: ${fn.source} (lines ${fn.startLine}-${fn.endLine})`,
        fn.code,
        '',
    ].join('\n');
    fs.writeFileSync(path.join(outputDir, outputName), contents);
}

const index = functions
    .map(fn => `${fn.source}:${fn.startLine}\t${fn.signature}`)
    .join('\n');
fs.writeFileSync(path.join(outputDir, 'functions.txt'), `${index}\n`);

console.log(`Extracted ${functions.length} functions from ${sourceFiles.length} JavaScript files into ${path.relative(rootDir, outputDir)}/`);