export-json.ts

import { Anthropic } from "@anthropic-ai/sdk"
import os from "os"
import * as path from "path"
import * as vscode from "vscode"

import type { ExtendedContentBlock } from "./export-markdown"
import { t } from "@shofer/core"
import { stringifyJsonToFile } from "../../utils/exportJsonWorker"

/**
 * JSON task trace export — produces a structured, machine-readable trace
 * of the entire LLM conversation enriched with per-request token usage,
 * cost, and tool call metadata.
 */

// ── Types ─────────────────────────────────────────────────────

export interface JsonExportCall {
	/** 1-based index of this API call within the task. */
	index: number
	/** API protocol from getApiProtocol(): "anthropic" or "openai". */
	apiProtocol?: string
	/** Model ID used for this request. */
	model?: string
	/** Input tokens consumed. */
	inputTokens: number
	/** Output tokens produced. */
	outputTokens: number
	/** Cache write tokens (prompt caching). */
	cacheWriteTokens: number
	/** Cache read tokens (prompt caching). */
	cacheReadTokens: number
	/** Estimated cost in USD. */
	costUsd: number
	/** Whether the request was cancelled mid-stream. */
	cancelled?: boolean
	/** Reason for cancellation if cancelled. */
	cancelReason?: string
	/** Error message if the stream failed. */
	streamingFailedMessage?: string
	/** The messages sent in this API request (user → assistant). */
	messages: Anthropic.Messages.MessageParam[]
	/** Tool calls extracted from the assistant response. */
	toolCalls: JsonExportToolCall[]
	/** Extended thinking / chain-of-thought content if present. */
	reasoning?: string
	/** Number of retries before this attempt (0 = first try). */
	retryAttempt?: number
	/**
	 * Whole milliseconds the request took, request open → stream end. Absent
	 * when the request never reached a stream end (still in flight when the
	 * task was persisted, or the host died holding it).
	 */
	durationMs?: number
	/**
	 * Whole milliseconds from request open to the first stream chunk of any
	 * kind. Absent when the stream produced nothing at all.
	 */
	firstChunkMs?: number
	/**
	 * Whole milliseconds the model spent reasoning, summed over every reasoning
	 * interval below. Absent means there was NO reasoning phase; it is never
	 * zero.
	 */
	thinkingMs?: number
	/**
	 * Every reasoning interval, as `[startMs, endMs]` pairs of whole
	 * milliseconds from the request's open — the same basis as `firstChunkMs`
	 * and `durationMs`. Absent means no reasoning was observed; never empty.
	 */
	reasoningIntervalsMs?: Array<[number, number]>
	/** Structured error information if this call failed. */
	error?: {
		message: string
		type?: string
		statusCode?: number
		stack?: string
	}
	/** Serialised wire-level request metadata captured before the call. */
	wireRequest?: string
	/** Present when tokens are estimated via char/4 heuristic rather than provider usage chunks. */
	_tokensEstimated?: true
}

export interface JsonExportToolCall {
	name: string
	id: string
	input: Record<string, unknown>
	result?: {
		content: unknown
		isError: boolean
	}
}

export interface JsonExportTrace {
	/** Export format version for forward compatibility. */
	version: 1
	/** Task identifier. */
	taskId: string
	/**
	 * Display title of the task (`HistoryItem.name`) — the short, curated name
	 * shown in the task list, as opposed to `task` (the full first-message
	 * description). Set via `set_task_title`, or locked by a parent that spawned
	 * this task with `new_task`'s `title`. Omitted when the task has no title.
	 */
	title?: string
	/** Human-readable task description. */
	task: string
	/** Mode slug (e.g. "code", "architect"). */
	mode?: string
	/** Timestamp when the task was created (ISO 8601). */
	createdAt: string
	/** Individual API calls in chronological order. */
	calls: JsonExportCall[]
	/** Aggregate token usage across all calls. */
	totalTokens: {
		input: number
		output: number
		cacheWrite: number
		cacheRead: number
	}
	/** Total estimated cost in USD. */
	totalCostUsd: number
	/** Number of API calls. */
	totalCalls: number
	/** Number of tool calls across all API calls. */
	totalToolCalls: number
	/**
	 * Nested traces for this task's spawned sub-tasks, in `childIds` order. Each
	 * child may in turn carry its own `subtasks`, so the field is the complete
	 * descendant tree. Present (and non-empty) only when the task spawned
	 * children.
	 */
	subtasks?: JsonExportTrace[]
}

