/**
 * JSONL-backed session tree manager.
 *
 * Reads, writes, forks, labels, and exports persisted session entries while preserving branch structure.
 */
import { randomUUID } from "node:crypto";
import {
  closeSync,
  chmodSync,
  existsSync,
  mkdirSync,
  openSync,
  readdirSync,
  readFileSync,
  readSync,
  statSync,
  writeFileSync,
} from "node:fs";
import { readdir, readFile, stat } from "node:fs/promises";
import { join, resolve } from "node:path";
import {
  appendJsonlEntrySync,
  writeJsonlEntriesSync,
} from "../../config/sessions/transcript-jsonl.js";
import { CURRENT_SESSION_VERSION } from "../../config/sessions/version.js";
import type { ImageContent, Message, TextContent } from "../../llm/types.js";
import { getAgentDir as getDefaultAgentDir, getSessionsDir } from "../config.js";
import {
  type AgentMessage,
  buildSessionContext as buildCoreSessionContext,
  type SessionTreeEntry as CoreSessionTreeEntry,
  uuidv7,
} from "../runtime/index.js";
import type { BashExecutionMessage, CustomMessage } from "./messages.js";

export { CURRENT_SESSION_VERSION };

export interface SessionHeader {
  type: "session";
  version?: number; // v1 sessions don't have this
  id: string;
  timestamp: string;
  cwd: string;
  parentSession?: string;
}

export interface NewSessionOptions {
  id?: string;
  parentSession?: string;
}

export interface SessionEntryBase {
  type: string;
  id: string;
  parentId: string | null;
  timestamp: string;
}

export interface SessionMessageEntry extends SessionEntryBase {
  type: "message";
  message: AgentMessage;
}

export interface ThinkingLevelChangeEntry extends SessionEntryBase {
  type: "thinking_level_change";
  thinkingLevel: string;
}

export interface ModelChangeEntry extends SessionEntryBase {
  type: "model_change";
  provider: string;
  modelId: string;
}

export interface CompactionEntry<T = unknown> extends SessionEntryBase {
  type: "compaction";
  summary: string;
  firstKeptEntryId: string;
  tokensBefore: number;
  /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */
  details?: T;
  /** True if generated by an extension, undefined/false if runtime-generated (backward compatible) */
  fromHook?: boolean;
}

export interface BranchSummaryEntry<T = unknown> extends SessionEntryBase {
  type: "branch_summary";
  fromId: string;
  summary: string;
  /** Extension-specific data (not sent to LLM) */
  details?: T;
  /** True if generated by an extension, false if runtime-generated */
  fromHook?: boolean;
}

/**
 * Custom entry for extensions to store extension-specific data in the session.
 * Use customType to identify your extension's entries.
 *
 * Purpose: Persist extension state across session reloads. On reload, extensions can
 * scan entries for their customType and reconstruct internal state.
 *
 * Does NOT participate in LLM context (ignored by buildSessionContext).
 * For injecting content into context, see CustomMessageEntry.
 */
export interface CustomEntry<T = unknown> extends SessionEntryBase {
  type: "custom";
  customType: string;
  data?: T;
}

/** Label entry for user-defined bookmarks/markers on entries. */
export interface LabelEntry extends SessionEntryBase {
  type: "label";
  targetId: string;
  label: string | undefined;
}

/** Session metadata entry (e.g., user-defined display name). */
export interface SessionInfoEntry extends SessionEntryBase {
  type: "session_info";
  name?: string;
}

/**
 * Custom message entry for extensions to inject messages into LLM context.
 * Use customType to identify your extension's entries.
 *
 * Unlike CustomEntry, this DOES participate in LLM context.
 * The content is converted to a user message in buildSessionContext().
 * Use details for extension-specific metadata (not sent to LLM).
 *
 * display controls TUI rendering:
 * - false: hidden entirely
 * - true: rendered with distinct styling (different from user messages)
 */
export interface CustomMessageEntry<T = unknown> extends SessionEntryBase {
  type: "custom_message";
  customType: string;
  content: string | (TextContent | ImageContent)[];
  details?: T;
  display: boolean;
}

/** Session entry - has id/parentId for tree structure (returned by "read" methods in SessionManager) */
export type SessionEntry =
  | SessionMessageEntry
  | ThinkingLevelChangeEntry
  | ModelChangeEntry
  | CompactionEntry
  | BranchSummaryEntry
  | CustomEntry
  | CustomMessageEntry
  | LabelEntry
  | SessionInfoEntry;

/** Raw file entry (includes header) */
export type FileEntry = SessionHeader | SessionEntry;

/** Tree node for getTree() - defensive copy of session structure */
export interface SessionTreeNode {
  entry: SessionEntry;
  children: SessionTreeNode[];
  /** Resolved label for this entry, if any */
  label?: string;
  /** Timestamp of the latest label change for this entry, if any */
  labelTimestamp?: string;
}

export interface SessionContext {
  messages: AgentMessage[];
  thinkingLevel: string;
  model: { provider: string; modelId: string } | null;
}

export interface SessionInfo {
  path: string;
  id: string;
  /** Working directory where the session was started. Empty string for old sessions. */
  cwd: string;
  /** User-defined display name from session_info entries. */
  name?: string;
  /** Path to the parent session (if this session was forked). */
  parentSessionPath?: string;
  created: Date;
  modified: Date;
  messageCount: number;
  firstMessage: string;
  allMessagesText: string;
}

export type ReadonlySessionManager = Pick<
  SessionManager,
  | "getCwd"
  | "getSessionDir"
  | "getSessionId"
  | "getSessionFile"
  | "getLeafId"
  | "getLeafEntry"
  | "getEntry"
  | "getLabel"
  | "getBranch"
  | "getHeader"
  | "getEntries"
  | "getTree"
  | "getSessionName"
