import { existsSync, mkdirSync, readFileSync, statSync, writeFileSync } from "node:fs";
import { dirname, join } from "node:path";
import { parse } from "smol-toml";
import { getAgentDir } from "../config.js";

export const RLM_MODELS_FILE_NAME = "rlm-models.toml";

export const DEFAULT_RLM_MODELS_TOML = `# ~/.prime/agent/rlm-models.toml
# Categorized models for rlm.spawn()

scouts = [
  "google-antigravity/gemini-3.8-flash-low",
  "openai-codex/gpt-5.6-luna",
  "opencode-go/deepseek-v4.1-flash",
  "deepseek/deepseek-flash",
]

coders = [
  "google-antigravity/gemini-3.8-flash-medium",
  "google-antigravity/gemini-3.8-flash-high",
  "openai-codex/gpt-5.6-sol",
  "opencode-go/deepseek-v4.1-flash",
  "deepseek/deepseek-flash",
]

researchers = [
  "google-antigravity/gemini-3.8-flash-medium",
  "openai-codex/gpt-5.6-luna",
  "openai-codex/gpt-5.6-sol",
  "opencode-go/deepseek-v4.1-flash",
  "deepseek/deepseek-flash",
]

thinkers = [
  "google-antigravity/gemini-3.8-flash-high",
  "openai-codex/gpt-5.6-sol",
  "opencode-go/deepseek-v4.1-flash",
  "deepseek/deepseek-flash",
]

reviewers = [
  "google-antigravity/gemini-3.8-flash-high",
  "openai-codex/gpt-5.6-sol",
  "opencode-go/deepseek-v4.1-flash",
  "deepseek/deepseek-flash",
]
`;

export function getRlmModelsFilePath(overridePath?: string | null): string {
	if (overridePath) return overridePath;
	return join(getAgentDir(), RLM_MODELS_FILE_NAME);
}

export class RlmModelsStore {
	private _cachedModels: Record<string, string[]> | null = null;
	private _lastMtimeMs: number = -1;
	private _filePath: string;

	constructor(filePath?: string | null) {
		this._filePath = getRlmModelsFilePath(filePath);
	}

	getModels(): Record<string, string[]> {
		this._ensureFileExists();
		const stat = statSync(this._filePath);
		if (this._cachedModels && stat.mtimeMs === this._lastMtimeMs) {
			return { ...this._cachedModels };
		}
		const content = readFileSync(this._filePath, "utf8");
		const parsed = parse(content);
		const models: Record<string, string[]> = {};
		for (const [key, value] of Object.entries(parsed)) {
			if (Array.isArray(value)) {
				models[key] = value.filter((v) => typeof v === "string" || typeof v === "number").map((v) => String(v));
			}
		}
		this._cachedModels = models;
		this._lastMtimeMs = stat.mtimeMs;
		return { ...models };
	}

	private _ensureFileExists(): void {
		if (existsSync(this._filePath)) return;
		try {
			mkdirSync(dirname(this._filePath), { recursive: true });
			writeFileSync(this._filePath, DEFAULT_RLM_MODELS_TOML, "utf8");
		} catch {
			// Best-effort write if dir cannot be created
		}
	}
}