/**
 * Trajectory Recorder — Step-by-step execution trace for SF units.
 *
 * Purpose: Provide RA.Aid-style trajectory recording with per-step
 * tool execution details, costs, tokens, errors, and results.
 *
 * Consumer: research units, planning units, execute-task units, /trajectory command.
 *
 * Design:
 * - SQLite-backed with JSONB-like flexibility for tool params/results
 * - Session + unit scoped
 * - Cost and token tracking per step
 * - Error recording with full context
 * - Exportable for analysis and debugging
 */

import { debugLog } from "./debug-logger.js";
import { isDbAvailable } from "./sf-db.js";
import { logWarning } from "./workflow-logger.js";

const TRAJECTORY_TABLE = "sf_trajectory";

/**
 * Trajectory step types.
 */
export const STEP_TYPES = {
	TOOL_CALL: "tool_call",
	TOOL_RESULT: "tool_result",
	LLM_CALL: "llm_call",
	LLM_RESPONSE: "llm_response",
	USER_INPUT: "user_input",
	STAGE_TRANSITION: "stage_transition",
	ERROR: "error",
	CHECKPOINT: "checkpoint",
};

/**
 * Ensure the trajectory table exists.
 */
export function ensureTrajectoryTable(db) {
	if (!db) return false;
	try {
		db.exec(`
			CREATE TABLE IF NOT EXISTS ${TRAJECTORY_TABLE} (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				session_id TEXT NOT NULL,
				unit_id TEXT,
				step_type TEXT NOT NULL CHECK(step_type IN ('tool_call', 'tool_result', 'llm_call', 'llm_response', 'user_input', 'stage_transition', 'error', 'checkpoint')),
				step_number INTEGER NOT NULL,
				tool_name TEXT,
				tool_params TEXT, -- JSON
				tool_result TEXT, -- JSON or text
				llm_model TEXT,
				llm_prompt TEXT,
				llm_response TEXT,
				input_tokens INTEGER,
				output_tokens INTEGER,
				step_cost REAL,
				is_error INTEGER NOT NULL DEFAULT 0,
				error_message TEXT,
				error_type TEXT,
				error_details TEXT,
				metadata TEXT, -- JSON: {duration_ms, retry_count, etc}
				created_at TEXT NOT NULL DEFAULT (datetime('now'))
			);
			CREATE INDEX IF NOT EXISTS idx_trajectory_session ON ${TRAJECTORY_TABLE}(session_id);
			CREATE INDEX IF NOT EXISTS idx_trajectory_unit ON ${TRAJECTORY_TABLE}(unit_id);
			CREATE INDEX IF NOT EXISTS idx_trajectory_type ON ${TRAJECTORY_TABLE}(step_type);
			CREATE INDEX IF NOT EXISTS idx_trajectory_created ON ${TRAJECTORY_TABLE}(created_at DESC);
		`);
		return true;
	} catch (err) {
		logWarning("trajectory", "Failed to ensure trajectory table", {
			error: String(err),
		});
		return false;
	}
}

const stepCounter = new Map(); // sessionId -> current step number

function getNextStepNumber(sessionId) {
	const current = stepCounter.get(sessionId) || 0;
	const next = current + 1;
	stepCounter.set(sessionId, next);
	return next;
}

function resetStepCounter(sessionId) {
	stepCounter.set(sessionId, 0);
}

/**
 * Record a trajectory step.
 *
 * @param {object} params
 * @param {string} params.sessionId — required
 * @param {string} [params.unitId] — optional unit tag
 * @param {string} params.stepType — STEP_TYPES value
 * @param {string} [params.toolName] — tool name for tool_call/tool_result
 * @param {object} [params.toolParams] — tool parameters
 * @param {any} [params.toolResult] — tool result
 * @param {string} [params.llmModel] — model identifier
 * @param {string} [params.llmPrompt] — prompt text
 * @param {string} [params.llmResponse] — response text
 * @param {number} [params.inputTokens] — input token count
 * @param {number} [params.outputTokens] — output token count
 * @param {number} [params.stepCost] — cost in USD
 * @param {boolean} [params.isError=false] — whether this step errored
 * @param {string} [params.errorMessage] — error message
 * @param {string} [params.errorType] — error type/class
 * @param {string} [params.errorDetails] — detailed error info
 * @param {object} [params.metadata] — additional metadata
 * @param {object} [params.db] — database connection
 * @returns {{id: number}|null}
 */
