import { appendFile, mkdir, readFile, rm, writeFile } from 'node:fs/promises';
import * as path from 'node:path';
import { randomUUID } from 'node:crypto';
import { getPathToLeaf } from './path.js';
import {
  SessionTreeProviderError,
  type SessionTreeProvider,
} from './provider.js';
import { buildSessionTreeSnapshot } from './snapshot.js';
import type {
  AppendSessionTreeEntryInput,
  CreateSessionTreeInput,
  SessionTreeEntry,
  SessionTreeHeader,
  SessionTreeMessage,
  SessionTreeSession,
  SessionTreeSnapshot,
  UpdateSessionTreeEntryPatch,
} from './types.js';

const SESSION_FILE_VERSION = 1;
const HEADER_LINE_NUMBER = 1;

type SessionIdGenerator = () => string;
type EntryIdGenerator = (existingIds?: { has(id: string): boolean }) => string;
type NowGenerator = () => string;
type CwdGenerator = () => string;

interface FileSessionTreeProviderOptions {
  rootDir: string;
  sessionIdGenerator?: SessionIdGenerator;
  entryIdGenerator?: EntryIdGenerator;
  now?: NowGenerator;
  cwd?: string | CwdGenerator;
}

interface SessionFileState {
  header: SessionTreeHeader;
  entries: SessionTreeEntry[];
  sessionFile: string;
}

export class FileSessionTreeProvider implements SessionTreeProvider {
  private readonly rootDir: string;
  private readonly sessionIdGenerator: SessionIdGenerator;
  private readonly entryIdGenerator: EntryIdGenerator;
  private readonly now: NowGenerator;
  private readonly cwd: CwdGenerator;

  constructor(options: FileSessionTreeProviderOptions) {
    this.rootDir = options.rootDir;
    this.sessionIdGenerator = options.sessionIdGenerator ?? createDefaultSessionId;
    this.entryIdGenerator = options.entryIdGenerator ?? createDefaultEntryId;
    this.now = options.now ?? (() => new Date().toISOString());
    if (typeof options.cwd === 'function') {
      this.cwd = options.cwd;
    } else {
      const cwd = options.cwd ?? process.cwd();
      this.cwd = () => cwd;
    }
  }

  async createSession(input: CreateSessionTreeInput): Promise<SessionTreeSession> {
    const sessionId = input.sessionId ?? this.sessionIdGenerator();
    const sessionFile = this.getSessionFilePath(sessionId);
    const timestamp = input.createdAt ?? this.now();
    assertJsonSerializable(input.metadata, `createSession metadata for session ${sessionId}`);
    const header: SessionTreeHeader = {
      type: 'session',
      version: SESSION_FILE_VERSION,
      id: sessionId,
      timestamp,
      cwd: input.cwd ?? this.cwd(),
      parentSession: input.parentSessionId,
      title: input.title,
      agentId: input.agentId,
      userId: input.userId,
      metadata: cloneValue(input.metadata),
    };

    if (Object.prototype.hasOwnProperty.call(input, 'leafEntryId')) {
      header.leafId = input.leafEntryId ?? null;
    }

    if (input.status !== undefined) {
      header.status = input.status;
    }

    if (input.updatedAt !== undefined) {
      header.updatedAt = input.updatedAt;
    }

    await mkdir(this.rootDir, { recursive: true });
    await writeFile(sessionFile, `${JSON.stringify(header)}\n`, 'utf8');

    return buildSessionFromState({
      header,
      entries: [],
      sessionFile,
    });
  }

  async getSession(sessionId: string): Promise<SessionTreeSession | null> {
    const state = await this.readSessionStateOrNull(sessionId);
    return state ? buildSessionFromState(state) : null;
  }