>;

function createSessionId(): string {
  return uuidv7();
}

/** Generate a unique short ID (8 hex chars, collision-checked) */
function generateId(byId: { has(id: string): boolean }): string {
  for (let i = 0; i < 100; i++) {
    const id = randomUUID().slice(0, 8);
    if (!byId.has(id)) {
      return id;
    }
  }
  // Fallback to full UUID if somehow we have collisions
  return randomUUID();
}

/** Migrate v1 → v2: add id/parentId tree structure. Mutates in place. */
function migrateV1ToV2(entries: FileEntry[]): void {
  const ids = new Set<string>();
  let prevId: string | null = null;

  for (const entry of entries) {
    if (entry.type === "session") {
      entry.version = 2;
      continue;
    }

    entry.id = generateId(ids);
    entry.parentId = prevId;
    prevId = entry.id;

    // Convert firstKeptEntryIndex to firstKeptEntryId for compaction
    if (entry.type === "compaction") {
      const comp = entry as CompactionEntry & { firstKeptEntryIndex?: number };
      if (typeof comp.firstKeptEntryIndex === "number") {
        const targetEntry = entries[comp.firstKeptEntryIndex];
        if (targetEntry && targetEntry.type !== "session") {
          comp.firstKeptEntryId = targetEntry.id;
        }
        delete comp.firstKeptEntryIndex;
      }
    }
  }
}

/** Migrate v2 → v3: rename hookMessage role to custom. Mutates in place. */
function migrateV2ToV3(entries: FileEntry[]): void {
  for (const entry of entries) {
    if (entry.type === "session") {
      entry.version = 3;
      continue;
    }

    // Update message entries with hookMessage role
    if (entry.type === "message") {
      const msgEntry = entry;
      if (msgEntry.message && (msgEntry.message as { role: string }).role === "hookMessage") {
        (msgEntry.message as { role: string }).role = "custom";
      }
    }
  }
}

/**
 * Run all necessary migrations to bring entries to current version.
 * Mutates entries in place. Returns true if any migration was applied.
 */
function migrateToCurrentVersion(entries: FileEntry[]): boolean {
  const header = entries.find((e) => e.type === "session");
  const version = header?.version ?? 1;

  if (version >= CURRENT_SESSION_VERSION) {
    return false;
  }

  if (version < 2) {
    migrateV1ToV2(entries);
  }
  if (version < 3) {
    migrateV2ToV3(entries);
  }

  return true;
}

/** Exported for testing */
export function migrateSessionEntries(entries: FileEntry[]): void {
  migrateToCurrentVersion(entries);
}

/** Exported for compaction.test.ts */
export function parseSessionEntries(content: string): FileEntry[] {
  const entries: FileEntry[] = [];
  const lines = content.trim().split("\n");

  for (const line of lines) {
    if (!line.trim()) {
      continue;
    }
    try {
      const entry = JSON.parse(line) as FileEntry;
      entries.push(entry);
    } catch {
      // Skip malformed lines
    }
  }

  return entries;
}

export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEntry | null {
  for (let i = entries.length - 1; i >= 0; i--) {
    if (entries[i].type === "compaction") {
      return entries[i] as CompactionEntry;
    }
  }
  return null;
}

/**
 * Build the session context from entries using tree traversal.
 * If leafId is provided, walks from that entry to root.
 * Handles compaction and branch summaries along the path.
 */
export function buildSessionContext(
  entries: SessionEntry[],
  leafId?: string | null,
  byIdInput?: Map<string, SessionEntry>,
): SessionContext {
  let byId = byIdInput;
  // Build uuid index if not available
  if (!byId) {
    byId = new Map<string, SessionEntry>();
    for (const entry of entries) {
      byId.set(entry.id, entry);
    }
  }

  // Find leaf
  let leaf: SessionEntry | undefined;
  if (leafId === null) {
    // Explicitly null - return no messages (navigated to before first entry)
    return { messages: [], thinkingLevel: "off", model: null };
  }
  if (leafId) {
    leaf = byId.get(leafId);
  }
  if (!leaf) {
    // Fallback to last entry (when leafId is undefined)
    leaf = entries[entries.length - 1];
  }

  if (!leaf) {
    return { messages: [], thinkingLevel: "off", model: null };
  }

  // Walk from leaf to root, collecting path
  const path: SessionEntry[] = [];
  let current: SessionEntry | undefined = leaf;
  while (current) {
    path.unshift(current);
    current = current.parentId ? byId.get(current.parentId) : undefined;
  }

  return buildCoreSessionContext(path as CoreSessionTreeEntry[]) as SessionContext;
}

/**
 * Compute the default session directory for a cwd.
 * Encodes cwd into a safe directory name under ~/.openclaw/agent/sessions/.
 */
export function getDefaultSessionDir(cwd: string, agentDir: string = getDefaultAgentDir()): string {
  const safePath = `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`;
  const sessionDir = join(agentDir, "sessions", safePath);
  if (!existsSync(sessionDir)) {
    mkdirSync(sessionDir, { recursive: true });
  }
  return sessionDir;
}

/** Exported for testing */
export function loadEntriesFromFile(filePath: string): FileEntry[] {
  if (!existsSync(filePath)) {
    return [];
  }

  const content = readFileSync(filePath, "utf8");
  const entries = parseJsonlEntries(content);

  // Validate session header
  if (entries.length === 0) {
    return entries;
  }
  const header = entries[0];
  if (header.type !== "session" || typeof (header as { id?: unknown }).id !== "string") {
    return [];
  }

  return entries;
}