export function recordTrajectoryStep({
	sessionId,
	unitId,
	stepType,
	toolName,
	toolParams,
	toolResult,
	llmModel,
	llmPrompt,
	llmResponse,
	inputTokens,
	outputTokens,
	stepCost,
	isError = false,
	errorMessage,
	errorType,
	errorDetails,
	metadata,
	db,
}) {
	if (!sessionId || !stepType) {
		logWarning("trajectory", "recordTrajectoryStep missing required fields", {
			sessionId,
			stepType,
		});
		return null;
	}
	if (!Object.values(STEP_TYPES).includes(stepType)) {
		logWarning("trajectory", "Invalid step type", { stepType });
		return null;
	}

	const dbConn = db || (isDbAvailable() ? require("./sf-db.js").getDb() : null);
	if (!dbConn) return null;

	ensureTrajectoryTable(dbConn);

	const stepNumber = getNextStepNumber(sessionId);

	try {
		const result = dbConn
			.prepare(
				`INSERT INTO ${TRAJECTORY_TABLE} (
				session_id, unit_id, step_type, step_number, tool_name,
				tool_params, tool_result, llm_model, llm_prompt, llm_response,
				input_tokens, output_tokens, step_cost, is_error,
				error_message, error_type, error_details, metadata
			) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
			)
			.run(
				sessionId,
				unitId || null,
				stepType,
				stepNumber,
				toolName || null,
				toolParams ? JSON.stringify(toolParams) : null,
				toolResult !== undefined
					? typeof toolResult === "string"
						? toolResult
						: JSON.stringify(toolResult)
					: null,
				llmModel || null,
				llmPrompt || null,
				llmResponse || null,
				inputTokens || null,
				outputTokens || null,
				stepCost || null,
				isError ? 1 : 0,
				errorMessage || null,
				errorType || null,
				errorDetails || null,
				metadata ? JSON.stringify(metadata) : null,
			);

		return { id: Number(result.lastInsertRowid), stepNumber };
	} catch (err) {
		debugLog("trajectory", `record step failed: ${stepType} #${stepNumber}`, {
			error: String(err),
		});
		logWarning("trajectory", "recordTrajectoryStep failed", {
			error: String(err),
			sessionId,
			stepType,
		});
		return null;
	}
}

/**
 * Get trajectory steps with optional filtering.
 */
export function getTrajectory({
	sessionId,
	unitId,
	stepType,
	limit = 100,
	db,
}) {
	const dbConn = db || (isDbAvailable() ? require("./sf-db.js").getDb() : null);
	if (!dbConn) return [];

	ensureTrajectoryTable(dbConn);

	try {
		const conditions = ["session_id = ?"];
		const params = [sessionId];

		if (unitId) {
			conditions.push("unit_id = ?");
			params.push(unitId);
		}
		if (stepType) {
			conditions.push("step_type = ?");
			params.push(stepType);
		}

		const whereClause = conditions.join(" AND ");
		params.push(limit);

		const rows = dbConn
			.prepare(
				`SELECT * FROM ${TRAJECTORY_TABLE}
			 WHERE ${whereClause}
			 ORDER BY step_number ASC
			 LIMIT ?`,
			)
			.all(...params);

		return rows.map((r) => ({
			id: r.id,
			sessionId: r.session_id,
			unitId: r.unit_id,
			stepType: r.step_type,
			stepNumber: r.step_number,
			toolName: r.tool_name,
			toolParams: r.tool_params ? JSON.parse(r.tool_params) : null,
			toolResult: r.tool_result ? tryParseJson(r.tool_result) : null,
			llmModel: r.llm_model,
			llmPrompt: r.llm_prompt,
			llmResponse: r.llm_response,
			inputTokens: r.input_tokens,
			outputTokens: r.output_tokens,
			stepCost: r.step_cost,
			isError: !!r.is_error,
			errorMessage: r.error_message,
			errorType: r.error_type,
			errorDetails: r.error_details,
			metadata: r.metadata ? JSON.parse(r.metadata) : null,
			createdAt: r.created_at,
		}));
	} catch (err) {
		logWarning("trajectory", "getTrajectory failed", {
			error: String(err),
			sessionId,
		});
		return [];
	}
}