  async updateSession(sessionId: string, patch: Partial<SessionTreeSession>): Promise<SessionTreeSession> {
    const state = await this.requireSessionState(sessionId);
    assertJsonSerializable(patch.metadata, `updateSession metadata for session ${sessionId}`);
    const header: SessionTreeHeader = {
      ...state.header,
      cwd: patch.cwd ?? state.header.cwd,
      parentSession: patch.parentSessionId ?? state.header.parentSession,
      title: patch.title ?? state.header.title,
      agentId: patch.agentId ?? state.header.agentId,
      userId: patch.userId ?? state.header.userId,
      metadata: patch.metadata === undefined ? state.header.metadata : cloneValue(patch.metadata),
    };

    if (Object.prototype.hasOwnProperty.call(patch, 'leafEntryId')) {
      header.leafId = patch.leafEntryId ?? null;
    }

    if (patch.status !== undefined) {
      header.status = patch.status;
    }

    if (patch.updatedAt !== undefined) {
      header.updatedAt = patch.updatedAt;
    } else if (Object.keys(patch).length > 0) {
      header.updatedAt = this.now();
    }

    await this.writeSessionState({
      ...state,
      header,
    });

    return buildSessionFromState({
      ...state,
      header,
    });
  }

  async deleteSession(sessionId: string): Promise<void> {
    const sessionFile = this.getSessionFilePath(sessionId);
    await rm(sessionFile, { force: true });
  }

  async listEntries(sessionId: string): Promise<SessionTreeEntry[]> {
    const state = await this.requireSessionState(sessionId);
    return cloneEntries(state.entries);
  }

  async appendEntry(sessionId: string, input: AppendSessionTreeEntryInput): Promise<SessionTreeEntry> {
    const state = await this.requireSessionState(sessionId);
    this.assertAppendParentConstraint(state.entries, input.parentId);

    const existingIds = new Set(state.entries.map(entry => entry.id));
    const entryCandidate = {
      ...input,
      id: input.id ?? this.entryIdGenerator(existingIds),
      timestamp: input.timestamp ?? this.now(),
    } as SessionTreeEntry;
    assertJsonSerializable(entryCandidate, `appendEntry input for session ${sessionId}`);
    const entry = cloneValue(entryCandidate);

    if (existingIds.has(entry.id)) {
      throw new SessionTreeProviderError(`Entry ${entry.id} already exists in session ${sessionId}`);
    }

    const header: SessionTreeHeader = {
      ...state.header,
      leafId: entry.id,
      updatedAt: this.now(),
    };
    const entries = state.entries.concat(entry);

    await this.writeSessionState({
      ...state,
      header,
      entries,
    });

    return cloneEntry(entry);
  }

  async appendMessage(sessionId: string, message: SessionTreeMessage): Promise<SessionTreeEntry> {
    const session = await this.requireSession(sessionId);
    return this.appendEntry(sessionId, {
      parentId: session.leafEntryId,
      type: 'message',
      message,
    });
  }

  async appendThinkingLevelChange(sessionId: string, thinkingLevel: string): Promise<SessionTreeEntry> {
    const session = await this.requireSession(sessionId);
    return this.appendEntry(sessionId, {
      parentId: session.leafEntryId,
      type: 'thinking_level_change',
      thinkingLevel,
    });
  }

  async appendModelChange(sessionId: string, provider: string, modelId: string): Promise<SessionTreeEntry> {
    const session = await this.requireSession(sessionId);
    return this.appendEntry(sessionId, {
      parentId: session.leafEntryId,
      type: 'model_change',
      provider,
      modelId,
    });
  }

  async appendCompaction(sessionId: string, input: {
    summary: string;
    firstKeptEntryId: string;
    tokensBefore: number;
    details?: Record<string, unknown>;
  }): Promise<SessionTreeEntry> {
    const session = await this.requireSession(sessionId);
    return this.appendEntry(sessionId, {
      parentId: session.leafEntryId,
      type: 'compaction',
      summary: input.summary,
      firstKeptEntryId: input.firstKeptEntryId,
      tokensBefore: input.tokensBefore,
      details: input.details,
    });
  }

