import { getConfig } from '../../config';
import { logError, resolveEngine } from '../../utils';
import { estimateMessagesTokens, computeOutputBudget } from '../../lib/contextManager';
import { lmStudioSamplingFromConfig, lmStudioRespondExtrasFromConfig } from '../../lib/contextBuilders/lmStudioSampling';
import { AGENT_PROMPTS, type AgentRole, type AgentExecutorOptions, type ChatMessage } from '../../agent';

export interface CallRoleAgentDeps {
    getAbortSignal: () => AbortSignal | undefined;
    createStreamingRequest: (params: {
        baseUrl: string;
        modelName: string;
        reqMessages: ChatMessage[];
        temperature: number;
        maxTokens?: number;
        contextLength?: number;
    }) => Promise<{ response: Response; engine: 'lmstudio' | 'ollama'; apiUrl: string }>;
    options: AgentExecutorOptions;
}

export async function callRoleAgent(deps: CallRoleAgentDeps, role: AgentRole, prompt: string, modelName: string, options: any): Promise<string> {
    const persona = AGENT_PROMPTS[role];
    const { ollamaUrl, contextLength, maxOutputTokens, contextSafetyMargin, contextOverflowPolicy } = getConfig();

    const messages: ChatMessage[] = [
        { role: 'system', content: persona },
        { role: 'user', content: prompt }
    ];
    // Dynamic output cap so input + output stays within the context window.
    const inputTokens = estimateMessagesTokens(messages);
    const { maxOutputTokens: subMaxTokens } = computeOutputBudget(inputTokens, {
        contextLength, maxOutputTokens, safetyMargin: contextSafetyMargin, minOutputTokens: 512,
    });

    const engine = resolveEngine(ollamaUrl);
    let responseText = '';

    if (engine === 'lmstudio' && deps.options.lmStudioStreamer) {
        try {
            const stream = deps.options.lmStudioStreamer.stream({
                modelName,
                messages: messages.map((m) => ({ role: m.role, content: m.content })),
                temperature: 0.3,
                maxTokens: subMaxTokens,
                contextOverflowPolicy,
                ...lmStudioSamplingFromConfig(),
                ...lmStudioRespondExtrasFromConfig(),
                signal: deps.getAbortSignal(),
            });
            let subStopReason: string | undefined;
            for await (const { token, stopReason } of stream) {
                if (token) responseText += token;
                if (stopReason) subStopReason = stopReason;
            }
            // Sub-agent answers that got cut mid-sentence corrupt the pipeline silently
            // (Planner produces a half-step, Writer can't recover). Surface a warn log so
            // the operator can raise subMaxTokens or pick a less aggressive output budget.
            if (subStopReason && /maxPredicted|context|truncat/i.test(subStopReason)) {
                logError('Sub-agent answer hit a generation limit.', {
                    role, model: modelName, stopReason: subStopReason,
                    chars: responseText.length, maxTokens: subMaxTokens,
                });
            }
            return responseText;
        } catch (err: any) {
            if (err?.name === 'AbortError' || deps.getAbortSignal()?.aborted) return responseText;
            logError('LM Studio SDK callAgent stream failed.', { role, error: err?.message ?? String(err) });
            throw err;
        }
    }

    const request = await deps.createStreamingRequest({
        baseUrl: ollamaUrl,
        modelName: modelName,
        reqMessages: messages,
        temperature: 0.3, // Use lower temperature for planning and research
        maxTokens: subMaxTokens,
        contextLength
    });

    const reader = request.response.body?.getReader();
    if (!reader) throw new Error("Agent response body is not readable.");

    const decoder = new TextDecoder();
    try {
        while (true) {
            const { done, value } = await reader.read();
            if (done) break;
            const chunk = decoder.decode(value, { stream: true });
            const lines = chunk.split('\n');
            for (const line of lines) {
                const trimmed = line.trim();
                if (!trimmed || trimmed === 'data: [DONE]') continue;
                try {
                    const json = JSON.parse(trimmed.startsWith('data: ') ? trimmed.slice(6) : trimmed);
                    const content = json.choices?.[0]?.delta?.content || json.message?.content || '';
                    responseText += content;
                } catch (e) { }
            }
        }
    } finally {
        try { reader.releaseLock(); } catch { /* already released */ }
    }
    return responseText;
}