function parseJsonlEntries(content: string): FileEntry[] {
  const entries: FileEntry[] = [];
  const lines = content.trim().split("\n");

  for (const line of lines) {
    if (!line.trim()) {
      continue;
    }
    try {
      const entry = JSON.parse(line) as FileEntry;
      entries.push(entry);
    } catch {
      // Skip malformed lines
    }
  }

  return entries;
}

function hasReadableSessionHeader(entries: FileEntry[]): boolean {
  const header = entries[0];
  return header?.type === "session" && typeof (header as { id?: unknown }).id === "string";
}

function buildCorruptSessionBackupPath(filePath: string): string {
  const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
  return `${filePath}.corrupt-${timestamp}-${randomUUID().slice(0, 8)}.jsonl`;
}

function recoverCorruptSessionEntries(filePath: string, cwd: string): FileEntry[] | null {
  const content = readFileSync(filePath, "utf8");
  if (content.trim().length === 0) {
    return null;
  }

  const parsedEntries = parseJsonlEntries(content);
  const recoveredHeader = parsedEntries.find(
    (entry): entry is SessionHeader =>
      entry.type === "session" && typeof (entry as { id?: unknown }).id === "string",
  );
  const header: SessionHeader =
    recoveredHeader ??
    ({
      type: "session",
      version: CURRENT_SESSION_VERSION,
      id: createSessionId(),
      timestamp: new Date().toISOString(),
      cwd,
    } satisfies SessionHeader);
  const recoveredEntries = parsedEntries.filter((entry) => entry.type !== "session");

  const backupPath = buildCorruptSessionBackupPath(filePath);
  const backupMode = statSync(filePath).mode & 0o777;
  writeFileSync(backupPath, content, { encoding: "utf8", mode: backupMode || 0o600 });
  chmodSync(backupPath, backupMode || 0o600);
  return [header, ...recoveredEntries];
}

function isValidSessionFile(filePath: string): boolean {
  try {
    const fd = openSync(filePath, "r");
    const buffer = Buffer.alloc(512);
    const bytesRead = readSync(fd, buffer, 0, 512, 0);
    closeSync(fd);
    const firstLine = buffer.toString("utf8", 0, bytesRead).split("\n")[0];
    if (!firstLine) {
      return false;
    }
    const header = JSON.parse(firstLine);
    return header.type === "session" && typeof header.id === "string";
  } catch {
    return false;
  }
}

/** Exported for testing */
export function findMostRecentSession(sessionDir: string): string | null {
  try {
    const files = readdirSync(sessionDir)
      .filter((f) => f.endsWith(".jsonl"))
      .map((f) => join(sessionDir, f))
      .filter(isValidSessionFile)
      .map((path) => ({ path, mtime: statSync(path).mtime }))
      .toSorted((a, b) => b.mtime.getTime() - a.mtime.getTime());

    return files[0]?.path || null;
  } catch {
    return null;
  }
}

function isMessageWithContent(message: AgentMessage): message is Message {
  return typeof (message as Message).role === "string" && "content" in message;
}

function extractTextContent(message: Message): string {
  const content = message.content;
  if (typeof content === "string") {
    return content;
  }
  return content
    .filter((block): block is TextContent => block.type === "text")
    .map((block) => block.text)
    .join(" ");
}

function getLastActivityTime(entries: FileEntry[]): number | undefined {
  let lastActivityTime: number | undefined;

  for (const entry of entries) {
    if (entry.type !== "message") {
      continue;
    }

    const message = entry.message;
    if (!isMessageWithContent(message)) {
      continue;
    }
    if (message.role !== "user" && message.role !== "assistant") {
      continue;
    }

    const msgTimestamp = (message as { timestamp?: number }).timestamp;
    if (typeof msgTimestamp === "number") {
      lastActivityTime = Math.max(lastActivityTime ?? 0, msgTimestamp);
      continue;
    }

    const entryTimestamp = (entry as SessionEntryBase).timestamp;
    if (typeof entryTimestamp === "string") {
      const t = new Date(entryTimestamp).getTime();
      if (!Number.isNaN(t)) {
        lastActivityTime = Math.max(lastActivityTime ?? 0, t);
      }
    }
  }

  return lastActivityTime;
}

function getSessionModifiedDate(
  entries: FileEntry[],
  header: SessionHeader,
  statsMtime: Date,
): Date {
  const lastActivityTime = getLastActivityTime(entries);
  if (typeof lastActivityTime === "number" && lastActivityTime > 0) {
    return new Date(lastActivityTime);
  }

  const headerTime =
    typeof header.timestamp === "string" ? new Date(header.timestamp).getTime() : Number.NaN;
  return !Number.isNaN(headerTime) ? new Date(headerTime) : statsMtime;
}

async function buildSessionInfo(filePath: string): Promise<SessionInfo | null> {
  try {
    const content = await readFile(filePath, "utf8");
    const entries: FileEntry[] = [];
    const lines = content.trim().split("\n");

    for (const line of lines) {
      if (!line.trim()) {
        continue;
      }
      try {
        entries.push(JSON.parse(line) as FileEntry);
      } catch {
        // Skip malformed lines
      }
    }

    if (entries.length === 0) {
      return null;
    }
    const header = entries[0];
    if (header.type !== "session") {
      return null;
    }

    const stats = await stat(filePath);
    let messageCount = 0;
    let firstMessage = "";
    const allMessages: string[] = [];
    let name: string | undefined;

    for (const entry of entries) {
      // Extract session name (use latest, including explicit clears)
      if (entry.type === "session_info") {
        const infoEntry = entry;
        name = infoEntry.name?.trim() || undefined;
      }

      if (entry.type !== "message") {
        continue;
      }
      messageCount++;

      const message = entry.message;
      if (!isMessageWithContent(message)) {
        continue;
      }
      if (message.role !== "user" && message.role !== "assistant") {
        continue;
      }

      const textContent = extractTextContent(message);
      if (!textContent) {
        continue;
      }

      allMessages.push(textContent);
      if (!firstMessage && message.role === "user") {
        firstMessage = textContent;
      }
    }

    const cwd = typeof header.cwd === "string" ? header.cwd : "";
    const parentSessionPath = header.parentSession;

    const modified = getSessionModifiedDate(entries, header, stats.mtime);

    return {
      path: filePath,
      id: header.id,
      cwd,
      name,
      parentSessionPath,
      created: new Date(header.timestamp),
      modified,
      messageCount,
      firstMessage: firstMessage || "(no messages)",
      allMessagesText: allMessages.join(" "),
    };
  } catch {
    return null;
  }
}

