import { kody } from 'kody:runtime'
import { generateText, type ChatMessage, type ToolCall } from './complete.ts'
import { kodyAgentToolDefinitions } from './model-step.ts'
import { clean, stringify } from './validation.ts'

const defaultMaxSteps = 8
const defaultSearchLimit = 15

export type AgentTurnMessage = ChatMessage
export type AgentTurnInput = {
	messages: Array<AgentTurnMessage>
	system?: string
	sessionId?: string
	maxSteps?: number
	model?: string
	modelId?: string
	provider?: string
	conversationId?: string
	memoryContext?: {
		task?: string
		query?: string
		entities?: string[]
		constraints?: string[]
	}
	dryRun?: boolean
}

export type AgentToolTrace = {
	id: string
	toolName: string
	input: unknown
	output?: unknown
	error?: string
}

export type AgentTurnResult = {
	assistantText: string
	reasoningText?: string
	summary?: string | null
	continueRecommended?: boolean
	needsUserInput?: boolean
	stepsUsed?: number
	newInformation?: boolean
	stopReason?: string
	finishReason?: string
	toolCalls?: Array<AgentToolTrace>
	conversationId?: string
	text?: string
	response?: string
	error?: string
}

export type AgentTurnStreamEvent =
	| { type: 'assistant_delta'; text: string }
	| { type: 'reasoning_delta'; text: string }
	| { type: 'tool_call_started'; id: string; toolName: string; input: unknown }
	| {
			type: 'tool_call_finished'
			id: string
			toolName: string
			input: unknown
			output?: unknown
			error?: string
	  }
	| ({ type: 'turn_complete' } & AgentTurnResult)
	| { type: 'error'; message: string; phase: string }

export type AgentTurnOutput = {
	ok: boolean
	result?: AgentTurnResult
	error?: string
	events?: Array<AgentTurnStreamEvent>
	preview?: unknown
}

const knownToolNames = new Set(['search', 'execute'])

function parseArguments(value: unknown): Record<string, unknown> {
	if (!value) return {}
	if (typeof value === 'object' && !Array.isArray(value)) {
		return value as Record<string, unknown>
	}
	if (typeof value !== 'string') return {}
	try {
		const parsed = JSON.parse(value)
		return parsed && typeof parsed === 'object' && !Array.isArray(parsed)
			? (parsed as Record<string, unknown>)
			: {}
	} catch {
		return {}
	}
}

function stripToolCallPrefixes(text: string) {
	return clean(text)
		.replace(/^```(?:json)?\s*/i, '')
		.replace(/\s*```$/i, '')
		.replace(/^\[?\s*TOOL_CALLS?\s*\]?\s*/i, '')
		.replace(/^TOOL_CALLS?\s*[:=]\s*/i, '')
		.trim()
}

function extractBalancedJsonSlice(text: string, openChar: '[' | '{') {
	const closeChar = openChar === '[' ? ']' : '}'
	const start = text.indexOf(openChar)
	if (start < 0) return null
	let depth = 0
	let inString = false
	let escaped = false
	for (let i = start; i < text.length; i += 1) {
		const ch = text[i]
		if (inString) {
			if (escaped) escaped = false
			else if (ch === '\\') escaped = true
			else if (ch === '"') inString = false
			continue
		}
		if (ch === '"') {
			inString = true
			continue
		}
		if (ch === openChar) depth += 1
		else if (ch === closeChar) {
			depth -= 1
			if (depth === 0) return text.slice(start, i + 1)
		}
	}
	return null
}

function coerceToolCallItems(parsed: unknown): Array<Record<string, unknown>> | null {
	if (Array.isArray(parsed)) {
		if (parsed.length === 0) return null
		if (!parsed.every((item) => item && typeof item === 'object' && !Array.isArray(item))) {
			return null
		}
		return parsed as Array<Record<string, unknown>>
	}
	if (parsed && typeof parsed === 'object') {
		const record = parsed as Record<string, unknown>
		if (Array.isArray(record.tool_calls)) return coerceToolCallItems(record.tool_calls)
		if (Array.isArray(record.tools)) return coerceToolCallItems(record.tools)
		return [record]
	}
	return null
}