function tryParseJson(str) {
	try {
		return JSON.parse(str);
	} catch {
		return str;
	}
}

/**
 * Get trajectory summary for a session/unit.
 */
export function getTrajectorySummary(sessionId, unitId, db) {
	const dbConn = db || (isDbAvailable() ? require("./sf-db.js").getDb() : null);
	if (!dbConn) return null;

	ensureTrajectoryTable(dbConn);

	try {
		const conditions = ["session_id = ?"];
		const params = [sessionId];

		if (unitId) {
			conditions.push("unit_id = ?");
			params.push(unitId);
		}

		const whereClause = conditions.join(" AND ");

		const stats = dbConn
			.prepare(
				`SELECT
				COUNT(*) as total_steps,
				SUM(CASE WHEN is_error = 1 THEN 1 ELSE 0 END) as error_count,
				SUM(input_tokens) as total_input_tokens,
				SUM(output_tokens) as total_output_tokens,
				SUM(step_cost) as total_cost,
				COUNT(DISTINCT tool_name) as unique_tools,
				MAX(step_number) as max_step
			 FROM ${TRAJECTORY_TABLE}
			 WHERE ${whereClause}`,
			)
			.get(...params);

		return {
			totalSteps: stats.total_steps,
			errorCount: stats.error_count,
			totalInputTokens: stats.total_input_tokens || 0,
			totalOutputTokens: stats.total_output_tokens || 0,
			totalCost: stats.total_cost || 0,
			uniqueTools: stats.unique_tools,
			maxStep: stats.max_step,
		};
	} catch (err) {
		logWarning("trajectory", "getTrajectorySummary failed", {
			error: String(err),
		});
		return null;
	}
}

/**
 * Format trajectory for display/export.
 */
export function formatTrajectory(steps, options = {}) {
	if (!steps || steps.length === 0) return "No trajectory steps recorded.";

	const { maxSteps = 50, includeDetails = true } = options;
	const displaySteps = steps.slice(-maxSteps);

	const lines = [
		`Trajectory (${steps.length} steps total, showing last ${displaySteps.length}):`,
	];

	for (const step of displaySteps) {
		const prefix = step.isError ? "❌" : "✓";
		const cost = step.stepCost ? ` ($${step.stepCost.toFixed(4)})` : "";
		const tokens = step.inputTokens
			? ` [${step.inputTokens}/${step.outputTokens} tokens]`
			: "";

		if (step.stepType === "tool_call") {
			lines.push(
				`${prefix} [${step.stepNumber}] Tool: ${step.toolName}${cost}${tokens}`,
			);
			if (includeDetails && step.toolParams) {
				const params = JSON.stringify(step.toolParams).slice(0, 200);
				lines.push(`    Params: ${params}${params.length >= 200 ? "..." : ""}`);
			}
		} else if (step.stepType === "llm_call") {
			lines.push(
				`${prefix} [${step.stepNumber}] LLM: ${step.llmModel}${cost}${tokens}`,
			);
		} else if (step.stepType === "error") {
			lines.push(`${prefix} [${step.stepNumber}] Error: ${step.errorMessage}`);
			if (includeDetails && step.errorDetails) {
				lines.push(`    Details: ${step.errorDetails.slice(0, 300)}`);
			}
		} else {
			lines.push(
				`${prefix} [${step.stepNumber}] ${step.stepType}${cost}${tokens}`,
			);
		}
	}

	return lines.join("\n");
}

/**
 * Clear trajectory for a session.
 */
export function clearTrajectory(sessionId, db) {
	const dbConn = db || (isDbAvailable() ? require("./sf-db.js").getDb() : null);
	if (!dbConn) return 0;

	ensureTrajectoryTable(dbConn);

	try {
		const result = dbConn
			.prepare(`DELETE FROM ${TRAJECTORY_TABLE} WHERE session_id = ?`)
			.run(sessionId);
		resetStepCounter(sessionId);
		return result.changes;
	} catch (err) {
		logWarning("trajectory", "clearTrajectory failed", { error: String(err) });
		return 0;
	}
}