export type SessionListProgress = (loaded: number, total: number) => void;

const MAX_CONCURRENT_SESSION_INFO_LOADS = 10;

async function buildSessionInfosWithConcurrency(
  files: string[],
  onLoaded: () => void,
): Promise<(SessionInfo | null)[]> {
  const results: (SessionInfo | null)[] = Array.from({ length: files.length }, () => null);
  const inFlight = new Set<Promise<void>>();
  let nextIndex = 0;

  const startNext = (): void => {
    const index = nextIndex++;
    const file = files[index];
    if (!file) {
      return;
    }
    const task: Promise<void> = buildSessionInfo(file)
      .then((info) => {
        results[index] = info;
      })
      .catch(() => {
        results[index] = null;
      })
      .finally(() => {
        inFlight.delete(task);
        onLoaded();
      });
    inFlight.add(task);
  };

  while (nextIndex < files.length || inFlight.size > 0) {
    while (nextIndex < files.length && inFlight.size < MAX_CONCURRENT_SESSION_INFO_LOADS) {
      startNext();
    }
    if (inFlight.size > 0) {
      await Promise.race(inFlight);
    }
  }

  return results;
}

async function listSessionsFromDir(
  dir: string,
  onProgress?: SessionListProgress,
  progressOffset = 0,
  progressTotal?: number,
): Promise<SessionInfo[]> {
  const sessions: SessionInfo[] = [];
  if (!existsSync(dir)) {
    return sessions;
  }

  try {
    const dirEntries = await readdir(dir);
    const files = dirEntries.filter((f) => f.endsWith(".jsonl")).map((f) => join(dir, f));
    const total = progressTotal ?? files.length;

    let loaded = 0;
    const results = await buildSessionInfosWithConcurrency(files, () => {
      loaded++;
      onProgress?.(progressOffset + loaded, total);
    });
    for (const info of results) {
      if (info) {
        sessions.push(info);
      }
    }
  } catch {
    // Return empty list on error
  }

  return sessions;
}

/**
 * Manages conversation sessions as append-only trees stored in JSONL files.
 *
 * Each session entry has an id and parentId forming a tree structure. The "leaf"
 * pointer tracks the current position. Appending creates a child of the current leaf.
 * Branching moves the leaf to an earlier entry, allowing new branches without
 * modifying history.
 *
 * Use buildSessionContext() to get the resolved message list for the LLM, which
 * handles compaction summaries and follows the path from root to current leaf.
 */
export class SessionManager {
  private sessionId = "";
  private sessionFile: string | undefined;
  private sessionDir: string;
  private cwd: string;
  private shouldPersist: boolean;
  private flushed = false;
  private fileEntries: FileEntry[] = [];
  private byId: Map<string, SessionEntry> = new Map();
  private labelsById: Map<string, string> = new Map();
  private labelTimestampsById: Map<string, string> = new Map();
  private leafId: string | null = null;
  private recoveredCorruptHeader = false;

  private constructor(
    cwd: string,
    sessionDir: string,
    sessionFile: string | undefined,
    persist: boolean,
  ) {
    this.cwd = cwd;
    this.sessionDir = sessionDir;
    this.shouldPersist = persist;
    if (persist && sessionDir && !existsSync(sessionDir)) {
      mkdirSync(sessionDir, { recursive: true });
    }

    if (sessionFile) {
      this.setSessionFile(sessionFile);
    } else {
      this.newSession();
    }
  }

  /** Switch to a different session file (used for resume and branching) */
  setSessionFile(sessionFile: string): void {
    this.sessionFile = resolve(sessionFile);
    this.recoveredCorruptHeader = false;
    if (existsSync(this.sessionFile)) {
      this.fileEntries = loadEntriesFromFile(this.sessionFile);

      // If file was empty or corrupted (no valid header), truncate and start fresh
      // to avoid appending messages without a session header (which breaks the session)
      if (this.fileEntries.length === 0) {
        const recoveredEntries = recoverCorruptSessionEntries(this.sessionFile, this.cwd);
        if (recoveredEntries && hasReadableSessionHeader(recoveredEntries)) {
          this.fileEntries = recoveredEntries;
          const header = this.fileEntries.find((e) => e.type === "session");
          this.sessionId = header?.id ?? createSessionId();
          this.buildIndex();
          this.rewriteFile();
          this.recoveredCorruptHeader = true;
          this.flushed = true;
          return;
        }

        const explicitPath = this.sessionFile;
        this.newSession();
        this.sessionFile = explicitPath;
        this.rewriteFile();
        this.flushed = true;
        return;
      }

      const header = this.fileEntries.find((e) => e.type === "session");
      this.sessionId = header?.id ?? createSessionId();

      if (migrateToCurrentVersion(this.fileEntries)) {
        this.rewriteFile();
      }

      this.buildIndex();
      this.flushed = true;
    } else {
      const explicitPath = this.sessionFile;
      this.newSession();
      this.sessionFile = explicitPath; // preserve explicit path from --session flag
    }
  }

