import OpenAI from 'openai'
import type { StepTrace } from '@recombine-ai/engine'

import { getRunResultEntry } from './data'
import { decodeHtmlEntities } from '../lib/decode-html-entities'

const DEFAULT_ANALYZE_SYSTEM_MESSAGE =
	"Stop roleplay. Analyze the last response. Respond with a JSON array of direct quotes from the system prompt that you used to generate it. Keep the quotes short if possible, especially if you're quoting a list."

const DEFAULT_ANALYZE_RATIONALE_SYSTEM_MESSAGE =
	'Stop roleplay. Analyze the last response. Extremely briefly and concisely, explain why you decieded to respond like that.'

const DEFAULT_SUGGESTED_USER_RESPONSE_PROMPT =
	'Given this conversation:\n\n{{stringifiedConversation}}\n\nPropose a realistic plausible response from the user. Be brief'

export interface QuoteCount {
	quote: string
	count: number
}

export interface RationaleAttempt {
	attempt: number
	status: 'ok' | 'error'
	text: string
}

export interface AnalyzeRunResult {
	quotes: QuoteCount[]
	rationales: RationaleAttempt[]
	parallelCalls: number
	renderedPrompt: string
	stringifiedConversation: string
	raw: Record<string, unknown>
}

export async function generateSuggestedUserResponse(input: {
	model: string | undefined
	stringifiedConversation: string
	response: string
}): Promise<{ text: string }> {
	const apiKey = process.env.OPENAI_API_KEY
	if (!apiKey) {
		throw new Error('OPENAI_API_KEY is not set')
	}

	const conversation = decodeHtmlEntities(String(input.stringifiedConversation ?? ''))
	if (!conversation) {
		throw new Error('Missing conversation')
	}

	const decodedResponse = decodeHtmlEntities(String(input.response ?? ''))
	if (!decodedResponse.trim()) {
		throw new Error('Missing response')
	}

	const resolvedModel = resolveModel(input.model)
	if (!resolvedModel) {
		throw new Error('Missing model')
	}

	const agentLabel = inferAgentLabel(conversation)
	const fullConversation = appendResponseToConversation(conversation, decodedResponse, agentLabel)
	const client = new OpenAI({ apiKey })

	const completion = await client.chat.completions.create({
		model: resolvedModel,
		messages: [{ role: 'user', content: buildSuggestedUserResponsePrompt(fullConversation) }]
	})

	const content = completion.choices?.[0]?.message?.content
	if (!content) throw new Error('No suggested user response from the LLM')

	const text = normalizeSuggestedUserResponse(content)
	if (!text) throw new Error('Suggested user response was empty')

	return { text }
}

