import { transformWithEsbuild } from 'vite';
import MagicString from 'magic-string';
import { parse } from '@babel/parser';
import { walk } from 'estree-walker';
import crypto from 'node:crypto';
import { isSelectish, splitSqlStatements } from './lib/sql.js';
import { extractStaticString, getObjectPropValue } from './lib/ast.js';
import { emitS2QueryReplacementBuild, emitS2QueryReplacementDev } from './lib/codegen.js';
const defaultOptions = {
    endpoint: '/api/query',
    debounceMs: 300,
    debug: false,
    devMode: 'soft'
};
function parseAesKey(secret) {
    const s = String(secret).trim();
    if (/^[0-9a-f]{64}$/i.test(s))
        return Buffer.from(s, 'hex');
    const b64 = s.replace(/-/g, '+').replace(/_/g, '/');
    return Buffer.from(b64, 'base64');
}
function aesGcmEncryptToToken(key, payloadObj) {
    const V = 1;
    const iv = crypto.randomBytes(12);
    const cipher = crypto.createCipheriv('aes-256-gcm', key, iv);
    const plaintext = Buffer.from(JSON.stringify(payloadObj), 'utf8');
    const ct = Buffer.concat([cipher.update(plaintext), cipher.final()]);
    const tag = cipher.getAuthTag();
    const header = Buffer.allocUnsafe(1);
    header.writeUInt8(V, 0);
    const packed = Buffer.concat([header, iv, tag, ct]);
    return packed.toString('base64url');
}
function ensureNamedImportsFrom(code, s, source, names) {
    const ast = parse(code, { sourceType: 'module', plugins: ['typescript', 'jsx'] });
    let imp;
    for (const n of ast.program.body ?? []) {
        if (n.type === 'ImportDeclaration' && n.source?.value === source) {
            imp = n;
            break;
        }
    }
    if (!imp) {
        s.prepend(`import { ${names.join(', ')} } from '${source}'\n`);
        return;
    }
    const named = new Set();
    for (const spec of imp.specifiers ?? []) {
        if (spec.type === 'ImportSpecifier')
            named.add(spec.imported.name);
    }
    const missing = names.filter((n) => !named.has(n));
    if (missing.length === 0)
        return;
    const src = code.slice(imp.start, imp.end);
    const braceClose = src.lastIndexOf('}');
    if (braceClose === -1) {
        s.prepend(`import { ${missing.join(', ')} } from '${source}'\n`);
        return;
    }
    const insertPos = imp.start + braceClose;
    s.appendLeft(insertPos, `, ${missing.join(', ')}`);
}
function unwrapLikeTs(node) {
    let n = node;
    while (n) {
        if (n.type === 'TSAsExpression' || n.type === 'TSSatisfiesExpression') {
            n = n.expression;
            continue;
        }
        if (n.type === 'TypeCastExpression') {
            n = n.expression;
            continue;
        }
        if (n.type === 'ParenthesizedExpression') {
            n = n.expression;
            continue;
        }
        break;
    }
    return n;
}
function extractSqlLike(node) {
    const n = unwrapLikeTs(node);
    if (!n)
        return null;
    // '...'
    const direct = extractStaticString(n);
    if (direct != null)
        return direct;
    // ref('...') / shallowRef('...')
    if (n.type === 'CallExpression') {
        const callee = unwrapLikeTs(n.callee);
        const calleeName = callee?.type === 'Identifier'
            ? callee.name
            : callee?.type === 'MemberExpression' &&
                !callee.computed &&
                callee.property?.type === 'Identifier'
                ? callee.property.name
                : null;
        if (calleeName === 'ref' || calleeName === 'shallowRef') {
            const a0 = n.arguments?.[0];
            if (!a0)
                return null;
            return extractStaticString(a0);
        }
    }
    return null;
}
function resolveSqlFromNode(node, lookupInit) {
    const seen = new Set();
    const go = (n0, depth) => {
        if (!n0)
            return null;
        if (depth > 25)
            return null;
        const n = unwrapLikeTs(n0);
        const direct = extractSqlLike(n);
        if (direct != null)
            return direct;
        if (n?.type === 'Identifier' && typeof n.name === 'string') {
            if (seen.has(n.name))
                return null;
            seen.add(n.name);
            const init = lookupInit(n.name);
            if (!init)
                return null;
            return go(init, depth + 1);
        }
        return null;
    };
    return go(node, 0);
}
export default function dbbuilderTransform(opts = {}) {
    const options = { ...defaultOptions, ...opts };
    let useSoftDev = false;
    let keyBuf = null;
    return {
        name: 'dbbuilder-transform',
        enforce: 'post',
        configResolved(config) {
            const effectiveDevMode = opts.devMode ?? (config.command === 'serve' ? 'soft' : 'hard');
            useSoftDev = effectiveDevMode === 'soft';
            if (!useSoftDev) {
                // Prefer S2_AES_KEY, fallback to DAO_KEY for backwards compat.
                const secret = opts.secretKey ?? process.env.S2_AES_KEY ?? process.env.DAO_KEY ?? '';
                if (!secret) {
                    throw new Error('[dbbuilder-transform] Missing secretKey (or S2_AES_KEY / DAO_KEY env var).');
                }
                const k = parseAesKey(secret);
                if (k.length !== 32) {
                    throw new Error(`[dbbuilder-transform] secretKey must be 32 bytes for AES-256-GCM (got ${k.length}).`);
                }
                keyBuf = k;
            }
        },
        async transform(code, id) {
            if (!id || id.includes('\0'))
                return null;
            if (id.includes('/node_modules/'))
                return null;
            const bareId = id.split('?', 1)[0];
            const isVueScriptModule = id.includes('.vue?') && (id.includes('type=script') || id.includes('vue&type=script'));
            const isVueMainModule = bareId.endsWith('.vue') && !id.includes('?');
            const isJsTsModule = /\.(ts|tsx|js|jsx)$/.test(bareId);
            if (!isVueMainModule && !isVueScriptModule && !isJsTsModule)
                return null;
            // cheap prefilter to avoid parsing work
            if (!code.includes('s2client') || !code.includes('.query'))
                return null;
            if (opts.debug)
                console.log('[dbbuilder-transform] hit', id);
            const loader = isVueMainModule
                ? 'js'
                : id.includes('lang.tsx') || bareId.endsWith('.tsx')
                    ? 'tsx'
                    : id.includes('lang.ts') || bareId.endsWith('.ts')
                        ? 'ts'
                        : bareId.endsWith('.jsx')
                            ? 'jsx'
                            : 'js';
            const esbuildResult = await transformWithEsbuild(code, id, { loader });
            code = esbuildResult.code;
            const ast = parse(code, { sourceType: 'module', plugins: ['typescript', 'jsx'] });
            // find local identifier(s) for `s2client`
            const s2clientLocals = new Set();
            for (const n of ast.program.body ?? []) {
                if (n?.type !== 'ImportDeclaration')
                    continue;
                if (n.source?.value !== '@compilers/s2client')
                    continue;
                for (const spec of n.specifiers ?? []) {
                    if (spec.type === 'ImportSpecifier' && spec.imported?.name === 's2client') {
                        s2clientLocals.add(spec.local?.name ?? 's2client');
                    }
                }
            }
            if (s2clientLocals.size === 0)
                s2clientLocals.add('s2client');
            const s = new MagicString(code);
            let didChange = false;
            // parent tracking without mutating nodes
            const parentMap = new WeakMap();
            // simple lexical scope init lookup (enough for `const sql = '...'` cases)
            const scopes = [new Map()];
            const pushScope = () => scopes.push(new Map());
            const popScope = () => {
                if (scopes.length > 1)
                    scopes.pop();
            };
            const lookupInit = (name) => {
                for (let i = scopes.length - 1; i >= 0; i--) {
                    if (scopes[i].has(name))
                        return scopes[i].get(name) ?? null;
                }
                return null;
            };
            const isS2QueryCallee = (callee) => {
                const c = unwrapLikeTs(callee);
                return (c &&
                    c.type === 'MemberExpression' &&
                    !c.computed &&
                    c.object?.type === 'Identifier' &&
                    s2clientLocals.has(c.object.name) &&
                    c.property?.type === 'Identifier' &&
                    c.property.name === 'query');
            };
            const isExecuteChainForQueryCall = (queryCallNode) => {
                const p = unwrapLikeTs(parentMap.get(queryCallNode));
                if (!p ||
                    p.type !== 'MemberExpression' ||
                    p.computed ||
                    p.object !== queryCallNode ||
                    p.property?.type !== 'Identifier' ||
                    p.property.name !== 'execute') {
                    return false;
                }
                const gp = unwrapLikeTs(parentMap.get(p));
                if (!gp || gp.type !== 'CallExpression' || gp.callee !== p)
                    return false;
                return (gp.arguments ?? []).length === 0;
            };
            walk(ast, {
                enter(node, parent) {
                    parentMap.set(node, parent ?? null);
                    // scope tracking
                    if (node?.type === 'Program' ||
                        node?.type === 'BlockStatement' ||
                        node?.type === 'FunctionDeclaration' ||
                        node?.type === 'FunctionExpression' ||
                        node?.type === 'ArrowFunctionExpression') {
                        pushScope();
                    }
                    if (node?.type === 'VariableDeclarator' && node.id?.type === 'Identifier') {
                        scopes[scopes.length - 1].set(node.id.name, node.init ?? null);
                    }
                    if (node?.type !== 'CallExpression')
                        return;
                    let queryCall = null;
                    let rewriteStart = node.start;
                    let rewriteEnd = node.end;
                    let whenArgNode = null;
                    // NEW: track whether the matched expression included `.execute()`
                    let matchedExecute = false;
                    const callee = unwrapLikeTs(node.callee);
                    const isMemberCallNamed = (callNode, name) => {
                        const c = unwrapLikeTs(callNode?.callee);
                        return (c?.type === 'MemberExpression' &&
                            !c.computed &&
                            c.property?.type === 'Identifier' &&
                            c.property.name === name);
                    };
                    const getMemberObject = (callNode) => {
                        const c = unwrapLikeTs(callNode?.callee);
                        return c?.type === 'MemberExpression' ? unwrapLikeTs(c.object) : null;
                    };
                    const isChainedFromQueryCall = (queryNode) => {
                        const p = unwrapLikeTs(parentMap.get(queryNode));
                        if (!p || p.type !== 'MemberExpression' || p.object !== queryNode)
                            return false;
                        const prop = p.property;
                        if (prop?.type !== 'Identifier')
                            return false;
                        if (prop.name !== 'execute' && prop.name !== 'when')
                            return false;
                        const gp = unwrapLikeTs(parentMap.get(p));
                        return !!(gp && gp.type === 'CallExpression' && gp.callee === p);
                    };
                    // Case 1: s2client.query(...).when(pred).execute()
                    if (isMemberCallNamed(node, 'execute') && (node.arguments ?? []).length === 0) {
                        const obj1 = getMemberObject(node);
                        if (obj1?.type === 'CallExpression' && isMemberCallNamed(obj1, 'when')) {
                            const obj0 = getMemberObject(obj1);
                            if (obj0?.type === 'CallExpression' && isS2QueryCallee(obj0.callee)) {
                                queryCall = obj0;
                                whenArgNode = (obj1.arguments ?? [])[0] ?? null;
                                rewriteStart = node.start;
                                rewriteEnd = node.end;
                                matchedExecute = true; // NEW
                            }
                        }
                        else if (obj1?.type === 'CallExpression' && isS2QueryCallee(obj1.callee)) {
                            // Case 2: s2client.query(...).execute()
                            queryCall = obj1;
                            rewriteStart = node.start;
                            rewriteEnd = node.end;
                            matchedExecute = true; // NEW
                        }
                    }
                    // Case 3: s2client.query(...).when(pred)   (no execute)
                    if (!queryCall && isMemberCallNamed(node, 'when')) {
                        const obj0 = getMemberObject(node);
                        if (obj0?.type === 'CallExpression' && isS2QueryCallee(obj0.callee)) {
                            queryCall = obj0;
                            whenArgNode = (node.arguments ?? [])[0] ?? null;
                            rewriteStart = node.start;
                            rewriteEnd = node.end;
                            matchedExecute = false; // NEW (explicit)
                        }
                    }
                    // Case 4: direct s2client.query(...)
                    if (!queryCall && isS2QueryCallee(callee)) {
                        // if chained into .when or .execute, skip to avoid overlapping edits
                        if (isChainedFromQueryCall(node))
                            return;
                        queryCall = node;
                        rewriteStart = node.start;
                        rewriteEnd = node.end;
                        matchedExecute = false; // NEW
                    }
                    if (!queryCall)
                        return;
                    const args = queryCall.arguments ?? [];
                    const queryArg = args[0];
                    const paramsArg = args[1];
                    const optionsArg = args[2];
                    if (!queryArg)
                        return;
                    const rawSql = resolveSqlFromNode(queryArg, lookupInit);
                    if (rawSql == null) {
                        if (opts.debug) {
                            console.warn('[dbbuilder-transform] unresolved SQL:', id, code.slice(queryArg.start, queryArg.end));
                        }
                        return;
                    }
                    const stmts = splitSqlStatements(rawSql);
                    if (stmts.length === 0)
                        return;
                    const paramsSrc = paramsArg ? code.slice(paramsArg.start, paramsArg.end) : '{}';
                    // NEW: runtime expression for guard (compiled from `.when(...)`)
                    // default = always run
                    const whenRuntimeExpr = whenArgNode ? code.slice(whenArgNode.start, whenArgNode.end) : 'true';
                    // options.url / options.headers / options.callback (all optional)
                    let urlExpr = null;
                    let headersRuntimeExpr = 'undefined';
                    let callbackRuntimeExpr = 'null';
                    if (optionsArg) {
                        const urlNode = getObjectPropValue(optionsArg, 'url');
                        const headersNode = getObjectPropValue(optionsArg, 'headers');
                        const callbackNode = getObjectPropValue(optionsArg, 'callback');
                        const urlStatic = urlNode ? extractStaticString(urlNode) : null;
                        if (urlStatic != null)
                            urlExpr = JSON.stringify(urlStatic);
                        if (headersNode)
                            headersRuntimeExpr = code.slice(headersNode.start, headersNode.end);
                        if (callbackNode)
                            callbackRuntimeExpr = code.slice(callbackNode.start, callbackNode.end);
                    }
                    const endpointExpr = urlExpr ?? JSON.stringify(options.endpoint);
                    const debounceMsExpr = JSON.stringify(options.debounceMs);
                    const replacementRaw = useSoftDev
                        ? emitS2QueryReplacementDev({
                            paramsSrc,
                            whenRuntimeExpr,
                            sqls: stmts,
                            isSelects: stmts.map(isSelectish),
                            endpointExpr,
                            debounceMsExpr,
                            headersRuntimeExpr,
                            callbackRuntimeExpr
                        })
                        : emitS2QueryReplacementBuild({
                            paramsSrc,
                            whenRuntimeExpr,
                            tokens: stmts.map((sql) => {
                                if (!keyBuf)
                                    throw new Error('[dbbuilder-transform] Internal error: missing keyBuf.');
                                return aesGcmEncryptToToken(keyBuf, { v: 1, sql, isSelect: isSelectish(sql) });
                            }),
                            isSelects: stmts.map(isSelectish),
                            endpointExpr,
                            debounceMsExpr,
                            headersRuntimeExpr,
                            callbackRuntimeExpr
                        });
                    // NEW: inject `guarded: true` into the runtime args (unless already present)
                    let replacement = replacementRaw;
                    if (!replacement.includes('guarded:')) {
                        // expected shape from codegen: s2queryRuntime({ ... })
                        replacement = replacement.replace('s2queryRuntime({', 's2queryRuntime({\n    guarded: true,');
                    }
                    // NEW: preserve `.execute()` semantics
                    if (matchedExecute && !replacement.trimEnd().endsWith('.execute()')) {
                        replacement = `${replacement}.execute()`;
                    }
                    s.overwrite(rewriteStart, rewriteEnd, replacement);
                    didChange = true;
                    if (opts.debug) {
                        console.log('[dbbuilder-transform] REWRITE', {
                            id,
                            mode: useSoftDev ? 'dev-soft' : 'build-hard',
                            stmts: stmts.length
                        });
                    }
                    // IMPORTANT: prevent visiting children of this rewritten chain
                    // (avoids rewriting `.when(...)` after rewriting `.when(...).execute()`)
                    this.skip();
                },
                leave(node) {
                    if (node?.type === 'Program' ||
                        node?.type === 'BlockStatement' ||
                        node?.type === 'FunctionDeclaration' ||
                        node?.type === 'FunctionExpression' ||
                        node?.type === 'ArrowFunctionExpression') {
                        popScope();
                    }
                }
            });
            if (!didChange)
                return null;
            ensureNamedImportsFrom(code, s, '@compilers/s2client/runtime', ['s2queryRuntime']);
            return { code: s.toString(), map: s.generateMap({ hires: true }) };
        }
    };
}