  newSession(options?: NewSessionOptions): string | undefined {
    this.recoveredCorruptHeader = false;
    this.sessionId = options?.id ?? createSessionId();
    const timestamp = new Date().toISOString();
    const header: SessionHeader = {
      type: "session",
      version: CURRENT_SESSION_VERSION,
      id: this.sessionId,
      timestamp,
      cwd: this.cwd,
      parentSession: options?.parentSession,
    };
    this.fileEntries = [header];
    this.byId.clear();
    this.labelsById.clear();
    this.leafId = null;
    this.flushed = false;

    if (this.shouldPersist) {
      const fileTimestamp = timestamp.replace(/[:.]/g, "-");
      this.sessionFile = join(this.getSessionDir(), `${fileTimestamp}_${this.sessionId}.jsonl`);
    }
    return this.sessionFile;
  }

  private buildIndex(): void {
    this.byId.clear();
    this.labelsById.clear();
    this.labelTimestampsById.clear();
    this.leafId = null;
    for (const entry of this.fileEntries) {
      if (entry.type === "session") {
        continue;
      }
      this.byId.set(entry.id, entry);
      this.leafId = entry.id;
      if (entry.type === "label") {
        if (entry.label) {
          this.labelsById.set(entry.targetId, entry.label);
          this.labelTimestampsById.set(entry.targetId, entry.timestamp);
        } else {
          this.labelsById.delete(entry.targetId);
          this.labelTimestampsById.delete(entry.targetId);
        }
      }
    }
  }

  private rewriteFile(): void {
    if (!this.shouldPersist || !this.sessionFile) {
      return;
    }
    writeJsonlEntriesSync(this.sessionFile, this.fileEntries);
  }

  isPersisted(): boolean {
    return this.shouldPersist;
  }

  getCwd(): string {
    return this.cwd;
  }

  getSessionDir(): string {
    return this.sessionDir;
  }

  getSessionId(): string {
    return this.sessionId;
  }

  wasRecoveredFromCorruptHeader(): boolean {
    return this.recoveredCorruptHeader;
  }

  getSessionFile(): string | undefined {
    return this.sessionFile;
  }

  persist(entry: SessionEntry): void {
    if (!this.shouldPersist || !this.sessionFile) {
      return;
    }

    const hasAssistant = this.fileEntries.some(
      (e) => e.type === "message" && e.message.role === "assistant",
    );
    if (!hasAssistant) {
      // Mark as not flushed so when assistant arrives, all entries get written
      this.flushed = false;
      return;
    }

    if (!this.flushed) {
      writeJsonlEntriesSync(this.sessionFile, this.fileEntries);
      this.flushed = true;
    } else {
      appendJsonlEntrySync(this.sessionFile, entry);
    }
  }

  private appendEntry(entry: SessionEntry): void {
    this.fileEntries.push(entry);
    this.byId.set(entry.id, entry);
    this.leafId = entry.id;
    this.persist(entry);
  }

  /** Append a message as child of current leaf, then advance leaf. Returns entry id.
   * Does not allow writing CompactionSummaryMessage and BranchSummaryMessage directly.
   * Reason: we want these to be top-level entries in the session, not message session entries,
   * so it is easier to find them.
   * These need to be appended via appendCompaction() and appendBranchSummary() methods.
   */
  appendMessage(message: Message | CustomMessage | BashExecutionMessage): string {
    const entry: SessionMessageEntry = {
      type: "message",
      id: generateId(this.byId),
      parentId: this.leafId,
      timestamp: new Date().toISOString(),
      message,
    };
    this.appendEntry(entry);
    return entry.id;
  }

  /** Append a thinking level change as child of current leaf, then advance leaf. Returns entry id. */
  appendThinkingLevelChange(thinkingLevel: string): string {
    const entry: ThinkingLevelChangeEntry = {
      type: "thinking_level_change",
      id: generateId(this.byId),
      parentId: this.leafId,
      timestamp: new Date().toISOString(),
      thinkingLevel,
    };
    this.appendEntry(entry);
    return entry.id;
  }

  /** Append a model change as child of current leaf, then advance leaf. Returns entry id. */
  appendModelChange(provider: string, modelId: string): string {
    const entry: ModelChangeEntry = {
      type: "model_change",
      id: generateId(this.byId),
      parentId: this.leafId,
      timestamp: new Date().toISOString(),
      provider,
      modelId,
    };
    this.appendEntry(entry);
    return entry.id;
  }

  /** Append a compaction summary as child of current leaf, then advance leaf. Returns entry id. */
  appendCompaction(
    summary: string,
    firstKeptEntryId: string,
    tokensBefore: number,
    details?: unknown,
    fromHook?: boolean,
  ): string {
    const entry: CompactionEntry = {
      type: "compaction",
      id: generateId(this.byId),
      parentId: this.leafId,
      timestamp: new Date().toISOString(),
      summary,
      firstKeptEntryId,
      tokensBefore,
      details,
      fromHook,
    };
    this.appendEntry(entry);
    return entry.id;
  }

  /** Append a custom entry (for extensions) as child of current leaf, then advance leaf. Returns entry id. */
  appendCustomEntry(customType: string, data?: unknown): string {
    const entry: CustomEntry = {
      type: "custom",
      customType,
      data,
      id: generateId(this.byId),
      parentId: this.leafId,
      timestamp: new Date().toISOString(),
    };
    this.appendEntry(entry);
    return entry.id;
  }