export async function analyzeTraceQuotes(input: {
	model: string | undefined
	renderedPrompt: string
	stringifiedConversation: string
	response: string
	stepTrace?: StepTrace
	rawInput?: Record<string, unknown>
}): Promise<AnalyzeRunResult> {
	const apiKey = process.env.OPENAI_API_KEY
	if (!apiKey) {
		throw new Error('OPENAI_API_KEY is not set')
	}

	const parallelCalls = envIntOrDefault('BOSUN_ANALYZE_PARALLEL_CALLS', 5, { min: 1, max: 25 })

	const analyzeSystemMessage = envTextOrDefault(
		'BOSUN_ANALYZE_SYSTEM_MESSAGE',
		DEFAULT_ANALYZE_SYSTEM_MESSAGE
	)

	const analyzeRationaleSystemMessage = envTextOrDefault(
		'BOSUN_ANALYZE_RATIONALE_SYSTEM_MESSAGE',
		DEFAULT_ANALYZE_RATIONALE_SYSTEM_MESSAGE
	)

	const renderedPrompt = decodeHtmlEntities(String(input.renderedPrompt ?? ''))
	if (!renderedPrompt) {
		throw new Error('Missing rendered prompt')
	}

	const conversation = decodeHtmlEntities(String(input.stringifiedConversation ?? ''))
	if (!conversation) {
		throw new Error('Missing conversation')
	}

	const resolvedModel = resolveModel(input.model)
	if (!resolvedModel) {
		throw new Error('Missing model')
	}

	const agentLabel = inferAgentLabel(conversation)
	const decodedResponse = decodeHtmlEntities(String(input.response ?? ''))
	const conversationWithResponse = appendResponseToConversation(
		conversation,
		decodedResponse,
		agentLabel
	)

	const client = new OpenAI({ apiKey })

	const quoteCalls = Array.from({ length: parallelCalls }, () =>
		analyzeOnce(
			client,
			resolvedModel,
			renderedPrompt,
			conversationWithResponse,
			analyzeSystemMessage
		)
	)
	const rationaleCalls = Array.from({ length: parallelCalls }, () =>
		analyzeRationaleOnce(
			client,
			resolvedModel,
			renderedPrompt,
			conversationWithResponse,
			analyzeRationaleSystemMessage
		)
	)

	const [quoteSettled, rationaleSettled] = await Promise.all([
		Promise.allSettled(quoteCalls),
		Promise.allSettled(rationaleCalls)
	])

	const quoteAttemptDetails: Array<Record<string, unknown>> = []
	const counts = new Map<string, number>()
	let succeeded = 0
	let failed = 0
	for (let i = 0; i < quoteSettled.length; i++) {
		const res = quoteSettled[i]
		if (res.status === 'fulfilled') {
			succeeded += 1
			const uniqueQuotesThisAttempt = [
				...new Set(res.value.map((q) => normalizeQuote(String(q))).filter((q) => q.length > 0))
			]
			quoteAttemptDetails.push({ attempt: i + 1, status: 'ok', quotes: uniqueQuotesThisAttempt })
			for (const quote of uniqueQuotesThisAttempt) {
				counts.set(quote, (counts.get(quote) || 0) + 1)
			}
			continue
		}
		failed += 1
		quoteAttemptDetails.push({
			attempt: i + 1,
			status: 'error',
			error: toErrorString(res.reason)
		})
	}

	const quotes = [...counts.entries()]
		.map(([quote, count]) => ({ quote, count }))
		.sort((a, b) => b.count - a.count || a.quote.localeCompare(b.quote))

	const rationaleAttemptDetails: Array<Record<string, unknown>> = []
	const rationales: RationaleAttempt[] = []
	let rationaleSucceeded = 0
	let rationaleFailed = 0
	for (let i = 0; i < rationaleSettled.length; i++) {
		const res = rationaleSettled[i]
		if (res.status === 'fulfilled') {
			rationaleSucceeded += 1
			const text = normalizeRationale(res.value)
			rationales.push({ attempt: i + 1, status: 'ok', text })
			rationaleAttemptDetails.push({ attempt: i + 1, status: 'ok', text })
			continue
		}
		rationaleFailed += 1
		const err = toErrorString(res.reason)
		rationales.push({ attempt: i + 1, status: 'error', text: err })
		rationaleAttemptDetails.push({ attempt: i + 1, status: 'error', error: err })
	}

	if (succeeded === 0 && rationaleSucceeded === 0) {
		throw new Error('All analysis calls failed')
	}

	const sanitizedStepTrace = input.stepTrace ? sanitizeAnalyzedStepTrace(input.stepTrace) : null

	return {
		quotes,
		rationales,
		parallelCalls,
		renderedPrompt,
		stringifiedConversation: conversation,
		raw: {
			input: input.rawInput ?? null,
			quoteCalls: {
				requested: parallelCalls,
				succeeded,
				failed
			},
			rationaleCalls: {
				requested: parallelCalls,
				succeeded: rationaleSucceeded,
				failed: rationaleFailed
			},
			stepTrace: sanitizedStepTrace,
			analysis: {
				parallelCalls,
				quoteSystemMessage: analyzeSystemMessage,
				rationaleSystemMessage: analyzeRationaleSystemMessage,
				userMessage: {
					conversationChars: conversation.length,
					appendedAgentLabel: agentLabel,
					appendedResponseChars: String(input.response ?? '').length,
					finalConversationChars: conversationWithResponse.length
				},
				systemPrompt: {
					renderedPromptChars: renderedPrompt.length
				}
			},
			quoteAttempts: quoteAttemptDetails,
			rationaleAttempts: rationaleAttemptDetails
		}
	}
}