/**
 * Parsed ui_messages.json entry for `api_req_started`.
 * Initially contains only `{apiProtocol}`; enriched later with cost/tokens.
 */
interface UiApiReqStarted {
	say: "api_req_started"
	ts: number
	text: string // JSON string — see UiApiReqStartedPayload
}

interface UiApiReqStartedPayload {
	apiProtocol?: string
	model?: string
	// Stored field names from ShoferApiReqInfo in Task.ts:
	tokensIn?: number
	tokensOut?: number
	cacheWrites?: number
	cacheReads?: number
	cost?: number
	cancelled?: boolean
	cancelReason?: string
	streamingFailedMessage?: string
	retryAttempt?: number
	/** Present only once the payload has been rewritten at stream end. */
	durationMs?: number
	/** Present only once the payload has been rewritten at stream end, and only
	 *  when the stream produced at least one chunk. */
	firstChunkMs?: number
	/** Present only when at least one reasoning interval was observed;
	 *  absence means there was none, and it is never written as zero. It is the
	 *  SUM of `reasoningIntervalsMs` below. */
	thinkingMs?: number
	/** Present only when at least one reasoning interval was observed; the
	 *  `[startMs, endMs]` pairs themselves, never an empty array. */
	reasoningIntervalsMs?: Array<[number, number]>
	error?: {
		message: string
		type?: string
		statusCode?: number
		stack?: string
	}
	wireRequest?: string
}

// ── Token estimation fallback ─────────────────────────────────

/**
 * Rough token count from raw text.  Uses the character/4 heuristic
 * common for English prose; close enough for trace diagnostics when
 * the provider does not emit `usage` chunks in streaming mode.
 *
 * @param text - Input text to estimate
 * @returns Approximate token count
 */
function estimateTokens(text: string): number {
	if (!text) return 0
	return Math.ceil(text.length / 4)
}

/**
 * Estimate tokens for a content block.
 */
function estimateBlockTokens(block: Record<string, unknown>): number {
	if (typeof block.text === "string") return estimateTokens(block.text)
	if (typeof block.content === "string") return estimateTokens(block.content)
	// tool_use: count the JSON-serialized input
	if (block.type === "tool_use" && block.input) {
		return estimateTokens(JSON.stringify(block.input))
	}
	return 0
}

/**
 * Estimate total input + output tokens for an array of `MessageParam`
 * messages using the char/4 heuristic.
 */
function estimateMessageTokens(messages: Anthropic.Messages.MessageParam[]): {
	input: number
	output: number
} {
	let input = 0
	let output = 0
	for (const msg of messages) {
		const blocks = Array.isArray(msg.content) ? msg.content : [{ text: String(msg.content ?? "") }]
		for (const block of blocks) {
			const tokens = estimateBlockTokens(block as Record<string, unknown>)
			if (msg.role === "user") input += tokens
			else output += tokens
		}
	}
	return { input, output }
}

// ── Public API ────────────────────────────────────────────────

/**
 * Build a filename for the JSON export.
 */
export function getJsonExportFileName(dateTs: number): string {
	const date = new Date(dateTs)
	const month = date.toLocaleString("en-US", { month: "short" }).toLowerCase()
	const day = date.getDate()
	const year = date.getFullYear()
	let hours = date.getHours()
	const minutes = date.getMinutes().toString().padStart(2, "0")
	const seconds = date.getSeconds().toString().padStart(2, "0")
	const ampm = hours >= 12 ? "pm" : "am"
	hours = hours % 12
	hours = hours ? hours : 12
	return `shofer_task_${month}-${day}-${year}_${hours}-${minutes}-${seconds}-${ampm}.json`
}

