/*
 * Copyright OpenSearch Contributors
 * SPDX-License-Identifier: Apache-2.0
 */
import { ToolCallStatus } from '../../../types/index.js';
import { SubprocessConnector } from '../../../services/connectors/subprocess/SubprocessConnector.js';
/**
 * Default Pi configuration
 */
const PI_DEFAULT_CONFIG = {
    command: 'pi',
    args: ['--print', '--mode', 'json'],
    env: {},
    inputMode: 'stdin',
    outputParser: 'streaming',
    timeout: 600000, // 10 minutes
};
/**
 * Pi CLI Connector
 * Invokes pi.dev as a subprocess for agent evaluation
 */
export class PiConnector extends SubprocessConnector {
    constructor(config) {
        super({ ...PI_DEFAULT_CONFIG, ...config });
        this.type = 'pi';
        this.name = 'Pi (pi.dev)';
        this.traceContext = { propagateEnv: true, serviceName: 'pi-agent' };
        this.piOutputBuffer = '';
        this.piThinkingBuffer = '';
        this.piTextBuffer = '';
    }
    /**
     * Build prompt for Pi
     */
    buildPayload(request) {
        const parts = [];
        // Add system context if available
        if (request.testCase.context && request.testCase.context.length > 0) {
            parts.push('## Context');
            for (const ctx of request.testCase.context) {
                parts.push(`**${ctx.description}:**`);
                parts.push(ctx.value);
                parts.push('');
            }
        }
        // Add the main prompt
        parts.push('## Task');
        parts.push(request.testCase.initialPrompt);
        return parts.join('\n');
    }
    /**
     * Parse Pi streaming output (JSON format)
     * Pi's --mode json produces NDJSON lines
     */
    parseStreamingOutput(chunk, trajectory, onProgress) {
        this.piOutputBuffer += chunk;
        // Parse complete JSON lines (NDJSON format)
        const lines = this.piOutputBuffer.split('\n');
        this.piOutputBuffer = lines.pop() || ''; // Keep incomplete line
        for (const line of lines) {
            const trimmed = line.trim();
            if (!trimmed)
                continue;
            try {
                const event = JSON.parse(trimmed);
                const steps = this.parsePiEvent(event);
                for (const step of steps) {
                    trajectory.push(step);
                    onProgress?.(step);
                }
            }
            catch {
                // Not JSON, treat as plain text
                if (trimmed) {
                    const step = this.createStep('assistant', trimmed);
                    trajectory.push(step);
                    onProgress?.(step);
                }
            }
        }
    }
    /**
     * Parse a single JSON event from Pi output.
     *
     * Pi's --mode json NDJSON format uses these event types:
     *  - session, agent_start, agent_end — lifecycle (ignored)
     *  - turn_start, turn_end — turn boundaries
     *  - message_start / message_end — full message with content blocks
     *  - message_update — streaming deltas with assistantMessageEvent
     */
    parsePiEvent(event) {
        const steps = [];
        // message_end contains the full final message with all content blocks
        if (event.type === 'message_end' && event.message?.role === 'assistant') {
            const content = event.message.content;
            if (Array.isArray(content)) {
                for (const block of content) {
                    if (block.type === 'thinking' && block.thinking) {
                        steps.push(this.createStep('thinking', block.thinking));
                    }
                    else if (block.type === 'text' && block.text) {
                        steps.push(this.createStep('assistant', block.text));
                    }
                    else if (block.type === 'tool_use') {
                        steps.push(this.createStep('action', JSON.stringify(block.input || {}), {
                            toolName: block.name,
                            toolArgs: block.input,
                        }));
                    }
                }
            }
        }
        else if (event.type === 'message_update') {
            // Streaming deltas — accumulate into buffers
            const assistantEvent = event.assistantMessageEvent;
            if (assistantEvent?.type === 'text_delta' && assistantEvent.delta) {
                this.piTextBuffer += assistantEvent.delta;
            }
            else if (assistantEvent?.type === 'thinking_delta' && assistantEvent.delta) {
                this.piThinkingBuffer += assistantEvent.delta;
            }
        }
        else if (event.type === 'tool_result') {
            const content = event.content || event.output || JSON.stringify(event);
            steps.push(this.createStep('tool_result', typeof content === 'string' ? content : JSON.stringify(content), { status: event.is_error ? ToolCallStatus.FAILURE : ToolCallStatus.SUCCESS }));
        }
        else if (event.type === 'agent_end') {
            // Flush any remaining buffers on agent_end
            if (this.piThinkingBuffer) {
                steps.push(this.createStep('thinking', this.piThinkingBuffer));
                this.piThinkingBuffer = '';
            }
            if (this.piTextBuffer) {
                steps.push(this.createStep('response', this.piTextBuffer));
                this.piTextBuffer = '';
            }
        }
        return steps;
    }
    /**
     * Parse final output for Pi
     */
    parseResponse(data) {
        const steps = [];
        const response = data.stdout.trim();
        if (response) {
            steps.push(this.createStep('response', response));
        }
        if (data.exitCode !== 0 && data.stderr.trim()) {
            steps.push(this.createStep('tool_result', `Error: ${data.stderr.trim()}`, {
                status: ToolCallStatus.FAILURE,
            }));
        }
        return steps;
    }
    /**
     * Flush remaining buffers when the subprocess stream ends.
     */
    onBeforeStreamEnd(trajectory, onProgress) {
        if (this.piOutputBuffer.trim()) {
            try {
                const event = JSON.parse(this.piOutputBuffer.trim());
                const steps = this.parsePiEvent(event);
                for (const step of steps) {
                    trajectory.push(step);
                    onProgress?.(step);
                }
            }
            catch {
                const step = this.createStep('assistant', this.piOutputBuffer.trim());
                trajectory.push(step);
                onProgress?.(step);
            }
            this.piOutputBuffer = '';
        }
        if (this.piThinkingBuffer) {
            const step = this.createStep('thinking', this.piThinkingBuffer);
            trajectory.push(step);
            onProgress?.(step);
            this.piThinkingBuffer = '';
        }
        if (this.piTextBuffer) {
            const step = this.createStep('response', this.piTextBuffer);
            trajectory.push(step);
            onProgress?.(step);
            this.piTextBuffer = '';
        }
    }
    /**
     * Reset state and apply connectorConfig
     */
    async execute(endpoint, request, auth, onProgress, onRawEvent) {
        this.piOutputBuffer = '';
        this.piThinkingBuffer = '';
        this.piTextBuffer = '';
        // Save original config
        const originalArgs = this.config.args ? [...this.config.args] : [];
        const originalEnv = this.config.env ? structuredClone(this.config.env) : {};
        const originalTimeout = this.config.timeout;
        const originalWorkingDir = this.config.workingDir;
        // Apply connectorConfig
        const piConfig = request.connectorConfig;
        if (piConfig) {
            if (piConfig.env) {
                this.config.env = { ...this.config.env, ...piConfig.env };
            }
            if (piConfig.timeout !== undefined) {
                this.config.timeout = piConfig.timeout;
            }
            if (piConfig.workingDir) {
                this.config.workingDir = piConfig.workingDir;
            }
            // Build additional args from config
            const extraArgs = [];
            if (piConfig.packagePath) {
                // Pi uses --skill and --extension to load package components
                extraArgs.push('--skill', `${piConfig.packagePath}/skills/*`);
                extraArgs.push('--extension', `${piConfig.packagePath}/extensions/agent-health.ts`);
                extraArgs.push('--append-system-prompt', `${piConfig.packagePath}/prompts/agent-health.md`);
            }
            if (piConfig.model) {
                extraArgs.push('--model', piConfig.model);
            }
            if (piConfig.additionalArgs) {
                extraArgs.push(...piConfig.additionalArgs);
            }
            if (extraArgs.length > 0) {
                this.config.args = [...(this.config.args || []), ...extraArgs];
            }
        }
        // Pass --model flag from request if specified
        if (request.modelId) {
            this.config.args = [...(this.config.args || []), '--model', request.modelId];
        }
        // Inherit AWS credentials
        if (process.env.AWS_PROFILE) {
            this.config.env = { ...this.config.env, AWS_PROFILE: process.env.AWS_PROFILE };
        }
        if (process.env.AWS_REGION) {
            this.config.env = { ...this.config.env, AWS_REGION: process.env.AWS_REGION };
        }
        try {
            return await super.execute(endpoint, request, auth, onProgress, onRawEvent);
        }
        finally {
            // Restore config
            this.config.args = originalArgs;
            this.config.env = originalEnv;
            this.config.timeout = originalTimeout;
            this.config.workingDir = originalWorkingDir;
        }
    }
    /**
     * Health check - verify pi command exists
     */
    async healthCheck(endpoint, auth) {
        return super.healthCheck(endpoint || 'pi', auth);
    }
}
/** Singleton instance for registry */
export const piConnector = new PiConnector();
/**
 * Create a Pi connector with the Agent Health package pre-configured
 */
export function createAgentHealthPiConnector(packagePath) {
    const args = ['--print', '--mode', 'json'];
    if (packagePath) {
        args.push('--skill', `${packagePath}/skills/*`);
        args.push('--extension', `${packagePath}/extensions/agent-health.ts`);
        args.push('--append-system-prompt', `${packagePath}/prompts/agent-health.md`);
    }
    return new PiConnector({ args });
}
//# sourceMappingURL=PiConnector.js.map