import { spawnSync } from "node:child_process";
import { readFile } from "node:fs/promises";
import path from "node:path";
import { PI_AGENT_CAPABILITIES, PI_THINKING_LEVELS, } from "../shared/agent.js";
import { findPiLaunchProfile, listPiLaunchProfiles, piProfileAsLaunchProfile } from "./pi-launch.js";
import { resolvePiCli } from "./pi-cli.js";
import { PiRpcClient } from "./pi-rpc.js";
import { getPiSession, listPiSessions, listPiWorkspaces, readPiSessionRecord, resolvePiSessionDir, } from "./pi-state.js";
export class PiSessionService {
    config;
    sessionDir;
    cliPath;
    currentLaunchProfile;
    rpc = null;
    currentWorkspace;
    currentThreadId = null;
    currentSessionPath;
    currentModel;
    currentThinking;
    currentSessionName;
    processing = false;
    cachedStats = {};
    constructor(config) {
        this.config = config;
        const cli = resolvePiCli(process.env, config.piCliPath);
        if (!cli.path) {
            throw new Error("Pi CLI not found. Install Pi from https://pi.dev/ or set PI_CLI_PATH.");
        }
        this.cliPath = cli.path;
        this.sessionDir = resolvePiSessionDir({ sessionDir: config.piSessionDir });
        this.currentWorkspace = config.workspace;
        this.currentModel = config.piDefaultModel;
        this.currentThinking = config.piDefaultThinking;
        this.currentLaunchProfile = findPiLaunchProfile(config.piDefaultLaunchProfileId);
    }
    static async create(config, options) {
        const service = new PiSessionService(config);
        service.currentWorkspace = options?.workspace ?? config.workspace;
        service.currentModel = options?.model ?? config.piDefaultModel;
        service.currentThinking = options?.reasoningEffort ?? config.piDefaultThinking;
        service.currentLaunchProfile = findPiLaunchProfile(options?.launchProfileId ?? config.piDefaultLaunchProfileId);
        if (options?.sessionPath) {
            await service.switchSession(options.sessionPath);
            return service;
        }
        if (options?.resumeThreadId) {
            await service.resumeThread(options.resumeThreadId);
            return service;
        }
        if (!options?.deferThreadStart) {
            await service.newThread(service.currentWorkspace, service.currentModel);
        }
        return service;
    }
    getInfo() {
        return {
            agentId: "pi",
            agentLabel: "Pi",
            threadId: this.currentThreadId,
            workspace: this.currentWorkspace,
            model: this.currentModel,
            reasoningEffort: this.currentThinking,
            launchProfileId: this.currentLaunchProfile.id,
            launchProfileLabel: this.currentLaunchProfile.label,
            launchProfileBehavior: this.currentLaunchProfile.behavior,
            sandboxMode: "host",
            approvalPolicy: "never",
            fastMode: false,
            unsafeLaunch: false,
            sessionPath: this.currentSessionPath,
            sessionUsage: this.cachedStats.sessionUsage,
            contextUsage: this.cachedStats.contextUsage,
            capabilities: PI_AGENT_CAPABILITIES,
        };
    }
    isProcessing() {
        return this.processing;
    }
    getActiveThreadId() {
        return this.currentThreadId;
    }
    hasActiveThread() {
        return Boolean(this.currentThreadId || this.currentSessionPath);
    }
    getCurrentWorkspace() {
        return this.currentWorkspace;
    }
    async prompt(input, callbacks) {
        if (this.processing) {
            throw new Error("A Pi turn is already in progress");
        }
        await this.ensureSessionStarted();
        const rpc = this.getRpc();
        const promptPayload = await this.buildPromptPayload(input);
        const lastToolOutput = new Map();
        let didEnd = false;
        this.processing = true;
        const off = rpc.onEvent((event) => {
            try {
                switch (event.type) {
                    case "message_update":
                        handleMessageUpdate(event, callbacks);
                        break;
                    case "tool_execution_start": {
                        const toolCallId = stringValue(event.toolCallId) ?? "tool";
                        const toolName = stringValue(event.toolName) ?? "tool";
                        lastToolOutput.set(toolCallId, "");
                        callbacks.onToolStart(toolName, toolCallId);
                        break;
                    }
                    case "tool_execution_update": {
                        const toolCallId = stringValue(event.toolCallId) ?? "tool";
                        const text = extractContentText(objectValue(event.partialResult));
                        const previous = lastToolOutput.get(toolCallId) ?? "";
                        const delta = computeTextDelta(previous, text);
                        lastToolOutput.set(toolCallId, text);
                        if (delta) {
                            callbacks.onToolUpdate(toolCallId, delta);
                        }
                        break;
                    }
                    case "tool_execution_end": {
                        const toolCallId = stringValue(event.toolCallId) ?? "tool";
                        const resultText = extractContentText(objectValue(event.result));
                        const previous = lastToolOutput.get(toolCallId) ?? "";
                        const delta = computeTextDelta(previous, resultText);
                        if (delta) {
                            callbacks.onToolUpdate(toolCallId, delta);
                        }
                        callbacks.onToolEnd(toolCallId, Boolean(event.isError));
                        break;
                    }
                    case "turn_end":
                        this.refreshFromTurnEnd(event, callbacks);
                        break;
                    case "agent_end":
                        didEnd = true;
                        callbacks.onAgentEnd();
                        break;
                    case "extension_error": {
                        const toolCallId = stringValue(event.id) ?? "extension-error";
                        callbacks.onToolStart("extension_error", toolCallId);
                        callbacks.onToolUpdate(toolCallId, stringValue(event.error) ?? "Extension error");
                        callbacks.onToolEnd(toolCallId, true);
                        break;
                    }
                    default:
                        break;
                }
            }
            catch (error) {
                console.error("Failed to handle Pi RPC event:", error);
            }
        });
        try {
            await rpc.send({ type: "prompt", ...promptPayload }, 30_000);
            await this.waitForAgentEnd(() => didEnd);
            await this.refreshState().catch(() => { });
            await this.refreshStats().catch(() => { });
        }
        finally {
            off();
            this.processing = false;
        }
    }
    async abort() {
        if (!this.rpc) {
            this.processing = false;
            return;
        }
        await this.rpc.send({ type: "abort" }, 10_000).catch(() => { });
        this.processing = false;
    }
    async newThread(workspace, model) {
        this.ensureIdle("start a new Pi session");
        this.currentWorkspace = workspace ?? this.currentWorkspace;
        if (model) {
            this.currentModel = model;
        }
        this.currentThreadId = null;
        this.currentSessionPath = undefined;
        this.cachedStats = {};
        this.restartRpc();
        await this.refreshState();
        return this.getInfo();
    }
    async resumeThread(threadId) {
        return this.switchSession(threadId);
    }
    async switchSession(threadId) {
        const record = getPiSession(threadId, { sessionDir: this.sessionDir });
        if (!record) {
            throw new Error(`Unknown Pi session: ${threadId}`);
        }
        this.currentWorkspace = record.cwd;
        this.currentThreadId = record.id;
        this.currentSessionPath = record.sessionPath;
        this.currentModel = record.model ?? this.currentModel;
        this.currentThinking = record.reasoningEffort ?? this.currentThinking;
        this.cachedStats = {};
        if (this.rpc) {
            const result = await this.rpc.send({ type: "switch_session", sessionPath: record.sessionPath }, 30_000);
            if (objectValue(result.data)?.cancelled === true) {
                throw new Error("Pi session switch was cancelled by an extension");
            }
        }
        else {
            this.restartRpc();
        }
        await this.refreshState();
        await this.refreshStats().catch(() => { });
        return this.getInfo();
    }
    listAllSessions(limit) {
        return listPiSessions(limit ?? 20, { sessionDir: this.sessionDir });
    }
    listWorkspaces() {
        const workspaces = new Set(listPiWorkspaces({ sessionDir: this.sessionDir }));
        workspaces.add(this.currentWorkspace);
        workspaces.add(this.config.workspace);
        return [...workspaces].sort((left, right) => left.localeCompare(right));
    }
    async refreshModels() {
        // Pi models are read from the CLI on each listModels() call.
    }
    listModels() {
        const result = spawnSync(this.cliPath, ["--list-models"], {
            cwd: this.currentWorkspace,
            env: process.env,
            encoding: "utf8",
            timeout: 10_000,
        });
        if (result.status !== 0 || !result.stdout) {
            return this.currentModel ? [{ slug: this.currentModel, displayName: this.currentModel }] : [];
        }
        const records = [];
        for (const line of result.stdout
            .split(/\r?\n/)
            .slice(1)
            .map((line) => line.trim())
            .filter(Boolean)) {
            const parts = line.split(/\s+/);
            const provider = parts[0];
            const model = parts[1];
            if (!provider || !model) {
                continue;
            }
            const slug = `${provider}/${model}`;
            const contextWindow = parseCompactTokenCount(parts[2]);
            const maxOutputTokens = parseCompactTokenCount(parts[3]);
            const supportsThinking = parseYesNo(parts[4]);
            const supportsImages = parseYesNo(parts[5]);
            records.push({
                slug,
                displayName: slug,
                ...(contextWindow !== undefined ? { maxInputTokens: contextWindow, contextWindow } : {}),
                ...(maxOutputTokens !== undefined ? { maxOutputTokens } : {}),
                ...(supportsThinking !== undefined ? { supportsThinking } : {}),
                ...(supportsImages !== undefined ? { supportsImages } : {}),
            });
        }
        if (this.currentModel && !records.some((record) => record.slug === this.currentModel)) {
            records.unshift({ slug: this.currentModel, displayName: this.currentModel });
        }
        return records;
    }
    listLaunchProfiles() {
        return listPiLaunchProfiles();
    }
    getSessionRecord(threadId) {
        return getPiSession(threadId, { sessionDir: this.sessionDir });
    }
    setModel(slug) {
        this.currentModel = slug;
        this.restartRpcIfIdle();
        return slug;
    }
    async setModelForCurrentSession(slug) {
        this.ensureIdle("change Pi model");
        this.currentModel = slug;
        let appliedToActiveThread = false;
        if (this.rpc && this.currentThreadId) {
            const { provider, modelId } = splitPiModelSlug(slug);
            await this.rpc.send({ type: "set_model", provider, modelId }, 30_000);
            await this.refreshState().catch(() => { });
            appliedToActiveThread = true;
        }
        else {
            this.restartRpcIfIdle();
        }
        return { value: slug, appliedToActiveThread };
    }
    setReasoningEffort(effort) {
        this.currentThinking = normalizePiThinking(effort);
        this.restartRpcIfIdle();
    }
    async setReasoningEffortForCurrentSession(effort) {
        this.ensureIdle("change Pi thinking level");
        const level = normalizePiThinking(effort);
        this.currentThinking = level;
        let appliedToActiveThread = false;
        if (this.rpc && this.currentThreadId) {
            await this.rpc.send({ type: "set_thinking_level", level }, 30_000);
            await this.refreshState().catch(() => { });
            appliedToActiveThread = true;
        }
        else {
            this.restartRpcIfIdle();
        }
        return { value: level, appliedToActiveThread };
    }
    setLaunchProfile(profileId) {
        this.ensureIdle("change Pi profile");
        this.currentLaunchProfile = findPiLaunchProfile(profileId);
        this.restartRpcIfIdle();
        return piProfileAsLaunchProfile(this.currentLaunchProfile);
    }
    setFastMode() {
        throw new Error("Fast mode is only supported by Codex sessions");
    }
    getSelectedLaunchProfile() {
        return piProfileAsLaunchProfile(this.currentLaunchProfile);
    }
    syncFromAgentState(options = {}) {
        const before = {
            threadId: this.currentThreadId,
            workspace: this.currentWorkspace,
            model: this.currentModel,
            thinking: this.currentThinking,
            sessionPath: this.currentSessionPath,
        };
        const changedFields = new Set();
        const record = this.currentSessionPath
            ? readPiSessionRecord(this.currentSessionPath, path.basename(path.dirname(this.currentSessionPath)))
            : this.currentThreadId
                ? getPiSession(this.currentThreadId, { sessionDir: this.sessionDir })
                : null;
        if (record) {
            if (record.id !== this.currentThreadId)
                changedFields.add("thread");
            if (record.cwd && record.cwd !== this.currentWorkspace)
                changedFields.add("workspace");
            if (record.model && record.model !== this.currentModel)
                changedFields.add("model");
            if (record.reasoningEffort && record.reasoningEffort !== this.currentThinking)
                changedFields.add("reasoning");
            if (record.sessionPath !== this.currentSessionPath)
                changedFields.add("sessionPath");
            this.currentThreadId = record.id;
            this.currentWorkspace = record.cwd || this.currentWorkspace;
            this.currentSessionPath = record.sessionPath;
            this.currentModel = record.model ?? this.currentModel;
            this.currentThinking = record.reasoningEffort ?? this.currentThinking;
            if (changedFields.has("thread") || changedFields.has("sessionPath")) {
                this.cachedStats = {};
            }
        }
        const changed = changedFields.size > 0 ||
            before.threadId !== this.currentThreadId ||
            before.workspace !== this.currentWorkspace ||
            before.model !== this.currentModel ||
            before.thinking !== this.currentThinking ||
            before.sessionPath !== this.currentSessionPath;
        let reattached = false;
        if (changed && options.reattach && this.rpc && !this.processing) {
            this.restartRpc();
            reattached = true;
        }
        return {
            threadId: this.currentThreadId,
            changed,
            reattached,
            changedFields: [...changedFields],
            info: this.getInfo(),
        };
    }
    handback() {
        const threadId = this.currentThreadId;
        const workspace = this.currentWorkspace;
        const sessionPath = this.currentSessionPath;
        this.rpc?.stop();
        this.rpc = null;
        this.currentThreadId = null;
        this.currentSessionPath = undefined;
        return {
            threadId,
            workspace,
            command: sessionPath
                ? `cd ${shellQuote(workspace)} && pi --session ${shellQuote(sessionPath)}`
                : undefined,
            label: "Pi CLI",
        };
    }
    dispose() {
        this.rpc?.stop();
        this.rpc = null;
        this.processing = false;
    }
    async ensureSessionStarted() {
        if (this.currentThreadId || this.currentSessionPath) {
            this.getRpc().ensureStarted();
            await this.refreshState().catch(() => { });
            return;
        }
        await this.newThread(this.currentWorkspace, this.currentModel);
    }
    restartRpc() {
        this.rpc?.stop();
        this.rpc = new PiRpcClient({
            commandPath: this.cliPath,
            cwd: this.currentWorkspace,
            sessionDir: this.sessionDir,
            sessionPath: this.currentSessionPath,
            model: this.currentModel,
            thinking: this.currentThinking,
            ...this.currentLaunchProfile.cli,
            env: { PI_CODING_AGENT_SESSION_DIR: this.sessionDir },
        });
    }
    restartRpcIfIdle() {
        if (this.processing) {
            return;
        }
        this.restartRpc();
    }
    getRpc() {
        if (!this.rpc) {
            this.restartRpc();
        }
        return this.rpc;
    }
    ensureIdle(action) {
        if (this.processing) {
            throw new Error(`Cannot ${action} while a turn is in progress`);
        }
    }
    async refreshState() {
        const response = await this.getRpc().send({ type: "get_state" }, 30_000);
        const data = objectValue(response.data);
        if (!data) {
            return;
        }
        const model = objectValue(data.model);
        const provider = stringValue(model?.provider);
        const modelId = stringValue(model?.id);
        this.currentModel = provider && modelId ? `${provider}/${modelId}` : modelId ?? this.currentModel;
        this.currentThinking = stringValue(data.thinkingLevel) ?? this.currentThinking;
        this.currentSessionPath = stringValue(data.sessionFile) ?? this.currentSessionPath;
        this.currentThreadId = stringValue(data.sessionId) ?? this.currentThreadId;
        this.currentSessionName = stringValue(data.sessionName) ?? this.currentSessionName;
        if (this.currentSessionPath) {
            const record = readPiSessionRecord(this.currentSessionPath, path.basename(path.dirname(this.currentSessionPath)));
            if (record?.cwd) {
                this.currentWorkspace = record.cwd;
            }
        }
    }
    async refreshStats() {
        const response = await this.getRpc().send({ type: "get_session_stats" }, 30_000);
        const data = objectValue(response.data);
        if (!data) {
            return;
        }
        const tokens = objectValue(data.tokens);
        const contextUsage = objectValue(data.contextUsage);
        this.cachedStats = {
            sessionUsage: tokens
                ? {
                    input: numberValue(tokens.input) ?? 0,
                    output: numberValue(tokens.output) ?? 0,
                    cacheRead: numberValue(tokens.cacheRead) ?? 0,
                    cacheWrite: numberValue(tokens.cacheWrite) ?? 0,
                    total: numberValue(tokens.total) ?? 0,
                    cost: numberValue(data.cost) ?? undefined,
                }
                : undefined,
            contextUsage: contextUsage
                ? {
                    tokens: numberValue(contextUsage.tokens),
                    contextWindow: numberValue(contextUsage.contextWindow),
                    percent: numberValue(contextUsage.percent),
                }
                : undefined,
        };
    }
    async buildPromptPayload(input) {
        if (typeof input === "string") {
            return { message: input };
        }
        const textParts = [input.stagedFileInstructions, input.text].filter((part) => Boolean(part?.trim()));
        const images = await Promise.all((input.imagePaths ?? []).map(async (imagePath) => ({
            type: "image",
            data: (await readFile(imagePath)).toString("base64"),
            mimeType: mimeTypeForImage(imagePath),
        })));
        if (images.length > 0) {
            const imageSupport = this.currentModelSupportsImages();
            if (imageSupport === false) {
                throw new Error(`Current Pi model does not support image input: ${this.currentModel}`);
            }
        }
        return {
            message: textParts.join("\n\n") || "Please inspect the attached file(s).",
            ...(images.length > 0 ? { images } : {}),
        };
    }
    currentModelSupportsImages() {
        const model = this.currentModel;
        if (!model) {
            return null;
        }
        const record = this.listModels().find((candidate) => candidate.slug === model || candidate.slug.endsWith(`/${model}`));
        return record?.supportsImages ?? null;
    }
    refreshFromTurnEnd(event, callbacks) {
        const message = objectValue(event.message);
        const usage = objectValue(message?.usage);
        if (!usage) {
            return;
        }
        callbacks.onTurnComplete?.({
            inputTokens: numberValue(usage.input) ?? 0,
            cachedInputTokens: numberValue(usage.cacheRead) ?? 0,
            outputTokens: numberValue(usage.output) ?? 0,
        });
    }
    async waitForAgentEnd(done) {
        if (done()) {
            return;
        }
        await new Promise((resolve, reject) => {
            const startedAt = Date.now();
            const timer = setInterval(() => {
                if (done()) {
                    clearInterval(timer);
                    resolve();
                    return;
                }
                if (!this.processing) {
                    clearInterval(timer);
                    resolve();
                    return;
                }
                if (Date.now() - startedAt > 24 * 60 * 60 * 1000) {
                    clearInterval(timer);
                    reject(new Error("Timed out waiting for Pi agent to finish"));
                }
            }, 250);
            timer.unref?.();
        });
    }
}
function handleMessageUpdate(event, callbacks) {
    const update = objectValue(event.assistantMessageEvent);
    const updateType = stringValue(update?.type);
    if (updateType === "text_delta") {
        const delta = stringValue(update?.delta);
        if (delta) {
            callbacks.onTextDelta(delta);
        }
    }
    else if (updateType === "toolcall_end") {
        const toolCall = objectValue(update?.toolCall);
        const id = stringValue(toolCall?.id);
        const name = stringValue(toolCall?.name);
        if (id && name) {
            callbacks.onToolStart(name, id);
        }
    }
    else if (updateType === "error") {
        const toolId = "pi-message-error";
        callbacks.onToolStart("message_error", toolId);
        callbacks.onToolUpdate(toolId, stringValue(update?.reason) ?? "Message error");
        callbacks.onToolEnd(toolId, true);
    }
}
function objectValue(value) {
    return typeof value === "object" && value !== null ? value : null;
}
function stringValue(value) {
    return typeof value === "string" && value.trim() ? value : null;
}
function numberValue(value) {
    return typeof value === "number" && Number.isFinite(value) ? value : null;
}
function extractContentText(container) {
    const content = container?.content;
    if (typeof content === "string") {
        return content;
    }
    if (!Array.isArray(content)) {
        return "";
    }
    return content
        .map((entry) => {
        const block = objectValue(entry);
        return stringValue(block?.text) ?? "";
    })
        .join("");
}
function computeTextDelta(previousText, nextText) {
    return nextText.startsWith(previousText) ? nextText.slice(previousText.length) : nextText;
}
function splitPiModelSlug(slug) {
    const separatorIndex = slug.indexOf("/");
    if (separatorIndex === -1) {
        return { provider: "openai-codex", modelId: slug };
    }
    return {
        provider: slug.slice(0, separatorIndex),
        modelId: slug.slice(separatorIndex + 1),
    };
}
function normalizePiThinking(value) {
    if (PI_THINKING_LEVELS.includes(value)) {
        return value;
    }
    throw new Error(`Unsupported Pi thinking level: ${value}`);
}
function parseCompactTokenCount(value) {
    if (!value) {
        return undefined;
    }
    const match = value.match(/^(\d+(?:\.\d+)?)([KMB])?$/i);
    if (!match) {
        return undefined;
    }
    const number = Number(match[1]);
    const unit = match[2]?.toUpperCase();
    const multiplier = unit === "M" ? 1_000_000 : unit === "K" ? 1_000 : unit === "B" ? 1_000_000_000 : 1;
    return Math.round(number * multiplier);
}
function parseYesNo(value) {
    if (!value) {
        return undefined;
    }
    const normalized = value.toLowerCase();
    if (["yes", "true", "1"].includes(normalized)) {
        return true;
    }
    if (["no", "false", "0"].includes(normalized)) {
        return false;
    }
    return undefined;
}
function mimeTypeForImage(filePath) {
    switch (path.extname(filePath).toLowerCase()) {
        case ".png":
            return "image/png";
        case ".webp":
            return "image/webp";
        case ".gif":
            return "image/gif";
        default:
            return "image/jpeg";
    }
}
function shellQuote(value) {
    return `'${value.replace(/'/g, `'\\''`)}'`;
}