export async function analyzeRunQuotes(input: {
	agent: string
	testId: string
	index: number
	timestamp?: string
}): Promise<AnalyzeRunResult> {
	const run = await getRunResultEntry(input.agent, input.testId, input.index, input.timestamp)

	const decodedRenderedPrompt = decodeHtmlEntities(String(run.rendered_prompt ?? ''))
	if (!decodedRenderedPrompt || decodedRenderedPrompt.startsWith('NO RENDERED PROMPT')) {
		throw new Error('Missing main step rendered prompt for this run')
	}

	const decodedConversation = decodeHtmlEntities(String(run.stringified_conversation ?? ''))
	if (!decodedConversation || decodedConversation.startsWith('NO STRINGIFIED CONVERSATION')) {
		throw new Error('Missing main step conversation for this run')
	}

	const mainTrace = findMainTrace(run.traces)
	if (!mainTrace) {
		throw new Error('Main step trace not found for this run')
	}

	const resolvedModel = resolveModel(mainTrace.model)
	if (!resolvedModel) {
		throw new Error('Missing main step model for this run')
	}

	return analyzeTraceQuotes({
		model: resolvedModel,
		renderedPrompt: decodedRenderedPrompt,
		stringifiedConversation: decodedConversation,
		response: String(run.result || ''),
		stepTrace: mainTrace,
		rawInput: {
			agent: input.agent,
			testId: input.testId,
			index: input.index,
			timestamp: input.timestamp ?? null
		}
	})
}

async function analyzeOnce(
	client: OpenAI,
	model: string,
	renderedPrompt: string,
	conversationWithResponse: string,
	analyzeSystemMessage: string
): Promise<string[]> {
	const completion = await client.chat.completions.create({
		model,
		messages: [
			{ role: 'system', content: renderedPrompt },
			{ role: 'user', content: conversationWithResponse },
			{ role: 'system', content: analyzeSystemMessage }
		]
	})
	const content = completion.choices?.[0]?.message?.content
	if (!content) throw new Error('No analysis response from the LLM')
	return parseJsonArrayOfStrings(content)
}

async function analyzeRationaleOnce(
	client: OpenAI,
	model: string,
	renderedPrompt: string,
	conversationWithResponse: string,
	analyzeSystemMessage: string
): Promise<string> {
	const completion = await client.chat.completions.create({
		model,
		messages: [
			{ role: 'system', content: renderedPrompt },
			{ role: 'user', content: conversationWithResponse },
			{ role: 'system', content: analyzeSystemMessage }
		]
	})
	const content = completion.choices?.[0]?.message?.content
	if (!content) throw new Error('No rationale response from the LLM')
	return content
}

export function buildSuggestedUserResponsePrompt(stringifiedConversation: string): string {
	return DEFAULT_SUGGESTED_USER_RESPONSE_PROMPT.replace(
		'{{stringifiedConversation}}',
		stringifiedConversation
	)
}

function normalizeQuote(raw: string): string {
	let q = decodeHtmlEntities(raw).trim()
	q = trimWrappingQuotes(q)
	return q.trim()
}

export function normalizeSuggestedUserResponse(raw: string): string {
	let text = decodeHtmlEntities(String(raw ?? ''))
	text = stripCodeFences(text).trim()
	text = text.replace(/^user:\s*/i, '')
	text = trimWrappingQuotes(text)
	return text.trim()
}

function normalizeRationale(raw: string): string {
	const text = decodeHtmlEntities(String(raw ?? '')).trim()
	return stripCodeFences(text).trim()
}

function sanitizeAnalyzedStepTrace(trace: StepTrace): Record<string, unknown> {
	let clone: any
	try {
		clone = JSON.parse(JSON.stringify(trace))
	} catch {
		clone = { ...(trace as any) }
	}
	if (clone && typeof clone === 'object') {
		if ('renderedPrompt' in clone) delete clone.renderedPrompt
		if ('receivedPrompt' in clone) delete clone.receivedPrompt
	}
	return clone as Record<string, unknown>
}

function toErrorString(err: unknown): string {
	if (err instanceof Error) return `${err.name}: ${err.message}`
	try {
		return JSON.stringify(err)
	} catch {
		return String(err)
	}
}

function trimWrappingQuotes(text: string): string {
	const quotePairs: Array<[string, string]> = [
		['"', '"'],
		["'", "'"],
		['“', '”'],
		['‘', '’'],
		['«', '»'],
		['„', '“']
	]

	let out = text
	for (let i = 0; i < 5; i++) {
		const trimmed = out.trim()
		let changed = false
		for (const [open, close] of quotePairs) {
			if (trimmed.startsWith(open) && trimmed.endsWith(close)) {
				const next = trimmed.slice(open.length, trimmed.length - close.length)
				if (next === trimmed) continue
				out = next
				changed = true
				break
			}
		}
		if (!changed) {
			out = trimmed
			break
		}
	}
	return out
}

