import * as vscode from 'vscode';
import { logInfo, logError } from '../../utils';
import type { ChatMessage } from '../../agent';
import {
    extractVisibleFinal,
    shouldAutoContinue,
    mergeContinuationParts,
    buildContinuationUserPrompt,
    CONTINUATION_SYSTEM_PROMPT,
} from '../../core/responseRecovery';
import { isRestartedAnswer } from '../../lib/contextBuilders/outputSanitization';
import {
    estimateTokens,
    estimateMessagesTokens,
    computeOutputBudget,
    classifyStopReason,
    type ContextLimits,
} from '../../lib/contextManager';
import { recordTelemetry } from '../../core/telemetry';

/** Result shape of `extractVisibleFinal` — kept structural here to avoid a hard import dependency. */
type CleanedAnswer = ReturnType<typeof extractVisibleFinal>;

export interface ApplyAutoContinuationDeps {
    /** Bound `this.streamChatOnce` — same params shape as the original method. */
    streamChatOnce: (params: any) => Promise<{ text: string; stopReason?: string; aborted: boolean }>;
    /** Bound `this.isStaleRun`. */
    isStaleRun: (runId: number) => boolean;
    /** Live AbortSignal getter — controller is reassigned across turns. */
    getAbortSignal: () => AbortSignal | undefined;
    /** Webview for posting `autoContinue` UI updates. */
    getWebview: () => vscode.Webview | undefined;
}

export interface ApplyAutoContinuationInput {
    /** From `extractVisibleFinal(aiResponseText)`. The loop mutates `cleaned.visible`. */
    cleaned: CleanedAnswer;
    finishStopReason: string | undefined;
    prompt: string | null;
    chatHistory: ChatMessage[];  // used to find the original user prompt fallback
    maxOutputTokens: number;     // the first generation's output budget (carried for first round)
    ctxLimits: ContextLimits;
    config: any;                  // getConfig() — we read autoContinueOnOutputLimit, maxAutoContinuations, contextOverflowPolicy
    runId: number;
    useLmStudioSdk: boolean;
    engine: string;
    ollamaUrl: string;
    actualModel: string;
    temperature: number;
    postLiveDeltas: boolean;
}

export interface ApplyAutoContinuationResult {
    cleaned: CleanedAnswer;
    finishStopReason: string | undefined;
    continuationCount: number;
}

export async function applyAutoContinuation(
    deps: ApplyAutoContinuationDeps,
    input: ApplyAutoContinuationInput,
): Promise<ApplyAutoContinuationResult> {
    let {
        cleaned,
        finishStopReason,
        prompt,
        chatHistory,
        maxOutputTokens,
        ctxLimits,
        config,
        runId,
        useLmStudioSdk,
        engine,
        ollamaUrl,
        actualModel,
        temperature,
        postLiveDeltas,
    } = input;

    // (c) Auto-continuation — the visible answer hit the output-token ceiling.
    let continuationCount = 0;
    if (config.autoContinueOnOutputLimit && config.maxAutoContinuations > 0) {
        const originalUserPrompt = prompt || (chatHistory.find(m => m.role === 'user' && typeof m.content === 'string')?.content as string) || '';
        let lastOutputTokens = estimateTokens(cleaned.visible);
        let lastMaxOutputTokens = maxOutputTokens;   // budget the last round actually had (≠ first gen's after round 1)
        while (
            shouldAutoContinue(classifyStopReason(finishStopReason), cleaned.visible, lastOutputTokens, lastMaxOutputTokens)
            && continuationCount < config.maxAutoContinuations
            && !deps.getAbortSignal()?.aborted
            && !deps.isStaleRun(runId)
        ) {
            continuationCount++;
            const continuationStartMs = Date.now();
            deps.getWebview()?.postMessage({ type: 'autoContinue', value: `답변이 길어 이어서 정리하는 중입니다... (${continuationCount}/${config.maxAutoContinuations})` });
            try {
                const contMsgs: ChatMessage[] = [
                    { role: 'system', content: CONTINUATION_SYSTEM_PROMPT, internal: true },
                    { role: 'user', content: buildContinuationUserPrompt(originalUserPrompt, cleaned.visible) },
                ];
                lastMaxOutputTokens = computeOutputBudget(estimateMessagesTokens(contMsgs), ctxLimits).maxOutputTokens;
                // Stream the continuation through the same channel as the main turn so
                // the user sees the answer keep growing instead of freezing for 10–30s
                // while we silently call non-streaming. The trailing streamReplace
                // (after sanitize / merge) corrects any overlap the model re-emits.
                const cr = await deps.streamChatOnce({
                    runId, useLmStudioSdk, engine, ollamaUrl, modelName: actualModel,
                    messages: contMsgs,
                    temperature,
                    maxTokens: lastMaxOutputTokens,
                    contextLength: ctxLimits.contextLength,
                    contextOverflowPolicy: config.contextOverflowPolicy,
                    signal: deps.getAbortSignal()!,
                    postLiveDeltas,
                });
                if (cr.aborted) {
                    logInfo('Auto-continuation aborted mid-stream.', { model: actualModel, round: continuationCount });
                    break;
                }
                finishStopReason = cr.stopReason;
                const ccl = extractVisibleFinal(cr.text);
                if (!ccl.visible.trim()) {
                    logInfo('Continuation produced no visible text — stopping.', { model: actualModel, round: continuationCount });
                    break;
                }
                // A weak model often ignores "continue from here" and re-generates the
                // whole answer from the top. Discard such a restart instead of merging
                // it — otherwise the user gets the entire analysis twice.
                if (isRestartedAnswer(cleaned.visible, ccl.visible)) {
                    logInfo('Continuation restarted the answer instead of continuing — discarding it.', { model: actualModel, round: continuationCount });
                    break;
                }
                const before = cleaned.visible;
                cleaned = { ...cleaned, visible: mergeContinuationParts(cleaned.visible, ccl.visible), wasThoughtOnly: false };
                lastOutputTokens = estimateTokens(ccl.visible);
                logInfo('Auto-continued the answer.', { model: actualModel, round: continuationCount, addedChars: ccl.visible.length, totalChars: cleaned.visible.length, contStopReason: cr.stopReason, contMaxTokens: lastMaxOutputTokens });
                recordTelemetry({
                    kind: 'continuation',
                    durationMs: Date.now() - continuationStartMs,
                    model: actualModel, engine,
                    outputTokens: lastOutputTokens,
                    round: continuationCount,
                    stopReason: cr.stopReason,
                    note: `addedChars=${ccl.visible.length} mergedAdd=${cleaned.visible.length - before.length}`,
                });
                // Guard against a continuation that adds (almost) nothing new after dedup — stop instead of spinning.
                if (cleaned.visible.length - before.length < 20) {
                    logInfo('Continuation added negligible new text — stopping.', { model: actualModel, round: continuationCount });
                    break;
                }
            } catch (e: any) {
                logError('Auto-continuation failed.', { model: actualModel, round: continuationCount, error: e?.message ?? String(e) });
                break;
            }
        }
        if (deps.isStaleRun(runId)) return { cleaned, finishStopReason, continuationCount };
    }

    return { cleaned, finishStopReason, continuationCount };
}
