import { readFile, writeFile, access, mkdir } from 'node:fs/promises'
import path from 'node:path'
import { getRootPath } from '../lib/get-root-path'
import { z } from 'zod'

export const TestRowSchema = z.object({
	id: z.string(),
	pinned: z.enum(['TRUE', 'FALSE']),
	scenario: z.string(),
	author: z.string().optional(),
	category: z.string().optional(),
	assertion_type: z.string(),
	assertion_value: z.string(),
	conversation: z.string(),
	expected_output: z.string(),
	timeout: z.string(),
	similarity_threshold: z.string(),
	context: z.record(z.unknown())
})

export const SuiteMetaSchema = z.object({
	defaultContext: z.record(z.unknown()).optional(),
	repeats: z.number().optional(),
	concurrency: z.number().optional()
})

export const TestSuiteSchema = z.object({
	version: z.literal(1),
	meta: SuiteMetaSchema.optional(),
	tests: z.array(TestRowSchema)
})

export type TestRow = z.infer<typeof TestRowSchema>
export type SuiteMeta = z.infer<typeof SuiteMetaSchema>
export type TestSuite = z.infer<typeof TestSuiteSchema>

export async function getSuiteFilePath(agent: string): Promise<string> {
	const base = path.dirname(await getRootPath())
	return path.resolve(base, `${agent}.tests.jsonl`)
}

export async function readSuite(agent: string): Promise<TestSuite> {
	const jsonlPath = await getSuiteFilePath(agent)
	const jsonPathLegacy = jsonlPath.replace(/\.jsonl$/, '.json')
	// Prefer JSONL, fallback to legacy JSON
	try {
		await access(jsonlPath)
		const raw = await readFile(jsonlPath, 'utf8')
		const lines = raw.split(/\r?\n/).filter((l) => l.trim().length > 0)
		let meta: SuiteMeta | undefined
		const rows: TestRow[] = []
		for (const line of lines) {
			try {
				const obj = JSON.parse(line)
				if (obj && (typeof obj.version !== 'undefined' || typeof obj.meta !== 'undefined')) {
					meta = (obj.meta || {}) as SuiteMeta
					continue
				}
				rows.push(normalizeRow(obj))
			} catch {
				throw "Malformed line in test suite JSONL file: '" + line + "'"
			}
		}
		return { version: 1, meta: meta || {}, tests: rows }
	} catch {}

	try {
		await access(jsonPathLegacy)
		const raw = await readFile(jsonPathLegacy, 'utf8')
		const parsed = JSON.parse(raw)
		const tests: TestRow[] = Array.isArray(parsed?.tests)
			? parsed.tests.map((r: any) => normalizeRow(r))
			: []
		const suite: TestSuite = { version: 1, meta: parsed?.meta || {}, tests }
		return suite
	} catch (e: any) {
		throw new Error(`Test suite not found. Please create ${jsonlPath}.`)
	}
}

export async function writeSuite(agent: string, suite: TestSuite): Promise<void> {
	const jsonlPath = await getSuiteFilePath(agent)
	await mkdir(path.dirname(jsonlPath), { recursive: true })
	const header = { version: 1, meta: suite.meta || {} }
	const lines = [JSON.stringify(header), ...suite.tests.map((r) => JSON.stringify(normalizeRow(r)))]
	await writeFile(jsonlPath, lines.join('\n') + '\n', 'utf8')
}

export async function upsertSuiteMeta(agent: string, metaPatch: SuiteMeta): Promise<void> {
	const suite = await readSuite(agent).catch(() => ({ version: 1, tests: [] }) as TestSuite)
	suite.meta = { ...(suite.meta || {}), ...(metaPatch || {}) }
	await writeSuite(agent, suite)
}

export function normalizeRow(r: any): TestRow {
	return {
		id: String(r.id || randomUUID()),
		pinned: String(r.pinned || 'TRUE').toUpperCase() === 'TRUE' ? 'TRUE' : 'FALSE',
		scenario: String(r.scenario || ''),
		author: r.author != null ? String(r.author) : '',
		category: r.category != null ? String(r.category) : '',
		assertion_type: String(r.assertion_type || 'llm-rubric'),
		assertion_value: String(r.assertion_value || ''),
		conversation: String(r.conversation || ''),
		expected_output: String(r.expected_output || ''),
		timeout: String(r.timeout ?? '0'),
		similarity_threshold: String(r.similarity_threshold ?? '60'),
		context: r.context || {}
	}
}

function randomUUID() {
	const rnd = () =>
		Math.floor(Math.random() * 0xffffffff)
			.toString(16)
			.padStart(8, '0')
	return `${rnd()}-${rnd().slice(0, 4)}-${rnd().slice(0, 4)}-${rnd().slice(0, 4)}-${rnd()}${rnd()}`
}