  /** Append a session info entry (e.g., display name). Returns entry id. */
  appendSessionInfo(name: string): string {
    const entry: SessionInfoEntry = {
      type: "session_info",
      id: generateId(this.byId),
      parentId: this.leafId,
      timestamp: new Date().toISOString(),
      name: name.trim(),
    };
    this.appendEntry(entry);
    return entry.id;
  }

  /** Get the current session name from the latest session_info entry, if any. */
  getSessionName(): string | undefined {
    // Walk entries in reverse to find the latest session_info entry.
    // Empty names explicitly clear the session title.
    const entries = this.getEntries();
    for (let i = entries.length - 1; i >= 0; i--) {
      const entry = entries[i];
      if (entry.type === "session_info") {
        return entry.name?.trim() || undefined;
      }
    }
    return undefined;
  }

  /**
   * Append a custom message entry (for extensions) that participates in LLM context.
   * @param customType Extension identifier for filtering on reload
   * @param content Message content (string or TextContent/ImageContent array)
   * @param display Whether to show in TUI (true = styled display, false = hidden)
   * @param details Optional extension-specific metadata (not sent to LLM)
   * @returns Entry id
   */
  appendCustomMessageEntry(
    customType: string,
    content: string | (TextContent | ImageContent)[],
    display: boolean,
    details?: unknown,
  ): string {
    const entry: CustomMessageEntry = {
      type: "custom_message",
      customType,
      content,
      display,
      details,
      id: generateId(this.byId),
      parentId: this.leafId,
      timestamp: new Date().toISOString(),
    };
    this.appendEntry(entry);
    return entry.id;
  }

  // =========================================================================
  // Tree Traversal
  // =========================================================================

  getLeafId(): string | null {
    return this.leafId;
  }

  getLeafEntry(): SessionEntry | undefined {
    return this.leafId ? this.byId.get(this.leafId) : undefined;
  }

  getEntry(id: string): SessionEntry | undefined {
    return this.byId.get(id);
  }

  /**
   * Get all direct children of an entry.
   */
  getChildren(parentId: string): SessionEntry[] {
    const children: SessionEntry[] = [];
    for (const entry of this.byId.values()) {
      if (entry.parentId === parentId) {
        children.push(entry);
      }
    }
    return children;
  }

  /**
   * Get the label for an entry, if any.
   */
  getLabel(id: string): string | undefined {
    return this.labelsById.get(id);
  }

  /**
   * Set or clear a label on an entry.
   * Labels are user-defined markers for bookmarking/navigation.
   * Pass undefined or empty string to clear the label.
   */
  appendLabelChange(targetId: string, label: string | undefined): string {
    if (!this.byId.has(targetId)) {
      throw new Error(`Entry ${targetId} not found`);
    }
    const entry: LabelEntry = {
      type: "label",
      id: generateId(this.byId),
      parentId: this.leafId,
      timestamp: new Date().toISOString(),
      targetId,
      label,
    };
    this.appendEntry(entry);
    if (label) {
      this.labelsById.set(targetId, label);
      this.labelTimestampsById.set(targetId, entry.timestamp);
    } else {
      this.labelsById.delete(targetId);
      this.labelTimestampsById.delete(targetId);
    }
    return entry.id;
  }

  /**
   * Walk from entry to root, returning all entries in path order.
   * Includes all entry types (messages, compaction, model changes, etc.).
   * Use buildSessionContext() to get the resolved messages for the LLM.
   */
  getBranch(fromId?: string): SessionEntry[] {
    const path: SessionEntry[] = [];
    const startId = fromId ?? this.leafId;
    let current = startId ? this.byId.get(startId) : undefined;
    while (current) {
      path.unshift(current);
      current = current.parentId ? this.byId.get(current.parentId) : undefined;
    }
    return path;
  }

  /**
   * Build the session context (what gets sent to the LLM).
   * Uses tree traversal from current leaf.
   */
  buildSessionContext(): SessionContext {
    return buildSessionContext(this.getEntries(), this.leafId, this.byId);
  }

  /**
   * Get session header.
   */
  getHeader(): SessionHeader | null {
    const h = this.fileEntries.find((e) => e.type === "session");
    return h ? h : null;
  }

  /**
   * Get all session entries (excludes header). Returns a shallow copy.
   * The session is append-only: use appendXXX() to add entries, branch() to
   * change the leaf pointer. Entries cannot be modified or deleted.
   */
  getEntries(): SessionEntry[] {
    return this.fileEntries.filter((e): e is SessionEntry => e.type !== "session");
  }

  /**
   * Get the session as a tree structure. Returns a shallow defensive copy of all entries.
   * A well-formed session has exactly one root (first entry with parentId === null).
   * Orphaned entries (broken parent chain) are also returned as roots.
   */
  getTree(): SessionTreeNode[] {
    const entries = this.getEntries();
    const nodeMap = new Map<string, SessionTreeNode>();
    const roots: SessionTreeNode[] = [];

    // Create nodes with resolved labels
    for (const entry of entries) {
      const label = this.labelsById.get(entry.id);
      const labelTimestamp = this.labelTimestampsById.get(entry.id);
      nodeMap.set(entry.id, { entry, children: [], label, labelTimestamp });
    }

    // Build tree
    for (const entry of entries) {
      const node = nodeMap.get(entry.id)!;
      if (entry.parentId === null || entry.parentId === entry.id) {
        roots.push(node);
      } else {
        const parent = nodeMap.get(entry.parentId);
        if (parent) {
          parent.children.push(node);
        } else {
          // Orphan - treat as root
          roots.push(node);
        }
      }
    }

    // Sort children by timestamp (oldest first, newest at bottom)
    // Use iterative approach to avoid stack overflow on deep trees
    const stack: SessionTreeNode[] = [...roots];
    while (stack.length > 0) {
      const node = stack.pop()!;
      node.children.sort(
        (a, b) => new Date(a.entry.timestamp).getTime() - new Date(b.entry.timestamp).getTime(),
      );
      stack.push(...node.children);
    }

    return roots;
  }