  async appendBranchSummary(sessionId: string, input: {
    fromId: string;
    summary: string;
    details?: Record<string, unknown>;
  }): Promise<SessionTreeEntry> {
    const session = await this.requireSession(sessionId);
    return this.appendEntry(sessionId, {
      parentId: session.leafEntryId,
      type: 'branch_summary',
      fromId: input.fromId,
      summary: input.summary,
      details: input.details,
    });
  }

  async appendLabelChange(sessionId: string, targetId: string, label: string | undefined): Promise<SessionTreeEntry> {
    const session = await this.requireSession(sessionId);
    return this.appendEntry(sessionId, {
      parentId: session.leafEntryId,
      type: 'label',
      targetId,
      label,
    });
  }

  async appendSessionInfo(sessionId: string, name?: string): Promise<SessionTreeEntry> {
    const session = await this.requireSession(sessionId);
    return this.appendEntry(sessionId, {
      parentId: session.leafEntryId,
      type: 'session_info',
      name,
    });
  }

  async updateEntry(sessionId: string, id: string, patch: UpdateSessionTreeEntryPatch): Promise<SessionTreeEntry> {
    const state = await this.requireSessionState(sessionId);
    const index = state.entries.findIndex(entry => entry.id === id);

    if (index < 0) {
      throw new SessionTreeProviderError(`Entry ${id} not found in session ${sessionId}`);
    }

    const entry = state.entries[index]!;
    assertImmutableEntryPatch(entry, patch as Record<string, unknown>);
    assertJsonSerializable(patch, `updateEntry patch for entry ${id} in session ${sessionId}`);

    const updatedEntry = cloneValue({
      ...entry,
      ...cloneEntryPatch(patch),
      id: entry.id,
      type: entry.type,
      parentId: entry.parentId,
      timestamp: entry.timestamp,
    }) as SessionTreeEntry;

    const entries = state.entries.slice();
    entries[index] = updatedEntry;

    const header = {
      ...state.header,
      updatedAt: this.now(),
    };

    await this.writeSessionState({
      ...state,
      header,
      entries,
    });

    return cloneEntry(updatedEntry);
  }

  async switchLeaf(sessionId: string, leafEntryId: string | null): Promise<SessionTreeSession> {
    const state = await this.requireSessionState(sessionId);

    if (leafEntryId !== null && !state.entries.some(entry => entry.id === leafEntryId)) {
      throw new SessionTreeProviderError(`Leaf entry ${leafEntryId} not found in session ${sessionId}`);
    }

    const header: SessionTreeHeader = {
      ...state.header,
      leafId: leafEntryId,
      updatedAt: this.now(),
    };

    await this.writeSessionState({
      ...state,
      header,
    });

    return buildSessionFromState({
      ...state,
      header,
    });
  }

  async resetLeaf(sessionId: string): Promise<SessionTreeSession> {
    return this.switchLeaf(sessionId, null);
  }

  async createBranchedSession(
    sessionId: string,
    leafEntryId: string,
    input: CreateSessionTreeInput = {},
  ): Promise<SessionTreeSession> {
    const state = await this.requireSessionState(sessionId);
    const pathEntries = getPathToLeaf(state.entries, leafEntryId);
    const session = buildSessionFromState(state);
    const created = await this.createSession({
      ...input,
      provider: 'file',
      parentSessionId: input.parentSessionId ?? sessionId,
      cwd: input.cwd ?? session.cwd,
      agentId: input.agentId ?? session.agentId,
      userId: input.userId ?? session.userId,
      title: input.title ?? session.title,
      metadata: input.metadata ?? session.metadata,
    });

    for (const entry of pathEntries) {
      await this.appendEntry(created.sessionId, cloneEntry(entry) as AppendSessionTreeEntryInput);
    }

    await this.switchLeaf(created.sessionId, leafEntryId);
    return (await this.getSession(created.sessionId))!;
  }

  async getActivePath(sessionId: string): Promise<SessionTreeEntry[]> {
    const state = await this.requireSessionState(sessionId);
    const session = buildSessionFromState(state);

    if (session.leafEntryId === null) {
      return [];
    }

    return cloneEntries(getPathToLeaf(state.entries, session.leafEntryId));
  }

