import { type AgentState } from './agents'
import fs from 'fs'
import path from 'path'

export interface StatePersistence {
	load: (agentId: string) => Promise<AgentState[]>
	save: (agentId: string, states: AgentState[]) => Promise<void>
}

export type Context = Record<string, any>

export interface ContextPersistence {
	load: (agentId: string, contextName: string) => Promise<Context>
	save: (agentId: string, contextName: string, context: Context) => Promise<void>
}

export const fsStatePersistence: StatePersistence = {
	load: async (agentId: string) => {
		const filePath = path.join('.', 'data', 'history', `${agentId}.json`)
		return fsPersistence.load(filePath)
	},
	save: async (agentId: string, states: AgentState[]) => {
		const filePath = path.join('.', 'data', 'history', `${agentId}.json`)
		return fsPersistence.save(filePath, states)
	}
}

export const fsContextPersistence: ContextPersistence = {
	load: async (agentId: string, contextName: string) => {
		const filePath = path.join('.', 'data', 'contexts', contextName, `${agentId}.json`)
		return fsPersistence.load(filePath)
	},
	save: async (agentId: string, contextName: string, context: Context) => {
		const filePath = path.join('.', 'data', 'contexts', contextName, `${agentId}.json`)
		return fsPersistence.save(filePath, context)
	}
}

const fsPersistence = {
	load: async (filePath: string) => {
		try {
			const data = String(await fs.promises.readFile(filePath))
			return data ? JSON.parse(data) : ''
		} catch (error) {
			if ((error as any).code === 'ENOENT') {
				await createFileIfNotExists(filePath)
				return ''
			}
			console.error('error loading: ', filePath, String(await fs.promises.readFile(filePath)))
			throw error
		}
	},
	// TODO replace with JSONObject | JSONArray
	save: async (filePath: string, toSave: Object) => {
		await createFileIfNotExists(filePath)
		await fs.promises.writeFile(filePath, JSON.stringify(toSave, null, 2))
	}
}

/**
 * Ensures that a file exists by creating any necessary parent directories and the file itself
 * @param filePath Path to the file, may include nested directories
 */
async function createFileIfNotExists(filePath: string) {
	const directory = path.dirname(filePath)

	try {
		await fs.promises.access(directory)
	} catch (error) {
		// Directory doesn't exist, create it recursively
		await fs.promises.mkdir(directory, { recursive: true })
	}

	try {
		await fs.promises.access(filePath)
	} catch (error) {
		// File doesn't exist, create an empty file
		await fs.promises.writeFile(filePath, '')
	}
}
