/**
 * Structured Trace Data Model & Export
 *
 * Provides a hierarchical span model for tracing autonomous mode execution.
 * Spans form a tree: root session span → unit spans (milestone/slice/task) → tool spans.
 *
 * Two export modes:
 *   exportTrace(path)          — write to arbitrary path
 *   exportTraceToProject(dir)  — write to .sf/traces/ in project dir
 */

import { randomUUID } from "node:crypto";
import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { join } from "node:path";

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

export type SpanKind = "session" | "unit" | "tool";
export type SpanStatus =
	| "ok"
	| "error"
	| "cancelled"
	| "timeout"
	| "in_progress";

export interface TraceEvent {
	name: string;
	timestamp: number;
	attributes?: Record<string, string | number | boolean | null>;
}

export interface SpanAttributes {
	// Session-level
	projectRoot?: string;
	sessionId?: string;
	cwd?: string;
	command?: string;
	model?: string;
	inputTokens?: number;
	outputTokens?: number;
	cacheReadTokens?: number;
	cacheWriteTokens?: number;
	costUsd?: number;
	exitCode?: number;

	// Unit-level
	unitType?: "milestone" | "slice" | "task";
	unitId?: string;
	unitStatus?: SpanStatus;
	unitErrorReason?: string;

	// Tool-level
	toolName?: string;
	toolCallId?: string;
	toolStatus?: SpanStatus;
	toolError?: string;
	toolDurationMs?: number;
}

export interface Span {
	id: string;
	name: string;
	kind: SpanKind;
	status: SpanStatus;
	startTime: number;
	endTime?: number;
	attributes: SpanAttributes;
	children: Span[];
	events: TraceEvent[];
}

export interface Trace {
	id: string;
	version: number;
	projectRoot: string;
	sessionId?: string;
	startedAt: string;
	completedAt?: string;
	rootSpan: Span;
}

// ---------------------------------------------------------------------------
// Span helpers
// ---------------------------------------------------------------------------

/** Create a new span with a random UUID and current timestamp. */
export function createSpan(
	name: string,
	kind: SpanKind,
	attributes: SpanAttributes = {},
): Span {
	return {
		id: randomUUID(),
		name,
		kind,
		status: "in_progress",
		startTime: Date.now(),
		attributes,
		children: [],
		events: [],
	};
}

/** Mark a span as complete and record end time. */
export function endSpan(span: Span, status: SpanStatus = "ok"): Span {
	span.status = status;
	span.endTime = Date.now();
	return span;
}

/** Append a named event to a span with optional attributes. */
export function addEvent(
	span: Span,
	name: string,
	attributes?: Record<string, string | number | boolean | null>,
): void {
	span.events.push({
		name,
		timestamp: Date.now(),
		attributes,
	});
}

/** Append an error event to a span with message and optional stack. */
export function addError(span: Span, message: string, stack?: string): void {
	span.events.push({
		name: "error",
		timestamp: Date.now(),
		attributes: {
			message,
			...(stack ? { stack } : {}),
		},
	});
	span.status = "error";
	if (!span.endTime) span.endTime = Date.now();
}

// ---------------------------------------------------------------------------
// Trace helpers
// ---------------------------------------------------------------------------

/** Create a new trace with a root session span. */
export function createTrace(
	projectRoot: string,
	sessionId?: string,
	command?: string,
	model?: string,
): Trace {
	const rootSpan = createSpan(`session:${sessionId ?? "unknown"}`, "session", {
		sessionId,
		projectRoot,
		command,
		model,
	});
	return {
		id: randomUUID(),
		version: 1,
		projectRoot,
		sessionId,
		startedAt: new Date().toISOString(),
		rootSpan,
	};
}

/** Finalize a trace: set completedAt timestamp. */
export function finalizeTrace(trace: Trace): Trace {
	trace.completedAt = new Date().toISOString();
	return trace;
}

/** Find a span in the tree by ID (linear walk). */
export function findSpan(span: Span, id: string): Span | undefined {
	if (span.id === id) return span;
	for (const child of span.children) {
		const found = findSpan(child, id);
		if (found) return found;
	}
	return undefined;
}

/** Add a child span to a parent. */
export function addChildSpan(parent: Span, child: Span): void {
	parent.children.push(child);
}

/** Walk all spans in a trace (root first, depth-first). Yields each span. */
export function* walkSpans(span: Span): Generator<Span, void, unknown> {
	yield span;
	for (const child of span.children) {
		yield* walkSpans(child);
	}
}

// ---------------------------------------------------------------------------
// Export
// ---------------------------------------------------------------------------

/**
 * Serialize and write a trace to an arbitrary path.
 * Creates parent directories as needed.
 */
export function exportTrace(trace: Trace, path: string): void {
	const dir = join(path, "..");
	if (!existsSync(dir)) {
		mkdirSync(dir, { recursive: true });
	}
	writeFileSync(path, JSON.stringify(trace, null, 2), "utf-8");
}

/**
 * Serialize and write a trace to .sf/traces/ in the project root.
 * Filename: trace-<timestamp>.json
 */
export function exportTraceToProject(
	trace: Trace,
	projectRoot: string,
): string {
	const tracesDir = join(projectRoot, ".sf", "traces");
	if (!existsSync(tracesDir)) {
		mkdirSync(tracesDir, { recursive: true });
	}
	const filename = `trace-${Date.now()}.json`;
	const path = join(tracesDir, filename);
	writeFileSync(path, JSON.stringify(trace, null, 2), "utf-8");
	return path;
}

/**
 * Read a trace from disk.
 */
export function readTrace(path: string): Trace {
	return JSON.parse(readFileSync(path, "utf-8")) as Trace;
}