function parseTextToolCalls(text: string): Array<ToolCall> {
	const normalized = stripToolCallPrefixes(text)
	const candidates = new Set<string>()
	if (normalized) candidates.add(normalized)
	const arraySlice = extractBalancedJsonSlice(normalized, '[')
	if (arraySlice) candidates.add(arraySlice)
	const objectSlice = extractBalancedJsonSlice(normalized, '{')
	if (objectSlice) candidates.add(objectSlice)
	for (const candidate of candidates) {
		if (!(candidate.startsWith('[') || candidate.startsWith('{'))) continue
		let parsed: unknown
		try {
			parsed = JSON.parse(candidate)
		} catch {
			continue
		}
		const items = coerceToolCallItems(parsed)
		if (!items) continue
		const calls: Array<ToolCall> = []
		let valid = true
		for (let index = 0; index < items.length; index += 1) {
			const record = items[index]
			const name =
				typeof record.name === 'string'
					? record.name
					: typeof record.toolName === 'string'
						? record.toolName
						: ''
			if (!knownToolNames.has(name)) {
				valid = false
				break
			}
			const argsRaw = record.arguments ?? record.input ?? record.args ?? {}
			const input =
				typeof argsRaw === 'string'
					? parseArguments(argsRaw)
					: argsRaw && typeof argsRaw === 'object' && !Array.isArray(argsRaw)
						? (argsRaw as Record<string, unknown>)
						: {}
			if (name === 'search' && typeof input.query !== 'string') {
				valid = false
				break
			}
			if (name === 'execute' && typeof input.code !== 'string') {
				valid = false
				break
			}
			calls.push({
				toolCallId:
					typeof record.id === 'string' && record.id
						? record.id
						: 'text-tool-call-' + index,
				toolName: name,
				input,
			})
		}
		if (valid && calls.length > 0) return calls
	}
	return []
}

function unwrapMcpToolPayload(output: unknown): unknown {
	if (typeof output === 'string') {
		const trimmed = output.trim()
		if (trimmed.startsWith('{') || trimmed.startsWith('[')) {
			try {
				return unwrapMcpToolPayload(JSON.parse(trimmed))
			} catch {
				return output
			}
		}
		return output
	}
	if (!output || typeof output !== 'object' || Array.isArray(output)) return output
	const record = output as Record<string, unknown>
	if (Array.isArray(record.matches)) return output
	if (record.kind === 'entity' && (record.entityRef || record.id)) return output
	const envelopeHints = 'timing' in record || 'returnedBytes' in record || 'logs' in record
	if (!envelopeHints) return output
	if ('result' in record) {
		const inner = record.result
		if (inner && typeof inner === 'object' && !Array.isArray(inner)) {
			const innerRecord = inner as Record<string, unknown>
			return {
				...innerRecord,
				conversationId: record.conversationId ?? innerRecord.conversationId ?? null,
				error: typeof record.error === 'string' ? record.error : innerRecord.error,
			}
		}
		return {
			result: inner ?? null,
			error: typeof record.error === 'string' ? record.error : null,
			conversationId: record.conversationId ?? null,
		}
	}
	if ('error' in record) {
		return {
			error: record.error ?? null,
			conversationId: record.conversationId ?? null,
		}
	}
	return output
}

function projectMatch(match: unknown) {
	const item = match && typeof match === 'object' ? (match as Record<string, unknown>) : {}
	return {
		kind: item.kind ?? null,
		type: item.type ?? null,
		id: item.id ?? null,
		entityRef: item.entityRef ?? null,
		title: item.title ?? item.name ?? null,
		description: typeof item.description === 'string' ? item.description.slice(0, 240) : null,
		usage: typeof item.usage === 'string' ? item.usage.slice(0, 280) : null,
		executeExample:
			typeof item.executeExample === 'string' ? item.executeExample.slice(0, 500) : null,
	}
}