  // =========================================================================
  // Branching
  // =========================================================================

  /**
   * Start a new branch from an earlier entry.
   * Moves the leaf pointer to the specified entry. The next appendXXX() call
   * will create a child of that entry, forming a new branch. Existing entries
   * are not modified or deleted.
   */
  branch(branchFromId: string): void {
    if (!this.byId.has(branchFromId)) {
      throw new Error(`Entry ${branchFromId} not found`);
    }
    this.leafId = branchFromId;
  }

  /**
   * Reset the leaf pointer to null (before any entries).
   * The next appendXXX() call will create a new root entry (parentId = null).
   * Use this when navigating to re-edit the first user message.
   */
  resetLeaf(): void {
    this.leafId = null;
  }

  /**
   * Start a new branch with a summary of the abandoned path.
   * Same as branch(), but also appends a branch_summary entry that captures
   * context from the abandoned conversation path.
   */
  branchWithSummary(
    branchFromId: string | null,
    summary: string,
    details?: unknown,
    fromHook?: boolean,
  ): string {
    if (branchFromId !== null && !this.byId.has(branchFromId)) {
      throw new Error(`Entry ${branchFromId} not found`);
    }
    this.leafId = branchFromId;
    const entry: BranchSummaryEntry = {
      type: "branch_summary",
      id: generateId(this.byId),
      parentId: branchFromId,
      timestamp: new Date().toISOString(),
      fromId: branchFromId ?? "root",
      summary,
      details,
      fromHook,
    };
    this.appendEntry(entry);
    return entry.id;
  }

  /**
   * Create a new session file containing only the path from root to the specified leaf.
   * Useful for extracting a single conversation path from a branched session.
   * Returns the new session file path, or undefined if not persisting.
   */
  createBranchedSession(leafId: string): string | undefined {
    const previousSessionFile = this.sessionFile;
    const path = this.getBranch(leafId);
    if (path.length === 0) {
      throw new Error(`Entry ${leafId} not found`);
    }

    // Filter out LabelEntry from path - we'll recreate them from the resolved map
    const pathWithoutLabels = path.filter((e) => e.type !== "label");

    const newSessionId = createSessionId();
    const timestamp = new Date().toISOString();
    const fileTimestamp = timestamp.replace(/[:.]/g, "-");
    const newSessionFile = join(this.getSessionDir(), `${fileTimestamp}_${newSessionId}.jsonl`);

    const header: SessionHeader = {
      type: "session",
      version: CURRENT_SESSION_VERSION,
      id: newSessionId,
      timestamp,
      cwd: this.cwd,
      parentSession: this.shouldPersist ? previousSessionFile : undefined,
    };

    // Collect labels for entries in the path
    const pathEntryIds = new Set(pathWithoutLabels.map((e) => e.id));
    const labelsToWrite: Array<{ targetId: string; label: string; timestamp: string }> = [];
    for (const [targetId, label] of this.labelsById) {
      if (pathEntryIds.has(targetId)) {
        labelsToWrite.push({ targetId, label, timestamp: this.labelTimestampsById.get(targetId)! });
      }
    }

    if (this.shouldPersist) {
      // Build label entries
      const lastEntryId = pathWithoutLabels[pathWithoutLabels.length - 1]?.id || null;
      let parentId = lastEntryId;
      const labelEntries: LabelEntry[] = [];
      for (const { targetId, label, timestamp: labelTimestamp } of labelsToWrite) {
        const labelEntry: LabelEntry = {
          type: "label",
          id: generateId(new Set(pathEntryIds)),
          parentId,
          timestamp: labelTimestamp,
          targetId,
          label,
        };
        pathEntryIds.add(labelEntry.id);
        labelEntries.push(labelEntry);
        parentId = labelEntry.id;
      }

      this.fileEntries = [header, ...pathWithoutLabels, ...labelEntries];
      this.sessionId = newSessionId;
      this.sessionFile = newSessionFile;
      this.buildIndex();

      // Only write the file now if it contains an assistant message.
      // Otherwise defer to persist(), which creates the file on the
      // first assistant response, matching the newSession() contract
      // and avoiding the duplicate-header bug when persist()'s
      // no-assistant guard later resets flushed to false.
      const hasAssistant = this.fileEntries.some(
        (e) => e.type === "message" && e.message.role === "assistant",
      );
      if (hasAssistant) {
        this.rewriteFile();
        this.flushed = true;
      } else {
        this.flushed = false;
      }

      return newSessionFile;
    }

    // In-memory mode: replace current session with the path + labels
    const labelEntries: LabelEntry[] = [];
    let parentId = pathWithoutLabels[pathWithoutLabels.length - 1]?.id || null;
    for (const { targetId, label, timestamp: labelTimestamp } of labelsToWrite) {
      const labelEntry: LabelEntry = {
        type: "label",
        id: generateId(new Set([...pathEntryIds, ...labelEntries.map((e) => e.id)])),
        parentId,
        timestamp: labelTimestamp,
        targetId,
        label,
      };
      labelEntries.push(labelEntry);
      parentId = labelEntry.id;
    }
    this.fileEntries = [header, ...pathWithoutLabels, ...labelEntries];
    this.sessionId = newSessionId;
    this.buildIndex();
    return undefined;
  }

  /**
   * Create a new session.
   * @param cwd Working directory (stored in session header)
   * @param sessionDir Optional session directory. If omitted, uses default (~/.openclaw/agent/sessions/<encoded-cwd>/).
   */
  static create(cwd: string, sessionDir?: string): SessionManager {
    const dir = sessionDir ?? getDefaultSessionDir(cwd);
    return new SessionManager(cwd, dir, undefined, true);
  }

