import { executeNodeAsync, parseTarget } from "./util.js";
import { resetCompilationCache } from "./commands.js";
import { parseAgency } from "../parser.js";
import { exprParser } from "../parsers/parsers.js";
import { TypescriptPreprocessor } from "../preprocessors/typescriptPreprocessor.js";
import { buildCompilationUnit } from "../compilationUnit.js";
import { AgencyGenerator } from "../backends/agencyGenerator.js";
import { expressionToString } from "../utils/node.js";
import { DefaultOptimizerIO } from "./optimizerIO.js";
import fs from "fs";
const DEFAULT_OPTIONS = {
    iterations: 5,
    earlyStopThreshold: 0.02,
    earlyStopPatience: 2,
};
export function getPromptValue(target) {
    if (target.llmCall && target.llmCall.type === "functionCall" && target.llmCall.arguments[0]?.type === "string") {
        return target.llmCall.arguments[0].segments
            .map((s) => s.type === "text" ? s.value : `\${${expressionToString(s.expression)}}`)
            .join("");
    }
    return "";
}
export function updatePrompt(target, newPrompt) {
    if (target.llmCall && target.llmCall.type === "functionCall" && target.llmCall.arguments[0]?.type === "string") {
        target.llmCall.arguments[0].segments = parsePromptToSegments(newPrompt);
    }
}
export function writeBack(filename, program) {
    const generator = new AgencyGenerator();
    const result = generator.generate(program);
    fs.writeFileSync(filename, result.output);
}
export function parsePromptToSegments(prompt) {
    const segments = [];
    const regex = /\$\{([^}]+)\}/g;
    let lastIndex = 0;
    let match;
    while ((match = regex.exec(prompt)) !== null) {
        if (match.index > lastIndex) {
            segments.push({ type: "text", value: prompt.slice(lastIndex, match.index) });
        }
        const exprStr = match[1].trim();
        const parsed = exprParser(exprStr);
        if (!parsed.success) {
            throw new Error(`Failed to parse interpolation expression: ${exprStr}`);
        }
        segments.push({ type: "interpolation", expression: parsed.result });
        lastIndex = match.index + match[0].length;
    }
    if (lastIndex < prompt.length) {
        segments.push({ type: "text", value: prompt.slice(lastIndex) });
    }
    return segments;
}
export async function optimize(config, target, opts = {}, io) {
    const options = { ...DEFAULT_OPTIONS, ...opts };
    const _io = io || new DefaultOptimizerIO(config);
    const { filename, nodeName } = parseTarget(target);
    // 1. Parse the file and run preprocessor to attach tags to nodes
    const contents = fs.readFileSync(filename, "utf-8");
    const parsed = parseAgency(contents);
    if (!parsed.success) {
        console.error("Parse error:", parsed.message);
        process.exit(1);
    }
    const info = buildCompilationUnit(parsed.result);
    const preprocessor = new TypescriptPreprocessor(parsed.result, config, info);
    const program = preprocessor.preprocess();
    // 2. Find @goal and @optimize tags
    const goalTag = findGoalTag(program, nodeName);
    const optimizeTargets = findOptimizeTargets(program, nodeName);
    if (!goalTag) {
        console.error(`No @goal tag found on node "${nodeName}". Add @goal("description") to define what success looks like.`);
        process.exit(1);
    }
    if (optimizeTargets.length === 0) {
        console.error(`No @optimize tags found in node "${nodeName}". Mark statements with @optimize to tell the optimizer what to tune.`);
        process.exit(1);
    }
    if (optimizeTargets.length > 1) {
        console.error(`Multiple @optimize targets found in node "${nodeName}". Only one @optimize target is supported in v1.`);
        process.exit(1);
    }
    const optimizeTarget = optimizeTargets[0];
    const nonPromptKeys = (optimizeTarget.configKeys || []).filter(k => k !== "prompt");
    if (nonPromptKeys.length > 0) {
        console.warn(`Warning: @optimize(${nonPromptKeys.join(", ")}) — only prompt optimization is supported in v1. Parameter tuning for ${nonPromptKeys.join(", ")} will be ignored.`);
    }
    const goalArg = goalTag.arguments[0];
    if (!goalArg ||
        goalArg.type !== "string" ||
        goalArg.segments.length !== 1 ||
        goalArg.segments[0].type !== "text") {
        const where = goalArg?.loc
            ? ` (line ${goalArg.loc.line}, col ${goalArg.loc.col})`
            : "";
        throw new Error(`@goal(...) requires a single plain string-literal argument${where}; got ${goalArg ? goalArg.type : "no argument"}.`);
    }
    const goal = goalArg.segments[0].value;
    console.log(`Goal: ${goal}`);
    console.log(`Running up to ${options.iterations} iterations...\n`);
    const targetNode = program.nodes.find((n) => n.type === "graphNode" && n.nodeName === nodeName);
    const history = [];
    let bestScore = -Infinity;
    let stagnantIterations = 0;
    for (let iteration = 1; iteration <= options.iterations; iteration++) {
        console.log(`\n--- Iteration ${iteration}/${options.iterations} ---\n`);
        const input = await _io.getUserInput(nodeName, targetNode.parameters);
        const argsString = targetNode.parameters
            .map((p) => {
            const val = input[p.name];
            return typeof val === "string" ? JSON.stringify(val) : String(val);
        })
            .join(", ");
        resetCompilationCache();
        console.log("Running agent...");
        let result;
        try {
            result = await executeNodeAsync({
                config,
                agencyFile: filename,
                nodeName,
                hasArgs: targetNode.parameters.length > 0,
                argsString,
            });
        }
        catch (e) {
            console.error("Execution error:", e);
            continue;
        }
        console.log("\nOutput:");
        console.log(JSON.stringify(result.data, null, 2));
        const { score, feedback } = await _io.collectFeedback();
        if (feedback === "done") {
            console.log("\nStopping optimization.");
            break;
        }
        history.push({ input, output: result.data, score, feedback, promptUsed: getPromptValue(optimizeTarget) });
        if (score !== null) {
            const improvement = bestScore > -Infinity ? (score - bestScore) / Math.abs(bestScore) : 1;
            if (score > bestScore)
                bestScore = score;
            if (improvement < options.earlyStopThreshold) {
                stagnantIterations++;
            }
            else {
                stagnantIterations = 0;
            }
            if (stagnantIterations >= options.earlyStopPatience) {
                console.log(`\nScore hasn't improved significantly for ${options.earlyStopPatience} iterations. Stopping.`);
                break;
            }
        }
        if (iteration < options.iterations) {
            console.log("\nProposing improved prompt...");
            const proposed = await _io.proposeImprovement(getPromptValue(optimizeTarget), goal, history);
            if (await _io.confirmProposal(proposed)) {
                updatePrompt(optimizeTarget, proposed);
                writeBack(filename, program);
                console.log(`Updated ${filename}`);
            }
        }
    }
    // Final summary
    if (history.length > 0) {
        console.log("\n--- Optimization Summary ---");
        console.log(`Iterations: ${history.length}`);
        if (bestScore > -Infinity)
            console.log(`Best score: ${bestScore}/10`);
        console.log(`Final prompt: "${getPromptValue(optimizeTarget)}"`);
    }
}
export function findGoalTag(program, nodeName) {
    for (const node of program.nodes) {
        if (node.type === "graphNode" && node.nodeName === nodeName) {
            const tags = node.tags || [];
            return tags.find((t) => t.name === "goal") || null;
        }
    }
    return null;
}
export function findOptimizeTargets(program, nodeName) {
    const targets = [];
    for (const node of program.nodes) {
        if (node.type === "graphNode" && node.nodeName === nodeName) {
            collectOptimizeTargets(node.body, targets);
        }
    }
    return targets;
}
function collectOptimizeTargets(body, targets) {
    for (const node of body) {
        if (node.type !== "assignment" && node.type !== "functionCall")
            continue;
        const tags = node.tags || [];
        const optimizeTag = tags.find((t) => t.name === "optimize");
        if (optimizeTag) {
            let llmCall = null;
            if (node.type === "assignment" &&
                node.value?.type === "functionCall" &&
                node.value.functionName === "llm") {
                llmCall = node.value;
            }
            else if (node.type === "functionCall" && node.functionName === "llm") {
                llmCall = node;
            }
            const target = { node, llmCall, tag: optimizeTag };
            target.promptValue = getPromptValue(target);
            if (optimizeTag.arguments.length === 0) {
                target.configKeys = ["prompt"];
            }
            else {
                target.configKeys = optimizeTag.arguments.map((a) => {
                    if (a.type !== "variableName") {
                        const where = a.loc ? ` (line ${a.loc.line}, col ${a.loc.col})` : "";
                        throw new Error(`@optimize(...) arguments must be plain config-key identifiers${where}; got ${a.type}.`);
                    }
                    return a.value;
                });
            }
            targets.push(target);
        }
    }
}