/** Compact search payloads for weaker models / host tool loops. */
export function projectSearchOutput(output: unknown) {
	const unwrapped = unwrapMcpToolPayload(output)
	const record =
		unwrapped && typeof unwrapped === 'object' ? (unwrapped as Record<string, unknown>) : {}
	const nested =
		record.result && typeof record.result === 'object' && !Array.isArray(record.result)
			? (record.result as Record<string, unknown>)
			: null
	let matches: unknown[] = Array.isArray(record.matches)
		? record.matches
		: Array.isArray(nested?.matches)
			? nested.matches
			: []
	if (
		matches.length === 0 &&
		(record.kind === 'entity' || nested?.kind === 'entity') &&
		(record.entityRef || record.id || nested?.entityRef || nested?.id)
	) {
		matches = [nested?.kind === 'entity' ? nested : record]
	}
	const guidance =
		typeof record.guidance === 'string'
			? record.guidance
			: typeof nested?.guidance === 'string'
				? nested.guidance
				: null
	return {
		conversationId: record.conversationId ?? nested?.conversationId ?? null,
		guidance: guidance ? guidance.slice(0, 800) : null,
		matchCount: matches.length,
		matches: matches.slice(0, 8).map(projectMatch),
	}
}

/** Compact execute payloads for weaker models / host tool loops. */
export function projectExecuteOutput(output: unknown) {
	const unwrapped = unwrapMcpToolPayload(output)
	if (unwrapped == null) return null
	if (typeof unwrapped === 'string') return unwrapped.slice(0, 4000)
	if (typeof unwrapped !== 'object') return unwrapped
	const record = unwrapped as Record<string, unknown>
	const envelopeKeys = new Set(['result', 'conversationId', 'logs', 'returnedBytes', 'timing'])
	const looksLikeExecuteEnvelope =
		'result' in record &&
		!('matches' in record) &&
		record.kind !== 'entity' &&
		!('entityRef' in record) &&
		Object.keys(record).every((key) => envelopeKeys.has(key))
	const value = looksLikeExecuteEnvelope
		? { result: record.result, conversationId: record.conversationId ?? null }
		: unwrapped
	try {
		const serialized = JSON.stringify(value)
		if (serialized.length <= 4000) return value
		return { truncated: true, preview: serialized.slice(0, 4000) }
	} catch {
		return { error: 'Could not serialize execute output.' }
	}
}

async function callTool({
	toolName: name,
	args,
	conversationId,
	memoryContext,
}: {
	toolName: string
	args: Record<string, unknown>
	conversationId: string
	memoryContext?: AgentTurnInput['memoryContext']
}) {
	if (name === 'search') {
		const query = clean(args.query)
		const entity = args.entity
		if (!query && entity == null) throw new Error('search requires query or entity.')
		const payload: Record<string, unknown> = {
			limit:
				typeof args.limit === 'number'
					? Math.min(Math.max(Math.floor(args.limit), 1), 50)
					: defaultSearchLimit,
			conversationId,
			memoryContext,
		}
		if (query) payload.query = query
		if (entity != null) payload.entity = entity
		return await kody.search(payload as Parameters<typeof kody.search>[0])
	}
	if (name === 'execute') {
		const code = clean(args.code)
		if (!code) throw new Error('execute requires code.')
		return await kody.execute({
			code,
			params: args.params && typeof args.params === 'object' ? args.params : undefined,
			conversationId,
		})
	}
	throw new Error('Unknown Kody agent tool: ' + name)
}

function inferNeedsUserInput(text: string) {
	const lower = text.toLowerCase()
	return (
		lower.includes('?') &&
		(lower.includes('could you') ||
			lower.includes('can you clarify') ||
			lower.includes('what do you mean') ||
			lower.includes('which part'))
	)
}

/**
 * Run one tool-using agent turn via the configured provider.
 * Tools call kody.search/execute directly in this package runtime.
 *
 * @example
 * import runAgentTurn from 'kody:@kody/ai/turn'
 * const preview = await runAgentTurn({
 *   messages: [{ role: 'user', content: 'Search for weather helpers' }],
 *   dryRun: true,
 * })
 */