  /**
   * Open a specific session file.
   * @param path Path to session file
   * @param sessionDir Optional session directory for /new or /branch. If omitted, derives from file's parent.
   * @param cwdOverride Optional cwd override instead of the session header cwd.
   */
  static open(path: string, sessionDir?: string, cwdOverride?: string): SessionManager {
    // Extract cwd from session header if possible, otherwise use process.cwd()
    const entries = loadEntriesFromFile(path);
    const header = entries.find((e) => e.type === "session");
    const cwd = cwdOverride ?? header?.cwd ?? process.cwd();
    // If no sessionDir provided, derive from file's parent directory
    const dir = sessionDir ?? resolve(path, "..");
    return new SessionManager(cwd, dir, path, true);
  }

  /**
   * Continue the most recent session, or create new if none.
   * @param cwd Working directory
   * @param sessionDir Optional session directory. If omitted, uses default (~/.openclaw/agent/sessions/<encoded-cwd>/).
   */
  static continueRecent(cwd: string, sessionDir?: string): SessionManager {
    const dir = sessionDir ?? getDefaultSessionDir(cwd);
    const mostRecent = findMostRecentSession(dir);
    if (mostRecent) {
      return new SessionManager(cwd, dir, mostRecent, true);
    }
    return new SessionManager(cwd, dir, undefined, true);
  }

  /** Create an in-memory session (no file persistence) */
  static inMemory(cwd: string = process.cwd()): SessionManager {
    return new SessionManager(cwd, "", undefined, false);
  }

  /**
   * Fork a session from another project directory into the current project.
   * Creates a new session in the target cwd with the full history from the source session.
   * @param sourcePath Path to the source session file
   * @param targetCwd Target working directory (where the new session will be stored)
   * @param sessionDir Optional session directory. If omitted, uses default for targetCwd.
   */
  static forkFrom(sourcePath: string, targetCwd: string, sessionDir?: string): SessionManager {
    const sourceEntries = loadEntriesFromFile(sourcePath);
    if (sourceEntries.length === 0) {
      throw new Error(`Cannot fork: source session file is empty or invalid: ${sourcePath}`);
    }

    const sourceHeader = sourceEntries.find((e) => e.type === "session");
    if (!sourceHeader) {
      throw new Error(`Cannot fork: source session has no header: ${sourcePath}`);
    }

    const dir = sessionDir ?? getDefaultSessionDir(targetCwd);
    if (!existsSync(dir)) {
      mkdirSync(dir, { recursive: true });
    }

    // Create new session file with new ID but forked content
    const newSessionId = createSessionId();
    const timestamp = new Date().toISOString();
    const fileTimestamp = timestamp.replace(/[:.]/g, "-");
    const newSessionFile = join(dir, `${fileTimestamp}_${newSessionId}.jsonl`);

    // Write new header pointing to source as parent, with updated cwd
    const newHeader: SessionHeader = {
      type: "session",
      version: CURRENT_SESSION_VERSION,
      id: newSessionId,
      timestamp,
      cwd: targetCwd,
      parentSession: sourcePath,
    };
    appendJsonlEntrySync(newSessionFile, newHeader);

    // Copy all non-header entries from source
    for (const entry of sourceEntries) {
      if (entry.type !== "session") {
        appendJsonlEntrySync(newSessionFile, entry);
      }
    }

    return new SessionManager(targetCwd, dir, newSessionFile, true);
  }

  /**
   * List all sessions for a directory.
   * @param cwd Working directory (used to compute default session directory)
   * @param sessionDir Optional session directory. If omitted, uses default (~/.openclaw/agent/sessions/<encoded-cwd>/).
   * @param onProgress Optional callback for progress updates (loaded, total)
   */
  static async list(
    cwd: string,
    sessionDir?: string,
    onProgress?: SessionListProgress,
  ): Promise<SessionInfo[]> {
    const dir = sessionDir ?? getDefaultSessionDir(cwd);
    const sessions = await listSessionsFromDir(dir, onProgress);
    sessions.sort((a, b) => b.modified.getTime() - a.modified.getTime());
    return sessions;
  }

  /**
   * List all sessions across all project directories.
   * @param onProgress Optional callback for progress updates (loaded, total)
   */
  static async listAll(onProgress?: SessionListProgress): Promise<SessionInfo[]> {
    const sessionsDir = getSessionsDir();

    try {
      if (!existsSync(sessionsDir)) {
        return [];
      }
      const entries = await readdir(sessionsDir, { withFileTypes: true });
      const dirs = entries.filter((e) => e.isDirectory()).map((e) => join(sessionsDir, e.name));

      // Count total files first for accurate progress
      let totalFiles = 0;
      const dirFiles: string[][] = [];
      for (const dir of dirs) {
        try {
          const files = (await readdir(dir)).filter((f) => f.endsWith(".jsonl"));
          dirFiles.push(files.map((f) => join(dir, f)));
          totalFiles += files.length;
        } catch {
          dirFiles.push([]);
        }
      }

      // Process all files with progress tracking
      let loaded = 0;
      const sessions: SessionInfo[] = [];
      const allFiles = dirFiles.flat();

      const results = await buildSessionInfosWithConcurrency(allFiles, () => {
        loaded++;
        onProgress?.(loaded, totalFiles);
      });

      for (const info of results) {
        if (info) {
          sessions.push(info);
        }
      }

      sessions.sort((a, b) => b.modified.getTime() - a.modified.getTime());
      return sessions;
    } catch {
      return [];
    }
  }
}
