agent-v1

import { tool, type FileHandle } from "@lmstudio/sdk";
import { z } from "zod";
import { Database } from "duckdb";
import { loadExcelIntoDuckDb, getSampleRows, sanitizeTableName } from "../../utils/excel/loadExcel";

const dbCache = new Map<string, Database>();

function bigIntSafe(_key: string, value: unknown): unknown {
    if (typeof value === "bigint") {
        return value <= BigInt(Number.MAX_SAFE_INTEGER) && value >= BigInt(Number.MIN_SAFE_INTEGER)
            ? Number(value)
            : value.toString();
    }
    return value;
}

function cacheKey(file: FileHandle): string {
    return file.name;
}

function withLogging<T extends (...args: any[]) => Promise<string>>(name: string, fn: T): T {
    return (async (...args: Parameters<T>) => {
        console.log(`[TOOL CALL] ${name}`, JSON.stringify(args[0], null, 2));
        const result = await fn(...args);
        console.log(`[TOOL RESULT] ${name}`, result.length > 2000 ? result.slice(0, 2000) + "...(truncated)" : result);
        return result;
    }) as T;
}

export async function getOrCreateDb(file: FileHandle): Promise<Database> {
    const key = cacheKey(file);
    const cached = dbCache.get(key);
    if (cached) return cached;

    const db = new Database(":memory:");
    await loadExcelIntoDuckDb(db, file);
    dbCache.set(key, db);
    return db;
}

export function queryAll(db: any, sql: string): Promise<Record<string, unknown>[]> {
    return new Promise((resolve, reject) => {
        const conn = db.connect();
        conn.all(sql, (err: any, rows: any) => {
            conn.close(() => {
                if (err) reject(err);
                else resolve(rows);
            });
        });
    });
}

// Builds the SUB-AGENT's tools (list_spreadsheet_tables, query_spreadsheet) —
// these are what excelSubAgent.ts calls internally, NOT what the main model sees.
export async function buildExcelTools(excelFiles: FileHandle[]) {
    if (excelFiles.length !== 1) {
        throw new Error("buildExcelTools currently expects exactly one target file.");
    }
    const file = excelFiles[0];
    const db = await getOrCreateDb(file);

    const dbs = new Map<string, Database>();
    for (const f of excelFiles) {
        dbs.set(cacheKey(f), await getOrCreateDb(f));
    }

    const listTables = tool({
        name: "list_spreadsheet_tables",
        description:
            "Discovers the exact sanitized internal table (sheet) names and column structures required " +
            "to construct valid SQL. Guessing table names will result in syntax errors.",
        parameters: {},
        implementation: async () => {
            const tables = await queryAll(
                db,
                `SELECT table_name FROM information_schema.tables WHERE table_schema='main'`
            );

            const out = [];
            for (const t of tables) {
                const tableName = t.table_name as string;
                const sampleRows = await getSampleRows(db, tableName, 3);
                const profiles = await queryAll(db, `PRAGMA table_info("${tableName}")`);
                const [{ total_rows }] = await queryAll(db, `SELECT COUNT(*) AS total_rows FROM "${tableName}"`);
                const [{ exact_dupes }] = await queryAll(
                    db,
                    `SELECT (SELECT COUNT(*) FROM "${tableName}") - (SELECT COUNT(*) FROM (SELECT DISTINCT * FROM "${tableName}")) AS exact_dupes`
                );

                const columnProfiles = [];
                for (const p of profiles) {
                    const colName = p.name as string;
                    const [{ distinct_count, nulls }] = await queryAll(db,
                        `SELECT COUNT(DISTINCT "${colName}") as distinct_count, COUNT(*) - COUNT("${colName}") as nulls FROM "${tableName}"`
                    );
                    columnProfiles.push({
                        name: colName,
                        type: p.type,
                        null_count: Number(nulls),
                        distinct_count: Number(distinct_count),
                        likely_categorical: Number(distinct_count) < Number(total_rows) * 0.5,
                    });
                }

                out.push({
                    table: tableName,
                    total_rows: Number(total_rows),
                    exact_duplicate_rows: Number(exact_dupes),
                    columns: columnProfiles,
                    sample_rows: sampleRows,
                });
            }
            return JSON.stringify(out, bigIntSafe, 2);
        },
    });


    const queryData = tool({
        name: "query_spreadsheet",
        description:
            `Runs a read-only SQL query against the spreadsheet's tables. ` +
            `If this tool returns a SQL error, the error response includes the ACTUAL column names — ` +
            `read them and retry immediately with corrected names rather than asking the user.`,
        parameters: {
            sql: z.string(),
        },
        implementation: withLogging("query_spreadsheet", async ({ sql }) => {
            const clean = sql.trim().toLowerCase();
            if (!clean.startsWith("select") && !clean.startsWith("with")) {
                return "Error: only SELECT or WITH queries are permitted.";
            }
            try {
                const rows = await queryAll(db, sql);
                return JSON.stringify(rows.slice(0, 200), bigIntSafe, 2);
            } catch (e) {
                const message = (e as Error).message;
                const isSchemaError = /column|table|does not exist|not found|binder error/i.test(message);
                if (isSchemaError) {
                    try {
                        const tableMatch = sql.match(/from\s+"?(\w+)"?/i);
                        const tableName = tableMatch?.[1];
                        if (tableName) {
                            const cols = await queryAll(db, `PRAGMA table_info("${tableName}")`);
                            const colList = cols.map(c => `${c.name} (${c.type})`).join(", ");
                            return `SQL error: ${message}\n\n` +
                                `Actual columns in "${tableName}": ${colList}\n\n` +
                                `Retry the query using these exact column names.`;
                        }
                    } catch {
                        // fall through
                    }
                }
                return `SQL error: ${message}`;
            }
        }),
    });


    return [listTables, queryData];
}