/**
 * Build a complete JSON trace from the task's persisted files.
 *
 * @param taskId - The task identifier.
 * @param taskText - Human-readable task description.
 * @param mode - Mode slug (e.g. "code").
 * @param createdAt - ISO 8601 creation timestamp.
 * @param apiConversationHistory - The full API conversation (Anthropic format).
 * @param uiMessages - Parsed ui_messages.json array.
 * @param options - Extras: the display `title` (`HistoryItem.name`).
 */
export function buildJsonTrace(
	taskId: string,
	taskText: string,
	mode: string | undefined,
	createdAt: string,
	apiConversationHistory: Anthropic.Messages.MessageParam[],
	uiMessages: Array<{ type: string; say?: string; ask?: string; ts: number; text?: string }>,
	options?: { title?: string },
): JsonExportTrace {
	const calls: JsonExportCall[] = []

	// Index the api_req_started entries by their position in the UI message stream.
	const apiReqStartedEntries = uiMessages.filter(
		(m) => m.type === "say" && m.say === "api_req_started",
	) as UiApiReqStarted[]

	// Walk the API conversation history and partition it by API call.
	// An API call is: user message(s) → assistant message.
	// The user messages carry tool_results from the previous turn; the
	// assistant message carries text + tool_uses for the current turn.
	let currentCallStart = 0
	let callIndex = 0

	for (let i = 0; i < apiConversationHistory.length; i++) {
		const msg = apiConversationHistory[i]

		if (msg.role === "assistant") {
			// An assistant message closes an API call.
			// Collect all messages from currentCallStart through this assistant message.
			const callMessages = apiConversationHistory.slice(currentCallStart, i + 1)

			// Find the matching api_req_started entry.
			const reqMeta = apiReqStartedEntries[callIndex]
			let payload: UiApiReqStartedPayload = {}
			if (reqMeta?.text) {
				try {
					payload = JSON.parse(reqMeta.text)
				} catch {
					/* best effort */
				}
			}

			// Extract tool calls and reasoning from the assistant content.
			const toolCalls: JsonExportToolCall[] = []
			let reasoning: string | undefined

			if (Array.isArray(msg.content)) {
				for (const rawBlock of msg.content) {
					const block = rawBlock as ExtendedContentBlock
					if (block.type === "tool_use") {
						const resultMsg = apiConversationHistory[i + 1]
						let result: { content: unknown; isError: boolean } | undefined

						if (resultMsg?.role === "user" && Array.isArray(resultMsg.content)) {
							const matchingResult = resultMsg.content.find(
								(b) => b.type === "tool_result" && b.tool_use_id === block.id,
							)
							if (matchingResult && matchingResult.type === "tool_result") {
								result = {
									content: matchingResult.content,
									isError: matchingResult.is_error ?? false,
								}
							}
						}

						toolCalls.push({
							name: block.name,
							id: block.id,
							input: (block.input as Record<string, unknown>) || {},
							result,
						})
					} else if (block.type === "reasoning") {
						if ("text" in block && typeof block.text === "string") {
							reasoning = (reasoning || "") + block.text
						}
					} else if (block.type === "thinking") {
						// Anthropic extended thinking format
						const thinkingBlock = block as { thinking: string }
						if (typeof thinkingBlock.thinking === "string") {
							reasoning = (reasoning || "") + thinkingBlock.thinking
						}
					}
				}
			}

			calls.push({
				index: callIndex + 1,
				apiProtocol: payload.apiProtocol,
				model: payload.model,
				inputTokens: payload.tokensIn ?? 0,
				outputTokens: payload.tokensOut ?? 0,
				cacheWriteTokens: payload.cacheWrites ?? 0,
				cacheReadTokens: payload.cacheReads ?? 0,
				costUsd: payload.cost ?? 0,
				cancelled: payload.cancelled,
				cancelReason: payload.cancelReason,
				streamingFailedMessage: payload.streamingFailedMessage,
				messages: callMessages,
				toolCalls,
				reasoning: reasoning || undefined,
				retryAttempt: payload.retryAttempt,
				durationMs: payload.durationMs,
				firstChunkMs: payload.firstChunkMs,
				thinkingMs: payload.thinkingMs,
				reasoningIntervalsMs: payload.reasoningIntervalsMs,
				error: payload.error,
				wireRequest: payload.wireRequest,
			})

			callIndex++
			currentCallStart = i + 1
		}
	}

	// Handle api_req_started entries that have no matching assistant messages.
	// This covers error-only tasks where the API never returned a response,
	// e.g. connection failures, rate limits, or empty streams.
	// Without this, the export would show an empty `calls[]` array.
	while (callIndex < apiReqStartedEntries.length) {
		const reqMeta = apiReqStartedEntries[callIndex]
		let payload: UiApiReqStartedPayload = {}
		if (reqMeta?.text) {
			try {
				payload = JSON.parse(reqMeta.text)
			} catch {
				/* best effort */
			}
		}

		calls.push({
			index: callIndex + 1,
			apiProtocol: payload.apiProtocol,
			model: payload.model,
			inputTokens: payload.tokensIn ?? 0,
			outputTokens: payload.tokensOut ?? 0,
			cacheWriteTokens: payload.cacheWrites ?? 0,
			cacheReadTokens: payload.cacheReads ?? 0,
			costUsd: payload.cost ?? 0,
			cancelled: payload.cancelled,
			cancelReason: payload.cancelReason,
			streamingFailedMessage: payload.streamingFailedMessage,
			messages: [],
			toolCalls: [],
			retryAttempt: payload.retryAttempt,
			durationMs: payload.durationMs,
			firstChunkMs: payload.firstChunkMs,
			thinkingMs: payload.thinkingMs,
			reasoningIntervalsMs: payload.reasoningIntervalsMs,
			error: payload.error,
			wireRequest: payload.wireRequest,
		})

		callIndex++
	}

	// If the provider did not emit `usage` chunks (common with
	// streaming-only providers), fall back to char/4 heuristic so the
	// trace still carries useful token estimates.
	const allZeroTokens = calls.length > 0 && calls.every((c) => c.inputTokens === 0 && c.outputTokens === 0)
	if (allZeroTokens) {
		for (const call of calls) {
			const est = estimateMessageTokens(call.messages)
			call.inputTokens = est.input
			call.outputTokens = est.output
			// Mark as estimated so consumers can distinguish from real values.
			call._tokensEstimated = true
		}
	}

	// Compute aggregates.
	let totalInput = 0
	let totalOutput = 0
	let totalCacheWrite = 0
	let totalCacheRead = 0
	let totalCost = 0
	let totalToolCalls = 0

	for (const call of calls) {
		totalInput += call.inputTokens
		totalOutput += call.outputTokens
		totalCacheWrite += call.cacheWriteTokens
		totalCacheRead += call.cacheReadTokens
		totalCost += call.costUsd
		totalToolCalls += call.toolCalls.length
	}

	return {
		version: 1,
		taskId,
		...(options?.title ? { title: options.title } : {}),
		task: taskText,
		mode,
		createdAt,
		calls,
		totalTokens: {
			input: totalInput,
			output: totalOutput,
			cacheWrite: totalCacheWrite,
			cacheRead: totalCacheRead,
		},
		totalCostUsd: totalCost,
		totalCalls: calls.length,
		totalToolCalls,
	}
}

