import { spawn } from "node:child_process";
import { randomBytes, randomUUID } from "node:crypto";
import { existsSync, rmSync } from "node:fs";
import { createServer } from "node:net";
import { join } from "node:path";
import { createInterface } from "node:readline";
import { createDefaultCoreDeps } from "../runtime/deps.js";
import { createDiffReviewErrorResponse, createDiffReviewSuccessResponse, DIFF_REVIEW_CLIENT_METHODS, DIFF_REVIEW_ERROR_CODES, DIFF_REVIEW_PROTOCOL_VERSION, parseDiffReviewMessageLine, serializeDiffReviewMessage, } from "./protocol.js";
import { DiffReviewThread, } from "./review_thread.js";
import { captureDiffReviewSnapshot, } from "./snapshot.js";
class DiffReviewRequestError extends Error {
    code;
    constructor(code, message) {
        super(message);
        this.name = "DiffReviewRequestError";
        this.code = code;
    }
}
const DIFF_REVIEW_INITIALIZE_TIMEOUT_MS = 10_000;
const DIFF_REVIEW_CLOSE_TIMEOUT_MS = 1_000;
export class DiffReviewSession {
    sessionId;
    snapshot;
    persona;
    config;
    deps;
    toolExecutionBackend;
    createThreadSession;
    socketPath;
    authToken;
    server;
    threads = new Map();
    reviewAgentRecords = new Map();
    connections = new Set();
    uiStateListeners = new Set();
    uiState = {
        reviewAgents: [],
    };
    initializedConnection;
    initializeTimeout;
    completionResolver;
    completionPromise;
    completedResult;
    closed = false;
    constructor(options) {
        this.sessionId = `tau-diff-review-${randomUUID()}`;
        this.snapshot = options.snapshot;
        this.persona = options.persona;
        this.config = options.config;
        this.deps = options.deps ?? createDefaultCoreDeps();
        this.toolExecutionBackend = options.toolExecutionBackend;
        this.socketPath = join("/tmp", `tau-diff-${randomBytes(8).toString("hex")}.sock`);
        this.authToken = randomBytes(24).toString("hex");
        this.createThreadSession =
            options.createThread ??
                ((threadOptions) => new DiffReviewThread({
                    threadId: threadOptions.threadId,
                    snapshot: this.snapshot,
                    persona: this.persona,
                    config: this.config,
                    discoveredSkills: options.discoveredSkills,
                    includeAgentContext: options.includeAgentContext,
                    deps: this.deps,
                    toolExecutionBackend: this.toolExecutionBackend,
                    onUpdate: threadOptions.onUpdate,
                    forkFrom: threadOptions.forkFrom,
                }));
        this.completionPromise = new Promise((resolve) => {
            this.completionResolver = resolve;
        });
    }
    get result() {
        return this.completionPromise;
    }
    get protocolVersion() {
        return DIFF_REVIEW_PROTOCOL_VERSION;
    }
    getUiState() {
        return cloneDiffReviewUiState(this.uiState);
    }
    onUiStateChange(listener) {
        this.uiStateListeners.add(listener);
        listener(this.getUiState());
        return () => {
            this.uiStateListeners.delete(listener);
        };
    }
    get launchEnvironment() {
        return {
            TAU_DIFF_PROTOCOL_VERSION: String(DIFF_REVIEW_PROTOCOL_VERSION),
            TAU_DIFF_SOCKET: this.socketPath,
            TAU_DIFF_TOKEN: this.authToken,
            TAU_DIFF_SESSION_ID: this.sessionId,
            TAU_DIFF_REPO_ROOT: this.snapshot.repoRoot,
            TAU_DIFF_CWD: this.snapshot.cwd,
            TAU_DIFF_ARGS_JSON: JSON.stringify(this.snapshot.diffArgs),
        };
    }
    async start() {
        if (this.server) {
            return;
        }
        if (existsSync(this.socketPath)) {
            rmSync(this.socketPath, { force: true });
        }
        this.server = createServer((socket) => this.handleConnection(socket));
        await new Promise((resolve, reject) => {
            const server = this.server;
            if (!server) {
                reject(new Error("diff review server was not created"));
                return;
            }
            const onError = (error) => {
                server.off("listening", onListening);
                reject(error);
            };
            const onListening = () => {
                server.off("error", onError);
                resolve();
            };
            server.once("error", onError);
            server.once("listening", onListening);
            server.listen(this.socketPath);
        });
        this.armInitializeTimeout();
    }
    async launchTool(diffTool) {
        const env = {
            ...this.deps.env.env(),
            ...(diffTool.env ?? {}),
            ...this.launchEnvironment,
        };
        await new Promise((resolve, reject) => {
            const child = spawn(diffTool.command, diffTool.args ?? [], {
                cwd: this.snapshot.cwd,
                env,
                stdio: "ignore",
                detached: true,
            });
            const onError = (error) => reject(error);
            child.once("error", onError);
            child.once("spawn", () => {
                child.off("error", onError);
                child.unref();
                resolve();
            });
        });
    }
    async cancel(reason = "controller_cancelled") {
        if (this.completedResult) {
            return;
        }
        await this.complete({ status: "cancelled", reason });
    }
    async close() {
        if (this.closed) {
            return;
        }
        this.closed = true;
        this.disposeReviewAgents();
        await this.closeServer();
    }
    handleConnection(socket) {
        const readline = createInterface({ input: socket, crlfDelay: Number.POSITIVE_INFINITY });
        const connection = {
            socket,
            readline,
            initialized: false,
            writeQueue: Promise.resolve(),
            pendingResponses: new Map(),
            requestCounter: 0,
        };
        this.connections.add(connection);
        readline.on("line", (line) => {
            void this.handleLine(connection, line).catch((error) => {
                void this.sendError(connection, null, "internalError", "diff review request failed", {
                    cause: error instanceof Error ? error.message : String(error),
                });
            });
        });
        readline.on("error", () => { });
        socket.on("close", () => {
            readline.close();
            this.rejectPendingToolResponses(connection, new Error("diff review protocol connection closed"));
            this.connections.delete(connection);
            const lostInitializedConnection = this.initializedConnection === connection;
            if (lostInitializedConnection) {
                this.initializedConnection = undefined;
            }
            if (!this.completedResult && lostInitializedConnection) {
                void this.cancel("tool_disconnected");
                return;
            }
            if (!this.completedResult && this.connections.size === 0 && !this.initializedConnection) {
                void this.cancel("tool_disconnected");
            }
        });
        socket.on("error", () => { });
    }
    async handleLine(connection, line) {
        const parsed = parseDiffReviewMessageLine(line);
        if (!parsed.ok) {
            await this.sendMessage(connection, createDiffReviewErrorResponse(parsed.id, parsed.error.code, parsed.error.message, parsed.error.data));
            return;
        }
        if (parsed.message.type === "response") {
            this.handleToolResponse(connection, parsed.message);
            return;
        }
        const request = parsed.message;
        if (this.completedResult && request.method !== "initialize") {
            await this.sendError(connection, request.id, "sessionClosed", "diff review session is closed");
            return;
        }
        if (request.method !== "initialize" && !connection.initialized) {
            await this.sendError(connection, request.id, "notInitialized", "call initialize before other diff review methods");
            return;
        }
        switch (request.method) {
            case "initialize":
                await this.handleInitialize(connection, request);
                return;
            case "session.get_context":
                await this.respond(connection, request.id, request.method, {
                    sessionId: this.sessionId,
                    repoRoot: this.snapshot.repoRoot,
                    cwd: this.snapshot.cwd,
                    diffArgs: [...this.snapshot.diffArgs],
                    diffCommand: this.snapshot.toDiffCommand(),
                });
                return;
            case "session.list_files":
                await this.respond(connection, request.id, request.method, {
                    files: this.snapshot.files.map((file) => ({ ...file })),
                });
                return;
            case "session.get_diff":
                await this.handleGetDiff(connection, request);
                return;
            case "session.set_ui_text":
                this.updateDiffToolUiText(request.params.text);
                await this.respond(connection, request.id, request.method, { status: "updated" });
                return;
            case "thread.submit_message":
                await this.handleThreadSubmit(connection, request);
                return;
            case "session.return_review":
                await this.respond(connection, request.id, request.method, { status: "returned" });
                await this.complete({ status: "returned", review: request.params.review });
                return;
            case "session.cancel":
                await this.respond(connection, request.id, request.method, { status: "cancelled" });
                await this.cancel("tool_cancelled");
                return;
            case "session.close":
                await this.sendError(connection, request.id, "invalidRequest", "session.close is sent by Tau, not the diff tool");
                return;
        }
    }
    handleToolResponse(connection, message) {
        if (message.id === null) {
            return;
        }
        const pending = connection.pendingResponses.get(message.id);
        if (!pending) {
            return;
        }
        connection.pendingResponses.delete(message.id);
        if (message.ok) {
            pending.resolve(message);
            return;
        }
        pending.reject(new Error(message.error.message));
    }
    async handleInitialize(connection, request) {
        if (request.params.token !== this.authToken) {
            await this.sendError(connection, request.id, "unauthorized", "invalid diff review token");
            return;
        }
        if (this.initializedConnection && this.initializedConnection !== connection) {
            await this.sendError(connection, request.id, "invalidRequest", "diff review session already has an active client");
            return;
        }
        const result = {
            protocolVersion: DIFF_REVIEW_PROTOCOL_VERSION,
            sessionId: this.sessionId,
            methods: [...DIFF_REVIEW_CLIENT_METHODS],
            alreadyInitialized: connection.initialized,
        };
        connection.initialized = true;
        this.initializedConnection = connection;
        this.clearInitializeTimeout();
        await this.sendMessage(connection, createDiffReviewSuccessResponse(request.id, result));
    }
    async handleGetDiff(connection, request) {
        if (!request.params.path) {
            await this.respond(connection, request.id, request.method, {
                scope: "session",
                patch: this.snapshot.patch,
            });
            return;
        }
        const patch = this.snapshot.getFilePatch(request.params.path);
        if (patch === undefined) {
            await this.sendError(connection, request.id, "invalidParams", `unknown diff file '${request.params.path}'`);
            return;
        }
        await this.respond(connection, request.id, request.method, {
            scope: "file",
            path: request.params.path,
            patch,
        });
    }
    async handleThreadSubmit(connection, request) {
        let acquired;
        try {
            acquired = this.acquireThreadForSubmit(request.params);
            const response = await acquired.thread.submitMessage(request.params.message);
            if (this.completedResult || connection.socket.destroyed) {
                return;
            }
            await this.respond(connection, request.id, request.method, {
                threadId: acquired.threadId,
                response,
            });
        }
        catch (error) {
            if (this.completedResult || connection.socket.destroyed) {
                return;
            }
            if (error instanceof DiffReviewRequestError) {
                await this.sendError(connection, request.id, error.code, error.message);
                return;
            }
            await this.sendError(connection, request.id, "internalError", error instanceof Error ? error.message : String(error));
        }
        finally {
            if (acquired) {
                this.markReviewAgentIdle(acquired.threadId);
            }
        }
    }
    acquireThreadForSubmit(params) {
        if (params.threadId) {
            const existing = this.threads.get(params.threadId);
            if (!existing) {
                throw new DiffReviewRequestError("invalidParams", `unknown thread '${params.threadId}'`);
            }
            this.markReviewAgentRunning(params.threadId);
            return { threadId: params.threadId, thread: existing };
        }
        const threadId = randomUUID();
        const forkSource = params.forkFromThreadId
            ? this.threads.get(params.forkFromThreadId)
            : undefined;
        if (params.forkFromThreadId && !forkSource) {
            throw new DiffReviewRequestError("invalidParams", `unknown fork source thread '${params.forkFromThreadId}'`);
        }
        if (params.forkFromThreadId && this.hasActiveReviewAgentRequest(params.forkFromThreadId)) {
            throw new DiffReviewRequestError("invalidRequest", `thread '${params.forkFromThreadId}' already has an active request and cannot be used as a fork source`);
        }
        const created = this.createThreadSession({
            threadId,
            ...(forkSource ? { forkFrom: forkSource.createForkSource() } : {}),
            onUpdate: (update) => this.applyReviewAgentUpdate(threadId, update),
        });
        this.threads.set(threadId, created);
        this.markReviewAgentRunning(threadId);
        return { threadId, thread: created };
    }
    updateDiffToolUiText(text) {
        const nextText = text.trim() || undefined;
        if (this.uiState.diffToolUiText === nextText) {
            return;
        }
        this.uiState = {
            ...(nextText ? { diffToolUiText: nextText } : {}),
            reviewAgents: this.uiState.reviewAgents,
        };
        this.emitUiState();
    }
    ensureReviewAgentRecord(threadId) {
        const existing = this.reviewAgentRecords.get(threadId);
        if (existing) {
            return existing;
        }
        const created = {
            threadId,
            status: "idle",
            activeRequestCount: 0,
            costTotal: 0,
            usage: createEmptyReviewAgentUsage(this.persona.model.contextWindow),
        };
        this.reviewAgentRecords.set(threadId, created);
        return created;
    }
    markReviewAgentRunning(threadId) {
        const record = this.ensureReviewAgentRecord(threadId);
        if (record.activeRequestCount > 0) {
            throw new DiffReviewRequestError("invalidRequest", `thread '${threadId}' already has an active request`);
        }
        record.activeRequestCount = 1;
        if (record.status !== "running") {
            record.status = "running";
            this.syncReviewAgents();
        }
    }
    markReviewAgentIdle(threadId) {
        const record = this.ensureReviewAgentRecord(threadId);
        if (record.activeRequestCount > 0) {
            record.activeRequestCount -= 1;
        }
        if (record.activeRequestCount === 0 && record.status !== "idle") {
            record.status = "idle";
            this.syncReviewAgents();
        }
    }
    applyReviewAgentUpdate(threadId, update) {
        const record = this.ensureReviewAgentRecord(threadId);
        let changed = false;
        if (record.costTotal !== update.costTotal) {
            record.costTotal = update.costTotal;
            changed = true;
        }
        if (!isSameUsageSnapshot(record.usage, update.usage)) {
            record.usage = { ...update.usage };
            changed = true;
        }
        if (record.lastActivityText !== update.lastActivityText) {
            record.lastActivityText = update.lastActivityText;
            changed = true;
        }
        if (changed) {
            this.syncReviewAgents();
        }
    }
    hasActiveReviewAgentRequest(threadId) {
        return (this.reviewAgentRecords.get(threadId)?.activeRequestCount ?? 0) > 0;
    }
    interruptAllReviewAgents() {
        for (const thread of this.threads.values()) {
            thread.interrupt();
        }
    }
    disposeReviewAgents() {
        for (const thread of this.threads.values()) {
            thread.dispose();
        }
    }
    syncReviewAgents() {
        this.uiState = {
            ...this.uiState,
            reviewAgents: [...this.reviewAgentRecords.values()].map((agent) => ({
                threadId: agent.threadId,
                status: agent.status,
                costTotal: agent.costTotal,
                usage: { ...agent.usage },
                ...(agent.lastActivityText ? { lastActivityText: agent.lastActivityText } : {}),
            })),
        };
        this.emitUiState();
    }
    emitUiState() {
        const state = this.getUiState();
        for (const listener of this.uiStateListeners) {
            listener(state);
        }
    }
    armInitializeTimeout() {
        if (this.initializeTimeout || this.completedResult || this.initializedConnection) {
            return;
        }
        this.initializeTimeout = setTimeout(() => {
            this.initializeTimeout = undefined;
            if (this.completedResult || this.initializedConnection) {
                return;
            }
            void this.cancel("tool_disconnected");
        }, DIFF_REVIEW_INITIALIZE_TIMEOUT_MS);
    }
    clearInitializeTimeout() {
        if (!this.initializeTimeout) {
            return;
        }
        clearTimeout(this.initializeTimeout);
        this.initializeTimeout = undefined;
    }
    async complete(result) {
        if (this.completedResult) {
            return;
        }
        this.clearInitializeTimeout();
        this.completedResult = result;
        this.interruptAllReviewAgents();
        this.completionResolver?.(result);
        this.completionResolver = undefined;
        await this.requestToolShutdown();
        await this.close();
    }
    async requestToolShutdown() {
        const connection = this.initializedConnection;
        if (!connection || connection.socket.destroyed) {
            return;
        }
        try {
            await Promise.race([
                this.requestTool(connection, "session.close", {}),
                new Promise((_, reject) => {
                    setTimeout(() => {
                        reject(new Error("diff review tool close request timed out"));
                    }, DIFF_REVIEW_CLOSE_TIMEOUT_MS);
                }),
            ]);
        }
        catch {
            // ignore
        }
    }
    async requestTool(connection, method, params) {
        if (connection.socket.destroyed ||
            connection.socket.writableEnded ||
            !connection.socket.writable) {
            throw new Error("diff review protocol socket is not available");
        }
        const id = `tau-${++connection.requestCounter}`;
        const request = {
            version: DIFF_REVIEW_PROTOCOL_VERSION,
            type: "request",
            id,
            method,
            params,
        };
        return await new Promise((resolve, reject) => {
            connection.pendingResponses.set(id, {
                method,
                resolve: (message) => {
                    if (!message.ok) {
                        reject(new Error(message.error.message));
                        return;
                    }
                    resolve(message.result);
                },
                reject,
            });
            void this.sendMessage(connection, request).catch((error) => {
                const pending = connection.pendingResponses.get(id);
                if (!pending) {
                    return;
                }
                connection.pendingResponses.delete(id);
                pending.reject(error instanceof Error ? error : new Error(String(error)));
            });
        });
    }
    async closeServer() {
        const server = this.server;
        this.server = undefined;
        this.clearInitializeTimeout();
        const connectionClosePromises = [];
        for (const connection of this.connections) {
            connection.readline.close();
            this.rejectPendingToolResponses(connection, new Error("diff review protocol connection closed"));
            if (!connection.socket.destroyed) {
                connectionClosePromises.push(new Promise((resolve) => {
                    const socket = connection.socket;
                    const destroyTimeout = setTimeout(() => {
                        socket.destroy();
                    }, 100);
                    socket.once("close", () => {
                        clearTimeout(destroyTimeout);
                        resolve();
                    });
                    socket.end();
                }));
            }
        }
        this.connections.clear();
        this.initializedConnection = undefined;
        await Promise.all(connectionClosePromises);
        if (server) {
            await new Promise((resolve) => {
                server.close(() => resolve());
            });
        }
        if (existsSync(this.socketPath)) {
            rmSync(this.socketPath, { force: true });
        }
    }
    rejectPendingToolResponses(connection, error) {
        for (const [id, pending] of connection.pendingResponses) {
            connection.pendingResponses.delete(id);
            pending.reject(error);
        }
    }
    async respond(connection, id, _method, result) {
        await this.sendMessage(connection, createDiffReviewSuccessResponse(id, result));
    }
    async sendError(connection, id, code, message, data) {
        const resolvedCode = DIFF_REVIEW_ERROR_CODES[code];
        await this.sendMessage(connection, createDiffReviewErrorResponse(id, resolvedCode, message, data));
    }
    async sendMessage(connection, message) {
        if (connection.socket.destroyed ||
            connection.socket.writableEnded ||
            !connection.socket.writable) {
            return;
        }
        const payload = serializeDiffReviewMessage(message);
        const writePromise = connection.writeQueue.then(() => new Promise((resolve, reject) => {
            if (connection.socket.destroyed ||
                connection.socket.writableEnded ||
                !connection.socket.writable) {
                resolve();
                return;
            }
            let settled = false;
            const finish = (callback) => {
                if (settled) {
                    return;
                }
                settled = true;
                connection.socket.off("close", onClose);
                connection.socket.off("drain", onDrain);
                connection.socket.off("error", onError);
                callback();
            };
            const onClose = () => {
                finish(resolve);
            };
            const onDrain = () => {
                finish(resolve);
            };
            const onError = (error) => {
                finish(() => {
                    reject(error);
                });
            };
            connection.socket.once("close", onClose);
            connection.socket.once("error", onError);
            try {
                if (connection.socket.write(payload)) {
                    finish(resolve);
                    return;
                }
                connection.socket.once("drain", onDrain);
            }
            catch (error) {
                finish(() => {
                    reject(error instanceof Error ? error : new Error(String(error)));
                });
            }
        }));
        connection.writeQueue = writePromise.catch(() => { });
        try {
            await writePromise;
        }
        catch {
            if (!connection.socket.destroyed) {
                connection.socket.destroy();
            }
        }
    }
}
function createEmptyReviewAgentUsage(contextWindow) {
    return {
        input: 0,
        output: 0,
        cacheRead: 0,
        cacheWrite: 0,
        contextWindowUsageTokens: 0,
        contextWindow,
    };
}
function isSameUsageSnapshot(left, right) {
    return (left.input === right.input &&
        left.output === right.output &&
        left.cacheRead === right.cacheRead &&
        left.cacheWrite === right.cacheWrite &&
        left.contextWindowUsageTokens === right.contextWindowUsageTokens &&
        left.contextWindow === right.contextWindow);
}
function cloneDiffReviewUiState(state) {
    return {
        ...(state.diffToolUiText ? { diffToolUiText: state.diffToolUiText } : {}),
        reviewAgents: state.reviewAgents.map((agent) => ({
            threadId: agent.threadId,
            status: agent.status,
            costTotal: agent.costTotal,
            usage: { ...agent.usage },
            ...(agent.lastActivityText ? { lastActivityText: agent.lastActivityText } : {}),
        })),
    };
}
export async function startDiffReviewSession(options) {
    throwIfDiffReviewStartAborted(options.signal);
    const snapshot = await captureDiffReviewSnapshot({
        cwd: options.cwd,
        source: options.source,
        signal: options.signal,
        deps: options.deps,
    });
    throwIfDiffReviewStartAborted(options.signal);
    const session = new DiffReviewSession({
        snapshot,
        persona: options.persona,
        config: options.config,
        discoveredSkills: options.discoveredSkills,
        includeAgentContext: options.includeAgentContext,
        deps: options.deps,
        toolExecutionBackend: options.toolExecutionBackend,
    });
    await session.start();
    try {
        throwIfDiffReviewStartAborted(options.signal);
        await session.launchTool(options.diffTool);
    }
    catch (error) {
        await session.cancel("launch_failed");
        throw error;
    }
    return {
        session,
        result: session.result,
    };
}
function throwIfDiffReviewStartAborted(signal) {
    if (!signal?.aborted) {
        return;
    }
    throw new Error("diff review start aborted");
}
//# sourceMappingURL=session.js.map