KnowledgeBase/server/scripts/test-local-import.ts

import axios from 'axios';
import * as fs from 'fs';
import * as path from 'path';
import * as os from 'os';

async function testLocalImport() {
    const baseURL = 'http://localhost:3001/api';
    const username = process.argv[2] || 'admin';
    const password = process.argv[3];
    const sourceFolder = process.argv[4];
    const tenantId = process.argv[5];

    if (!password) {
        console.error('Usage: ts-node scripts/test-local-import.ts <username> <password> [sourceFolder] [tenantId]');
        process.exit(1);
    }

    try {
        // 1. Login to get JWT Token
        console.log(`Logging in as ${username} to ${baseURL}...`);
        const loginRes = await axios.post(`${baseURL}/auth/login`, {
            username,
            password
        });
        const jwtToken = loginRes.data.access_token;
        console.log('Login successful.');

        // 2. Get API Key using JWT Token
        console.log('Retrieving API key...');
        const apiKeyRes = await axios.get(`${baseURL}/auth/api-key`, {
            headers: { Authorization: `Bearer ${jwtToken}` }
        });
        const apiKey = apiKeyRes.data.apiKey;
        console.log('API Key retrieved:', apiKey);

        // From now on, using x-api-key for authentication
        const authHeaders: any = { 'x-api-key': apiKey };
        if (tenantId) {
            authHeaders['x-tenant-id'] = tenantId;
            console.log(`Target tenant set to: ${tenantId}`);
        }

        // 3. Prepare folder structure
        let targetPath = sourceFolder;
        let isTemp = false;

        if (!targetPath) {
            isTemp = true;
            targetPath = path.join(os.tmpdir(), `aurak-test-${Date.now()}`);
            const subDir = path.join(targetPath, 'subfolder');

            fs.mkdirSync(targetPath, { recursive: true });
            fs.mkdirSync(subDir, { recursive: true });

            fs.writeFileSync(path.join(targetPath, 'root-file.md'), '# Root File\nContent in root.', 'utf8');
            fs.writeFileSync(path.join(subDir, 'sub-file.txt'), 'Content in subfolder.', 'utf8');

            console.log(`Created temporary test structure at: ${targetPath}`);
        } else {
            console.log(`Using provided source folder: ${targetPath}`);
            if (!fs.existsSync(targetPath)) {
                throw new Error(`The specified folder does not exist: ${targetPath}`);
            }
        }

        // 4. Initial check for embedding models
        const modelsRes = await axios.get(`${baseURL}/models`, {
            headers: authHeaders
        });
        const embeddingModel = modelsRes.data.find((m: any) => m.type === 'embedding' && m.isEnabled !== false);

        if (!embeddingModel) {
            throw new Error('No enabled embedding model found');
        }

        console.log(`Using embedding model: ${embeddingModel.id}`);

        // 5. Call local-folder import endpoint
        console.log('Triggering local folder import...');
        const importRes = await axios.post(`${baseURL}/upload/local-folder`, {
            sourcePath: targetPath,
            embeddingModelId: embeddingModel.id,
            useHierarchy: true
        }, {
            headers: authHeaders
        });

        console.log('Import response:', importRes.data);

        // 6. Verification
        if (isTemp) {
            console.log('Waiting for background processing (10s)...');
            await new Promise(resolve => setTimeout(resolve, 10000));

            const kbRes = await axios.get(`${baseURL}/knowledge-bases`, {
                headers: authHeaders
            });

            const importedFiles = kbRes.data.filter((f: any) =>
                f.originalName === 'root-file.md' || f.originalName === 'sub-file.txt'
            );

            console.log(`Found ${importedFiles.length} imported files in KB.`);
            if (importedFiles.length === 2) {
                console.log('SUCCESS: All files imported.');
            } else {
                console.log('FAILURE: Not all files were imported.');
            }
        } else {
            console.log('Custom folder import triggered. Please check the UI or database for results.');
        }

    } catch (error: any) {
        if (error.response) {
            console.error(`Test failed with status ${error.response.status}:`, JSON.stringify(error.response.data));
        } else {
            console.error('Test failed:', error.message);
        }
        process.exit(1);
    }
}


testLocalImport();