  async getSnapshot(sessionId: string): Promise<SessionTreeSnapshot> {
    const state = await this.requireSessionState(sessionId);
    const session = buildSessionFromState(state);

    return cloneSnapshot(buildSessionTreeSnapshot(session, cloneEntries(state.entries)));
  }

  getSessionFilePath(sessionId: string): string {
    return path.join(this.rootDir, `${Buffer.from(sessionId, 'utf8').toString('base64url')}.jsonl`);
  }

  private async readSessionStateOrNull(sessionId: string): Promise<SessionFileState | null> {
    const sessionFile = this.getSessionFilePath(sessionId);

    try {
      const content = await readFile(sessionFile, 'utf8');
      return parseSessionFile(content, sessionFile, sessionId);
    } catch (error) {
      if (isNodeError(error) && error.code === 'ENOENT') {
        return null;
      }
      throw error;
    }
  }

  private async requireSessionState(sessionId: string): Promise<SessionFileState> {
    const state = await this.readSessionStateOrNull(sessionId);
    if (!state) {
      throw new SessionTreeProviderError(`Session ${sessionId} not found`);
    }
    return state;
  }

  private async requireSession(sessionId: string): Promise<SessionTreeSession> {
    const state = await this.requireSessionState(sessionId);
    return buildSessionFromState(state);
  }

  private assertAppendParentConstraint(
    entries: SessionTreeEntry[],
    parentId: string | null,
  ): void {
    if (parentId === null) {
      if (entries.some(entry => entry.parentId === null)) {
        throw new SessionTreeProviderError('Session already has a root entry');
      }
      return;
    }

    if (!entries.some(entry => entry.id === parentId)) {
      throw new SessionTreeProviderError(`Parent entry ${parentId} not found`);
    }
  }

  private async writeSessionState(state: SessionFileState): Promise<void> {
    const content = serializeSessionFile(state.header, state.entries);
    await writeFile(state.sessionFile, content, 'utf8');
  }
}

function parseSessionFile(content: string, sessionFile: string, expectedSessionId: string): SessionFileState {
  const rawLines = content.split('\n');
  const lastContentLineIndex = findLastContentLineIndex(rawLines);
  if (lastContentLineIndex < 0) {
    throw new SessionTreeProviderError(`Session file ${sessionFile} is empty`);
  }

  const header = parseJsonLine<SessionTreeHeader>(rawLines[0] ?? '', sessionFile, HEADER_LINE_NUMBER, 'header');
  assertValidHeader(header, sessionFile, expectedSessionId);

  const entries: SessionTreeEntry[] = [];
  let firstMiddleBlankLine: number | undefined;

  for (let index = 1; index <= lastContentLineIndex; index += 1) {
    const line = rawLines[index] ?? '';
    const lineNumber = index + 1;

    if (line.trim().length === 0) {
      firstMiddleBlankLine ??= lineNumber;
      continue;
    }

    entries.push(parseJsonLine<SessionTreeEntry>(
      line,
      sessionFile,
      lineNumber,
      'entry',
    ));
  }

  if (firstMiddleBlankLine !== undefined) {
    throw new SessionTreeProviderError(
      `Invalid blank line in session file ${sessionFile} at line ${firstMiddleBlankLine}: blank lines are only allowed at the end`,
    );
  }

  return {
    header: cloneValue(header),
    entries: cloneEntries(entries),
    sessionFile,
  };
}

function findLastContentLineIndex(lines: string[]): number {
  for (let index = lines.length - 1; index >= 0; index -= 1) {
    if ((lines[index] ?? '').trim().length > 0) {
      return index;
    }
  }

  return -1;
}

function serializeSessionFile(header: SessionTreeHeader, entries: SessionTreeEntry[]): string {
  const lines = [JSON.stringify(header), ...entries.map(entry => JSON.stringify(entry))];
  return `${lines.join('\n')}\n`;
}