function findMainTrace(traces: StepTrace[] | undefined): StepTrace | undefined {
	if (!Array.isArray(traces) || traces.length === 0) return undefined
	for (let i = traces.length - 1; i >= 0; i--) {
		const t = traces[i]
		if (!t) continue
		if (t.name === 'main' || t.name === 'mainReply') return t
	}
	return undefined
}

function resolveModel(model: string | undefined): string | null {
	if (!model) return null
	if (model === 'default') return 'gpt-4o-2024-08-06'

	// Some traces stringify adapter options instead of a bare model name.
	const trimmed = model.trim()
	if (trimmed.startsWith('{') && trimmed.endsWith('}')) {
		try {
			const parsed = JSON.parse(trimmed) as unknown
			if (parsed && typeof parsed === 'object' && 'model' in parsed) {
				const m = (parsed as any).model
				if (typeof m === 'string' && m.trim().length > 0) return m.trim()
			}
		} catch {
			// ignore
		}
	}

	return trimmed.length > 0 ? trimmed : null
}

function appendResponseToConversation(
	conversation: string,
	response: string,
	agentLabel: string
): string {
	const base = conversation.trimEnd()
	const prefix = base.length > 0 ? base + '\n' : ''
	return `${prefix}${agentLabel}: ${response}`
}

function inferAgentLabel(conversation: string): string {
	const lines = conversation.split(/\r?\n/)
	const counts = new Map<string, number>()
	const seen: string[] = []

	for (const line of lines) {
		const m = line.match(/^([^:\n]{1,50}):\s/)
		if (!m) continue
		const label = m[1]
		seen.push(label)
		counts.set(label, (counts.get(label) || 0) + 1)
	}

	const userLabel = counts.has('User') ? 'User' : seen[0]
	const sorted = [...counts.entries()].sort((a, b) => b[1] - a[1])
	for (const [label] of sorted) {
		if (!label) continue
		if (label === userLabel) continue
		if (label.toLowerCase() === 'system') continue
		if (label.toLowerCase().startsWith('proposed')) continue
		return label
	}

	return 'Agent'
}

function parseJsonArrayOfStrings(raw: string): string[] {
	const text = stripCodeFences(raw).trim()

	const parsed = tryParseJson(text) ?? tryParseJson(extractJsonArray(text))
	if (!parsed) {
		throw new Error('Failed to parse analysis response as JSON array')
	}

	if (Array.isArray(parsed)) {
		return parsed.map((v) => String(v))
	}

	if (parsed && typeof parsed === 'object') {
		const quotes = (parsed as any).quotes
		if (Array.isArray(quotes)) return quotes.map((v: unknown) => String(v))
	}

	throw new Error('Analysis response was not a JSON array')
}

function stripCodeFences(text: string): string {
	const trimmed = text.trim()
	const m = trimmed.match(/^```(?:[a-z0-9_-]+)?\s*([\s\S]*?)\s*```$/i)
	return m ? m[1] : trimmed
}

function tryParseJson(text: string): unknown | null {
	try {
		return JSON.parse(text)
	} catch {
		return null
	}
}

function extractJsonArray(text: string): string {
	const start = text.indexOf('[')
	const end = text.lastIndexOf(']')
	if (start === -1 || end === -1 || end < start) return text
	return text.slice(start, end + 1)
}

function envTextOrDefault(key: string, defaultValue: string): string {
	const raw = (process.env[key] ?? '').trim()
	if (!raw) return defaultValue
	return raw.replace(/\\n/g, '\n')
}

function envIntOrDefault(
	key: string,
	defaultValue: number,
	opts: { min: number; max: number }
): number {
	const raw = (process.env[key] ?? '').trim()
	if (!raw) return defaultValue
	const n = Number(raw)
	if (!Number.isFinite(n)) return defaultValue
	const rounded = Math.round(n)
	if (!Number.isFinite(rounded)) return defaultValue
	return Math.max(opts.min, Math.min(opts.max, rounded))
}