/**
 * Build a task's trace plus the full descendant tree of its spawned sub-tasks.
 *
 * `loadTask(id)` returns the single-task trace and that task's direct `childIds`
 * (the persisted parent→child links). Each
 * child is loaded and walked the same way, so the result carries the complete
 * descendant tree under nested `subtasks` (in `childIds` order).
 *
 * `visited` is seeded with the root and grows as the walk proceeds, so a
 * malformed/looping `childIds` chain can neither recurse forever nor duplicate a
 * task across the export. A child whose `loadTask` throws (already deleted /
 * corrupt history) is skipped via `onSkip` rather than aborting the whole export.
 *
 * The I/O lives entirely in the injected `loadTask`, keeping this walker pure and
 * unit-testable. Because a large tree means many sequential `loadTask` reads,
 * `onProgress` fires after each node (with a running count) so callers can drive a
 * progress UI, and `isCancelled` is polled before each read so the walk can be
 * abandoned promptly — on cancellation it returns the partial tree built so far
 * (the caller is expected to discard it).
 */
export interface BuildJsonTraceTreeOptions {
	/** Called when a child can't be loaded (deleted / corrupt); the walk continues. */
	onSkip?: (id: string, error: unknown) => void
	/** Called after each node is loaded, with the cumulative number of nodes exported. */
	onProgress?: (id: string, exportedCount: number) => void
	/** Polled before each node load; returning true abandons the walk (partial tree). */
	isCancelled?: () => boolean
}