function parseJsonLine<T>(
  line: string,
  sessionFile: string,
  lineNumber: number,
  kind: 'header' | 'entry',
): T {
  try {
    return JSON.parse(line) as T;
  } catch (error) {
    const reason = error instanceof Error ? error.message : 'Unknown JSON parse error';
    throw new SessionTreeProviderError(
      `Invalid ${kind} JSON in session file ${sessionFile} at line ${lineNumber}: ${reason}`,
    );
  }
}

function assertValidHeader(
  header: SessionTreeHeader,
  sessionFile: string,
  expectedSessionId: string,
): void {
  const context = `session file ${sessionFile} at line ${HEADER_LINE_NUMBER}`;

  if (header.type !== 'session') {
    throw new SessionTreeProviderError(`Invalid header in ${context}: field type must be "session"`);
  }

  if (header.version !== SESSION_FILE_VERSION) {
    throw new SessionTreeProviderError(
      `Invalid header in ${context}: field version must be ${SESSION_FILE_VERSION}`,
    );
  }

  if (typeof header.id !== 'string' || header.id.length === 0) {
    throw new SessionTreeProviderError(`Invalid header in ${context}: field id must be a non-empty string`);
  }

  if (typeof header.timestamp !== 'string' || header.timestamp.length === 0) {
    throw new SessionTreeProviderError(
      `Invalid header in ${context}: field timestamp must be a non-empty string`,
    );
  }

  if (typeof header.cwd !== 'string' || header.cwd.length === 0) {
    throw new SessionTreeProviderError(`Invalid header in ${context}: field cwd must be a non-empty string`);
  }

  if (header.id !== expectedSessionId) {
    throw new SessionTreeProviderError(
      `Invalid header in ${context}: requested session ${expectedSessionId} but file stores ${header.id}`,
    );
  }
}

function buildSessionFromState(state: SessionFileState): SessionTreeSession {
  const rootEntryId = state.entries.find(entry => entry.parentId === null)?.id ?? null;
  const leafEntryId = resolveLeafEntryId(state.header, state.entries);
  const updatedAt = state.header.updatedAt
    ?? state.entries[state.entries.length - 1]?.timestamp
    ?? state.header.timestamp;

  return cloneSession({
    sessionId: state.header.id,
    provider: 'file',
    rootEntryId,
    leafEntryId,
    cwd: state.header.cwd,
    sessionFile: state.sessionFile,
    parentSessionId: state.header.parentSession,
    agentId: state.header.agentId,
    userId: state.header.userId,
    title: state.header.title,
    status: state.header.status ?? 'active',
    metadata: cloneValue(state.header.metadata),
    createdAt: state.header.timestamp,
    updatedAt,
  });
}

function resolveLeafEntryId(header: SessionTreeHeader, entries: SessionTreeEntry[]): string | null {
  if (Object.prototype.hasOwnProperty.call(header, 'leafId') && header.leafId === null) {
    return null;
  }

  if (header.leafId && entries.some(entry => entry.id === header.leafId)) {
    return header.leafId;
  }

  return entries[entries.length - 1]?.id ?? null;
}

function assertImmutableEntryPatch(
  entry: SessionTreeEntry,
  patch: Record<string, unknown>,
): void {
  const immutableKeys: Array<keyof Pick<SessionTreeEntry, 'id' | 'type' | 'parentId' | 'timestamp'>> = [
    'id',
    'type',
    'parentId',
    'timestamp',
  ];

  for (const key of immutableKeys) {
    if (Object.prototype.hasOwnProperty.call(patch, key) && patch[key] !== entry[key]) {
      throw new SessionTreeProviderError(`Entry field ${key} is immutable`);
    }
  }

  const allowedKeys = getAllowedPatchKeys(entry.type);
  for (const key of Object.keys(patch)) {
    if (!allowedKeys.has(key)) {
      throw new SessionTreeProviderError(
        `Entry field ${key} cannot be updated for entry type ${entry.type}`,
      );
    }
  }
}