export async function runAgentTurn(
	input: AgentTurnInput = { messages: [] },
): Promise<AgentTurnOutput> {
	try {
		if (!Array.isArray(input.messages) || input.messages.length === 0) {
			throw new Error('messages must include at least one message.')
		}
		if (input.dryRun === true) {
			const preview = await generateText({
				...input,
				system: typeof input.system === 'string' ? input.system : undefined,
				tools: kodyAgentToolDefinitions(),
				dryRun: true,
			})
			return {
				ok: true,
				result: { assistantText: '', toolCalls: [] },
				events: [],
				preview,
			}
		}
		const conversationId = input.conversationId || crypto.randomUUID()
		const maxSteps = Math.max(1, Math.min(Number(input.maxSteps) || defaultMaxSteps, 50))
		const messages: Array<ChatMessage> = []
		if (input.system) messages.push({ role: 'system', content: input.system })
		for (const message of input.messages) {
			messages.push({
				role: message.role || 'user',
				content: String(message.content ?? ''),
			})
		}
		const events: Array<AgentTurnStreamEvent> = []
		const toolCalls: Array<AgentToolTrace> = []
		let assistantText = ''
		let reasoningText = ''
		let finishReason = 'stop'
		let stepsUsed = 0

		for (let step = 0; step < maxSteps; step += 1) {
			const result = await generateText({
				...input,
				messages,
				tools: kodyAgentToolDefinitions(),
			})
			if ('dryRun' in result) {
				return { ok: true, result: { assistantText: '', conversationId }, events: [] }
			}
			if (!result.ok) throw new Error(result.error || 'Model step failed.')
			stepsUsed = step + 1
			finishReason = result.finishReason
			const structuredCalls = result.toolCalls
			const recoveredCalls =
				structuredCalls.length === 0 && result.text ? parseTextToolCalls(result.text) : []
			const calls = structuredCalls.length > 0 ? structuredCalls : recoveredCalls
			if (!calls || calls.length === 0) {
				if (result.text) {
					assistantText += result.text
					events.push({ type: 'assistant_delta', text: result.text })
				}
				if (result.reasoningText) reasoningText += result.reasoningText
				break
			}
			messages.push({
				role: 'assistant',
				content: recoveredCalls.length > 0 ? null : result.text || null,
				tool_calls: calls.map((call) => ({
					id: call.toolCallId,
					type: 'function',
					function: { name: call.toolName, arguments: stringify(call.input) },
				})),
			})
			for (const call of calls) {
				const trace: AgentToolTrace = {
					id: call.toolCallId,
					toolName: call.toolName,
					input: call.input,
				}
				events.push({
					type: 'tool_call_started',
					id: call.toolCallId,
					toolName: call.toolName,
					input: call.input,
				})
				try {
					const output = await callTool({
						toolName: call.toolName,
						args: call.input,
						conversationId,
						memoryContext: input.memoryContext,
					})
					trace.output =
						call.toolName === 'search'
							? projectSearchOutput(output)
							: projectExecuteOutput(output)
				} catch (error) {
					trace.error = error instanceof Error ? error.message : String(error)
				}
				toolCalls.push(trace)
				events.push({
					type: 'tool_call_finished',
					id: call.toolCallId,
					toolName: call.toolName,
					input: call.input,
					output: trace.output,
					error: trace.error,
				})
				messages.push({
					role: 'tool',
					tool_call_id: call.toolCallId,
					name: call.toolName,
					content: stringify(trace.output ?? { error: trace.error ?? null }).slice(0, 8000),
				})
			}
		}

		const text = clean(assistantText)
		const turnResult: AgentTurnResult = {
			assistantText: text,
			reasoningText: clean(reasoningText),
			summary: null,
			continueRecommended: false,
			needsUserInput: inferNeedsUserInput(text),
			stepsUsed,
			newInformation: true,
			stopReason: inferNeedsUserInput(text)
				? 'needs_user'
				: toolCalls.some((call) => call.error)
					? 'tool_error'
					: finishReason === 'tool_calls' || stepsUsed >= maxSteps
						? 'budget_exhausted'
						: 'completed',
			finishReason,
			toolCalls,
			conversationId,
		}
		events.push({ type: 'turn_complete', ...turnResult })
		return { ok: true, result: turnResult, events }
	} catch (error) {
		return { ok: false, error: error instanceof Error ? error.message : String(error) }
	}
}

export default runAgentTurn