export async function buildJsonTraceTree(
	rootId: string,
	loadTask: (id: string) => Promise<{ trace: JsonExportTrace; childIds: string[] }>,
	options: BuildJsonTraceTreeOptions = {},
): Promise<JsonExportTrace> {
	const { onSkip, onProgress, isCancelled } = options
	const visited = new Set<string>([rootId])
	let exportedCount = 0

	const walk = async (id: string): Promise<JsonExportTrace> => {
		const { trace, childIds } = await loadTask(id)
		exportedCount++
		onProgress?.(id, exportedCount)

		const subtasks: JsonExportTrace[] = []
		for (const childId of childIds) {
			// Stop descending as soon as the caller cancels; the partial tree is
			// returned and discarded upstream.
			if (isCancelled?.()) {
				break
			}
			if (visited.has(childId)) {
				continue
			}
			visited.add(childId)
			try {
				subtasks.push(await walk(childId))
			} catch (error) {
				onSkip?.(childId, error)
			}
		}
		if (subtasks.length > 0) {
			trace.subtasks = subtasks
		}
		return trace
	}

	return walk(rootId)
}

/**
 * Prompt the user for a save location and write the JSON trace to disk.
 *
 * @param dateTs - Task creation timestamp (for filename).
 * @param trace - The built JSON export trace.
 * @param defaultUri - Default save URI.
 * @returns The URI of the saved file, or undefined if the user cancelled.
 */
/**
 * Above this serialized size, auto-opening the export in an editor tab would
 * make VS Code tokenize/fold a multi-MB document on the UI thread — itself a
 * freeze — so we surface an explicit open action instead.
 */
const LARGE_EXPORT_BYTES = 5 * 1024 * 1024

export async function downloadJsonTask(
	_dateTs: number,
	trace: JsonExportTrace,
	defaultUri: vscode.Uri,
): Promise<vscode.Uri | undefined> {
	const saveUri = await vscode.window.showSaveDialog({
		filters: { JSON: ["json"] },
		defaultUri,
	})

	if (!saveUri) {
		return undefined
	}

	// Serialize + write off the extension-host thread. A trace is the full
	// descendant task tree (every sub-task's history), so doing this inline
	// would block the event loop for seconds and freeze the webview. The progress
	// notification is indeterminate — JSON.stringify is atomic — but it now
	// actually animates, because the heavy work no longer holds the main thread.
	const bytes = await vscode.window.withProgress(
		{
			location: vscode.ProgressLocation.Notification,
			title: t("common:info.exporting_task_json_writing"),
		},
		() => stringifyJsonToFile(trace, saveUri.fsPath),
	)

	if (bytes <= LARGE_EXPORT_BYTES) {
		// Small enough to open instantly — auto-open for convenience.
		vscode.window.showTextDocument(saveUri, { preview: true })
	} else {
		// Opening a multi-MB JSON document makes VS Code tokenize/fold it on the
		// UI thread — itself a freeze — so offer an explicit open action instead.
		const open = t("common:info.export_task_json_open")
		const reveal = t("common:info.export_task_json_reveal")
		void vscode.window
			.showInformationMessage(
				t("common:info.export_task_json_large", { mb: (bytes / (1024 * 1024)).toFixed(1) }),
				open,
				reveal,
			)
			.then((choice) => {
				if (choice === open) {
					vscode.window.showTextDocument(saveUri, { preview: false })
				} else if (choice === reveal) {
					void vscode.commands.executeCommand("revealFileInOS", saveUri)
				}
			})
	}

	return saveUri
}