function getAllowedPatchKeys(entryType: SessionTreeEntry['type']): ReadonlySet<string> {
  switch (entryType) {
    case 'message':
      return new Set(['message', 'metadata']);
    case 'custom_message':
    case 'subagent_batch':
    case 'subagent_result':
      return new Set(['content', 'details', 'display', 'metadata']);
    case 'custom':
      return new Set(['data', 'metadata']);
    case 'thinking_level_change':
    case 'model_change':
    case 'compaction':
    case 'branch_summary':
    case 'label':
    case 'session_info':
      return new Set(['metadata']);
    default: {
      const exhaustive: never = entryType;
      return exhaustive;
    }
  }
}

function cloneSnapshot(snapshot: SessionTreeSnapshot): SessionTreeSnapshot {
  return {
    session: cloneSession(snapshot.session),
    entries: cloneEntries(snapshot.entries),
    activePath: cloneEntries(snapshot.activePath),
    childrenByParentId: cloneValue(snapshot.childrenByParentId),
    labelsByEntryId: cloneValue(snapshot.labelsByEntryId),
    runtimeContext: cloneValue(snapshot.runtimeContext),
  };
}

function cloneEntries(entries: SessionTreeEntry[]): SessionTreeEntry[] {
  return entries.map(entry => cloneEntry(entry));
}

function cloneEntry(entry: SessionTreeEntry): SessionTreeEntry {
  return cloneValue(entry);
}

function cloneEntryPatch(patch: UpdateSessionTreeEntryPatch): UpdateSessionTreeEntryPatch {
  return cloneValue(patch);
}

function cloneSession(session: SessionTreeSession): SessionTreeSession {
  return cloneValue(session);
}

function cloneValue<T>(value: T): T {
  if (value === undefined) {
    return value;
  }
  return JSON.parse(JSON.stringify(value)) as T;
}

function assertJsonSerializable(value: unknown, context: string, path = '$'): void {
  if (value === undefined) {
    return;
  }

  if (value === null) {
    return;
  }

  const valueType = typeof value;
  if (valueType === 'string' || valueType === 'boolean') {
    return;
  }

  if (valueType === 'number') {
    if (Number.isFinite(value)) {
      return;
    }
    throw new SessionTreeProviderError(`${context} contains a non-finite number at ${path}`);
  }

  if (valueType === 'bigint' || valueType === 'function' || valueType === 'symbol') {
    throw new SessionTreeProviderError(`${context} contains non-JSON-serializable ${valueType} at ${path}`);
  }

  if (Array.isArray(value)) {
    value.forEach((item, index) => {
      if (item === undefined) {
        throw new SessionTreeProviderError(`${context} contains undefined array item at ${path}[${index}]`);
      }
      assertJsonSerializable(item, context, `${path}[${index}]`);
    });
    return;
  }

  if (valueType === 'object') {
    const prototype = Object.getPrototypeOf(value);
    if (prototype !== Object.prototype && prototype !== null) {
      const constructorName = prototype?.constructor?.name ?? 'object';
      throw new SessionTreeProviderError(
        `${context} contains non-plain object ${constructorName} at ${path}`,
      );
    }

    for (const [key, nestedValue] of Object.entries(value as Record<string, unknown>)) {
      if (nestedValue === undefined) {
        continue;
      }
      assertJsonSerializable(nestedValue, context, `${path}.${key}`);
    }
    return;
  }

  throw new SessionTreeProviderError(`${context} contains unsupported value at ${path}`);
}

function isNodeError(error: unknown): error is NodeJS.ErrnoException {
  return typeof error === 'object' && error !== null && 'code' in error;
}

function createDefaultSessionId(): string {
  return randomUUID();
}

function createDefaultEntryId(existingIds?: { has(id: string): boolean }): string {
  while (true) {
    const id = `entry_${randomUUID().replace(/-/g, '').slice(0, 12)}`;
    if (!existingIds?.has(id)) {
      return id;
    }
  }
}
