import Dexie, { type EntityTable, type Table } from 'dexie';
import { DSL_VERSION } from '@openmaic/dsl';
import type {
  Scene,
  SceneType,
  SceneContent,
  Whiteboard,
  VideoManifest,
  GeneratedAgentConfig,
} from '@/lib/types/stage';
import type { Action } from '@/lib/types/action';
import type {
  SessionType,
  SessionStatus,
  SessionConfig,
  ToolCallRecord,
  ToolCallRequest,
  ChatSession,
} from '@/lib/types/chat';
import type { SceneOutline } from '@/lib/types/generation';
import type { VoiceDesign } from '@/lib/audio/voice-design';
import type { UIMessage } from 'ai';
import type { AgentEditSessionRecord } from '@/lib/agent/client/agent-edit-session-types';
import { createLogger } from '@/lib/logger';
import { beginStageRuntimeDeletionSafely, getRuntimeStore } from '@/lib/runtime/store';
import type { RuntimeStore } from '@openmaic/storage';
import {
  withRuntimeStorageExclusiveLock,
  withRuntimeStorageExclusiveLockUntilSettled,
  withRuntimeStorageSharedLock,
} from './chat-storage-lock';
import type { ChatStorageOptions } from './chat-storage';
import type { AppDocument } from '@/lib/document-store';
import { BrowserKVStore } from '@openmaic/storage';
import { clearAssetPool } from '@/lib/media/asset-pool';

const log = createLogger('Database');

/**
 * Legacy Snapshot type for undo/redo functionality
 * Used by useSnapshotStore
 */
export interface Snapshot {
  id?: number;
  index: number;
  slides: Scene[];
}

/**
 * MAIC Local Database
 *
 * Uses IndexedDB to store all user data locally
 * - Does not delete expired data; all data is stored permanently
 * - Uses a fixed database name
 * - Supports multi-course management
 */

// ==================== Database Table Type Definitions ====================

/**
 * Stage table - Course basic info
 */
export interface StageRecord {
  id: string; // Primary key
  name: string;
  description?: string;
  createdAt: number; // timestamp
  updatedAt: number; // timestamp
  languageDirective?: string;
  style?: string;
  currentSceneId?: string;
  agentIds?: string[]; // Agent IDs selected at creation time
  videoManifest?: VideoManifest; // Generated video request manifest; non-indexed
  interactiveMode?: boolean; // Interactive Mode flag; non-indexed
  taskEngineMode?: boolean; // Vocational Task Engine flag; non-indexed
  generatedAgentConfigs?: GeneratedAgentConfig[]; // Editor-authored agent roster snapshot
}

/**
 * Folder table - User-created folders for grouping courses.
 *
 * Folder membership is device-local organization metadata, not part of the
 * course document itself (which is owned by the `@openmaic/storage`
 * DocumentStore in a separate database). It lives in this Dexie database
 * alongside the legacy tables. See {@link StageFolderMembership}.
 */
export interface FolderRecord {
  id: string; // Primary key
  name: string;
  order: number; // Sort order
  createdAt: number; // timestamp
  updatedAt: number; // timestamp
}

/**
 * Stage→folder membership mapping. `stageId` is the primary key so each course
 * has at most one row; a missing row (or `folderId === undefined`) means the
 * course is unfiled. This is intentionally separate from both the legacy
 * `stages` table (a migration mirror that nothing writes) and the
 * DocumentStore stage row (version-independent document content), so folder
 * grouping never touches document semantics.
 */
export interface StageFolderMembership {
  stageId: string; // Primary key (FK -> DocumentStore stage id)
  folderId?: string; // FK -> folders.id; undefined = unfiled
  updatedAt: number; // timestamp
}

/**
 * Scene table - Scene/page data
 */
export interface SceneRecord {
  id: string; // Primary key
  stageId: string; // Foreign key -> stages.id
  type: SceneType;
  title: string;
  order: number; // Display order
  content: SceneContent; // Stored as JSON
  actions?: Action[]; // Stored as JSON
  whiteboard?: Whiteboard[]; // Stored as JSON
  createdAt: number;
  updatedAt: number;
}

/**
 * AudioFile table - Audio files (TTS)
 */
export interface AudioFileRecord {
  id: string; // Primary key (audioId)
  /** Stage ownership index. Absent on legacy rows; document walking remains their fallback. */
  stageId?: string;
  blob: Blob; // Audio binary data
  duration?: number; // Duration (seconds)
  format: string; // mp3, wav, etc.
  text?: string; // Corresponding text content
  voice?: string; // Voice used
  createdAt: number;
  ossKey?: string; // Full CDN URL for this audio blob
}

/**
 * ImageFile table - Image files
 */
export interface ImageFileRecord {
  id: string; // Primary key
  blob: Blob | ArrayBuffer; // Image binary data
  filename: string; // Original filename
  mimeType: string; // image/png, image/jpeg, etc.
  size: number; // File size (bytes)
  createdAt: number;
}

/**
 * ChatSession table - Chat session data
 */
export interface ChatSessionRecord {
  id: string; // PK (session id)
  stageId: string; // FK -> stages.id
  type: SessionType;
  title: string;
  status: SessionStatus;
  messages: UIMessage[]; // JSON-safe serialized messages
  config: SessionConfig;
  toolCalls: ToolCallRecord[];
  pendingToolCalls: ToolCallRequest[];
  createdAt: number;
  updatedAt: number;
  sceneId?: string;
  lastActionIndex?: number;
}

/**
 * PlaybackState table - Playback state snapshot (at most one per stage)
 */
export interface PlaybackStateRecord {
  stageId: string; // PK
  sceneIndex: number;
  actionIndex: number;
  consumedDiscussions: string[];
  updatedAt: number;
}

/**
 * StageOutlines table - Persisted outlines for resume-on-refresh
 */
export interface StageOutlinesRecord {
  stageId: string; // Primary key (FK -> stages.id)
  outlines: SceneOutline[];
  // True once generation finished for this stage. Gates resume-on-mount so an
  // edited (e.g. slide-deleted) finished deck is not treated as "interrupted"
  // and regenerated. Optional for backward compat with pre-existing records.
  generationComplete?: boolean;
  createdAt: number;
  updatedAt: number;
}

/**
 * MediaFile table - AI-generated media files (images/videos)
 */
export interface MediaFileRecord {
  // Compound key: `${stageId}:${mediaRef}`. Successful and failed rows use
  // the same reference space (allocated id after allocation, legacy ref before it).
  id: string;
  stageId: string; // FK → stages.id
  /** Original gen_* reference retained after allocation for reload reconciliation. */
  placeholderRef?: string;
  type: 'image' | 'video';
  blob: Blob; // Media binary
  mimeType: string; // image/png, video/mp4
  size: number;
  poster?: Blob; // Video thumbnail blob
  prompt: string; // Original prompt (for retry)
  params: string; // JSON-serialized generation params
  error?: string; // If set, this is a failed task (blob is empty placeholder)
  errorCode?: string; // Structured error code (e.g. 'CONTENT_SENSITIVE')
  ossKey?: string; // Full CDN URL for this media blob
  posterOssKey?: string; // Full CDN URL for the poster blob
  createdAt: number;
}

/**
 * GeneratedAgent table - AI-generated agent profiles.
 *
 * LEGACY. The roster now persists on the stage document
 * (`stage.generatedAgentConfigs`); this table is kept only as a lazy-migration
 * source for classrooms whose roster (or voice fields) predate the
 * document-embedded model. Production access is migration reads plus deletion
 * hygiene: `deleteStageData` clears a deleted stage's rows (as does the
 * deprecated `deleteStageWithRelatedData` cascade). Nothing writes new rows;
 * do not add writers.
 */
export interface GeneratedAgentRecord {
  id: string; // PK: agent ID (e.g. "gen-abc123")
  stageId: string; // FK -> stages.id
  name: string;
  role: string; // 'teacher' | 'assistant' | 'student'
  persona: string;
  avatar: string;
  color: string;
  priority: number;
  voiceDesign?: VoiceDesign; // 3-layer vocal descriptor for auto voice
  createdAt: number;
}

/**
 * VoiceProfile table - Browser-local TTS voice profiles
 */
export interface VoiceProfileRecord {
  id: string;
  providerId: string;
  kind: 'prompt' | 'clone';
  name: string;
  voicePrompt?: string;
  promptText?: string;
  referenceAudio?: Blob;
  referenceAudioName?: string;
  referenceAudioMimeType?: string;
  createdAt: number;
  updatedAt: number;
}

/**
 * Cached reference clip for a registered auto voice (any TTS provider). The
 * clip is the source of truth; the deterministic `voiceId` is its key, enabling
 * register-on-invalid re-registration after backend GC/restart.
 */
export interface AutoVoiceCacheRecord {
  voiceId: string;
  referenceAudio: Blob;
  mimeType: string;
  updatedAt: number;
}

/** Build the compound primary key for mediaFiles: `${stageId}:${elementId}` */
export function mediaFileKey(stageId: string, elementId: string): string {
  return `${stageId}:${elementId}`;
}

// ==================== Database Definition ====================

const DATABASE_NAME = 'MAIC-Database';
const _DATABASE_VERSION = 17;

/**
 * MAIC Database Instance
 */
class MAICDatabase extends Dexie {
  // Table definitions
  stages!: EntityTable<StageRecord, 'id'>;
  scenes!: EntityTable<SceneRecord, 'id'>;
  audioFiles!: EntityTable<AudioFileRecord, 'id'>;
  imageFiles!: EntityTable<ImageFileRecord, 'id'>;
  snapshots!: EntityTable<Snapshot, 'id'>; // Undo/redo snapshots (legacy)
  chatSessions!: EntityTable<ChatSessionRecord, 'id'>;
  chatRestoreStaging!: Table<ChatSessionRecord, [string, string]>;
  playbackState!: EntityTable<PlaybackStateRecord, 'stageId'>;
  stageOutlines!: EntityTable<StageOutlinesRecord, 'stageId'>;
  mediaFiles!: EntityTable<MediaFileRecord, 'id'>;
  generatedAgents!: EntityTable<GeneratedAgentRecord, 'id'>;
  voiceProfiles!: EntityTable<VoiceProfileRecord, 'id'>;
  autoVoiceCache!: EntityTable<AutoVoiceCacheRecord, 'voiceId'>;
  agentEditSessions!: EntityTable<AgentEditSessionRecord, 'id'>;
  folders!: EntityTable<FolderRecord, 'id'>;
  stageFolders!: EntityTable<StageFolderMembership, 'stageId'>;

  constructor() {
    super(DATABASE_NAME);

    // Version 1: Initial schema
    this.version(1).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      // Previously had: messages, participants, discussions, sceneSnapshots
    });

    // Version 2: Remove unused tables
    this.version(2).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      // Delete removed tables
      messages: null,
      participants: null,
      discussions: null,
      sceneSnapshots: null,
    });

    // Version 3: Add chatSessions and playbackState tables
    this.version(3).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      chatSessions: 'id, stageId, [stageId+createdAt]',
      playbackState: 'stageId',
    });

    // Version 4: Add stageOutlines table for resume-on-refresh
    this.version(4).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      chatSessions: 'id, stageId, [stageId+createdAt]',
      playbackState: 'stageId',
      stageOutlines: 'stageId',
    });

    // Version 5: Add mediaFiles table for async media generation
    this.version(5).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      chatSessions: 'id, stageId, [stageId+createdAt]',
      playbackState: 'stageId',
      stageOutlines: 'stageId',
      mediaFiles: 'id, stageId, [stageId+type]',
    });

    // Version 6: Fix mediaFiles primary key — use compound key stageId:elementId
    // to prevent cross-course collisions (gen_img_1 is NOT globally unique)
    this.version(6)
      .stores({
        stages: 'id, updatedAt',
        scenes: 'id, stageId, order, [stageId+order]',
        audioFiles: 'id, createdAt',
        imageFiles: 'id, createdAt',
        snapshots: '++id',
        chatSessions: 'id, stageId, [stageId+createdAt]',
        playbackState: 'stageId',
        stageOutlines: 'stageId',
        mediaFiles: 'id, stageId, [stageId+type]',
      })
      .upgrade(async (tx) => {
        const table = tx.table('mediaFiles');
        const allRecords = await table.toArray();
        for (const rec of allRecords) {
          const newKey = `${rec.stageId}:${rec.id}`;
          // Skip if already migrated (idempotent)
          if (rec.id.includes(':')) continue;
          await table.delete(rec.id);
          await table.put({ ...rec, id: newKey });
        }
      });

    // Version 7: Add ossKey fields to mediaFiles and audioFiles for OSS storage plugin
    // Non-indexed optional fields — Dexie handles these transparently.
    this.version(7).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      chatSessions: 'id, stageId, [stageId+createdAt]',
      playbackState: 'stageId',
      stageOutlines: 'stageId',
      mediaFiles: 'id, stageId, [stageId+type]',
    });

    // Version 8: Add generatedAgents table for AI-generated agent profiles
    this.version(8).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      chatSessions: 'id, stageId, [stageId+createdAt]',
      playbackState: 'stageId',
      stageOutlines: 'stageId',
      mediaFiles: 'id, stageId, [stageId+type]',
      generatedAgents: 'id, stageId',
    });

    // Version 9: Migrate legacy `language` field to `languageDirective`
    // Old stages stored a BCP-47 locale code (e.g. "zh-CN"); new code expects a
    // natural-language directive. Convert known locales and drop the old field.
    const LOCALE_TO_DIRECTIVE: Record<string, string> = {
      'zh-CN': 'Deliver the entire course in Chinese (Simplified, zh-CN).',
      'en-US': 'Deliver the entire course in English (en-US).',
      'ja-JP': 'Deliver the entire course in Japanese (ja-JP).',
      'ru-RU': 'Deliver the entire course in Russian (ru-RU).',
    };
    this.version(9)
      .stores({
        stages: 'id, updatedAt',
        scenes: 'id, stageId, order, [stageId+order]',
        audioFiles: 'id, createdAt',
        imageFiles: 'id, createdAt',
        snapshots: '++id',
        chatSessions: 'id, stageId, [stageId+createdAt]',
        playbackState: 'stageId',
        stageOutlines: 'stageId',
        mediaFiles: 'id, stageId, [stageId+type]',
        generatedAgents: 'id, stageId',
      })
      .upgrade(async (tx) => {
        const table = tx.table('stages');
        await table.toCollection().modify((stage: Record<string, unknown>) => {
          const lang = stage.language as string | undefined;
          if (lang && !stage.languageDirective) {
            stage.languageDirective =
              LOCALE_TO_DIRECTIVE[lang] || `Deliver the entire course in ${lang}.`;
          }
          delete stage.language;
        });
      });

    // Version 10: Add browser-local voice profiles for serverless TTS voice storage.
    this.version(10).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      chatSessions: 'id, stageId, [stageId+createdAt]',
      playbackState: 'stageId',
      stageOutlines: 'stageId',
      mediaFiles: 'id, stageId, [stageId+type]',
      generatedAgents: 'id, stageId',
      voiceProfiles: 'id, providerId, kind, updatedAt',
    });

    // Version 11: Add auto-voice reference-clip cache (provider-neutral register-by-id).
    this.version(11).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      chatSessions: 'id, stageId, [stageId+createdAt]',
      playbackState: 'stageId',
      stageOutlines: 'stageId',
      mediaFiles: 'id, stageId, [stageId+type]',
      generatedAgents: 'id, stageId',
      voiceProfiles: 'id, providerId, kind, updatedAt',
      autoVoiceCache: 'voiceId, updatedAt',
    });

    // Version 12: Add agentEditSessions — multi-session AI-editing conversation
    // history per stage (replaces the single-thread localStorage store).
    this.version(12).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      chatSessions: 'id, stageId, [stageId+createdAt]',
      playbackState: 'stageId',
      stageOutlines: 'stageId',
      mediaFiles: 'id, stageId, [stageId+type]',
      generatedAgents: 'id, stageId',
      voiceProfiles: 'id, providerId, kind, updatedAt',
      autoVoiceCache: 'voiceId, updatedAt',
      agentEditSessions: 'id, stageId, [stageId+updatedAt]',
    });

    // Version 13 briefly added chatStorageLocks on the draft chat cutover
    // branch. Advance past it and remove the abandoned lease table so database
    // versions stay monotonic for anyone who opened that intermediate build.
    this.version(14).stores({
      stages: 'id, updatedAt',
      scenes: 'id, stageId, order, [stageId+order]',
      audioFiles: 'id, createdAt',
      imageFiles: 'id, createdAt',
      snapshots: '++id',
      chatSessions: 'id, stageId, [stageId+createdAt]',
      playbackState: 'stageId',
      stageOutlines: 'stageId',
      mediaFiles: 'id, stageId, [stageId+type]',
      generatedAgents: 'id, stageId',
      voiceProfiles: 'id, providerId, kind, updatedAt',
      autoVoiceCache: 'voiceId, updatedAt',
      agentEditSessions: 'id, stageId, [stageId+updatedAt]',
      chatStorageLocks: null,
    });

    // Version 15: backup restore staging must preserve chat IDs that are reused
    // in different stage partitions; the legacy chat table is keyed by id only.
    this.version(15).stores({
      chatRestoreStaging: '[stageId+id], stageId, [stageId+createdAt]',
    });

    // Version 16: make newly-written audio independently reclaimable by stage.
    // Legacy rows remain valid and are found through speech-action references.
    this.version(16).stores({
      audioFiles: 'id, stageId, createdAt',
    });

    // Version 17: Course folders — group courses into user-created folders.
    // `folders` holds folder metadata; `stageFolders` maps each course (by
    // DocumentStore stage id) to its folder. Neither touches the document
    // aggregate: folder grouping is device-local organization metadata kept in
    // this Dexie database alongside the legacy tables, so an existing course
    // with no membership row is simply unfiled (no upgrade callback needed).
    this.version(17).stores({
      folders: 'id, order',
      stageFolders: 'stageId, folderId',
    });
  }
}

// Create database instance
export const db = new MAICDatabase();

// ==================== Helper Functions ====================

/**
 * Initialize database
 * Call at application startup
 */
export async function initDatabase(): Promise<void> {
  try {
    await db.open();
    // Request persistent storage to prevent browser from evicting IndexedDB
    // under storage pressure (large media blobs can trigger LRU cleanup)
    void navigator.storage?.persist?.();
    log.info('Database initialized successfully');
  } catch (error) {
    log.error('Failed to initialize database:', error);
    throw error;
  }
}

/**
 * Clear database (optional)
 * Use with caution: deletes all data
 */
export async function clearDatabase(runtimeStore?: RuntimeStore): Promise<void> {
  // Clear the whole runtime database first, including rows orphaned by an
  // earlier best-effort stage deletion. This user-requested destructive action
  // must fail loud: reporting success while runtime data remains is misleading.
  await withRuntimeStorageExclusiveLock(async () => {
    const { bumpGeneration } = await import('@/lib/document-store/storage-generation');
    await bumpGeneration();
    await (runtimeStore ?? getRuntimeStore()).deleteAllRuntime();
    await deleteAllDocuments();
    await clearDocumentStoreKeys();
    await db.delete();
    await clearAssetPool();
  });
  log.info('Database cleared');
}

/** Delete every aggregate without bypassing the document store's open connection. */
export async function deleteAllDocuments(): Promise<void> {
  const { getDocumentStore } = await import('@/lib/document-store');
  const store = getDocumentStore();
  const documents = await store.listDocuments();
  await Promise.all(documents.map((document) => store.deleteDocument(document.id)));
}

/** Remove device-local metadata owned by the document cutover. */
export async function clearDocumentStoreKeys(): Promise<void> {
  if (typeof localStorage === 'undefined') return;
  const kv = new BrowserKVStore();
  for (const prefix of ['document-migration:', 'editor-current-scene:']) {
    const keys = await kv.keys(prefix, 'device');
    await Promise.all(keys.map((key) => kv.remove(key, 'device')));
  }
}

function toChatSessionRecord(stageId: string, session: ChatSession): ChatSessionRecord {
  return {
    id: session.id,
    stageId,
    type: session.type,
    title: session.title,
    status: session.status,
    messages: session.messages,
    config: session.config,
    toolCalls: session.toolCalls,
    pendingToolCalls: session.pendingToolCalls,
    createdAt: session.createdAt,
    updatedAt: session.updatedAt,
    sceneId: session.sceneId,
    lastActionIndex: session.lastActionIndex,
  };
}

/**
 * Export database contents (for backup)
 */
export async function exportDatabase(chatOptions: ChatStorageOptions = {}): Promise<{
  documents: AppDocument[];
  chatSessions: ChatSessionRecord[];
  playbackState: PlaybackStateRecord[];
}> {
  const {
    accessDocument,
    canonicalizeLegacyOutline,
    canonicalizeLegacyScene,
    canonicalizeLegacyStage,
    getDocumentStore,
    getLegacyDocumentStore,
  } = await import('@/lib/document-store');
  const documentStore = getDocumentStore();
  // Backups must not strand courses that have not yet been opened since cutover.
  // Route them through the normal lazy-migration seam before enumerating aggregates.
  const legacyStages = await getLegacyDocumentStore().listStages();
  await Promise.all(legacyStages.map((stage) => accessDocument(stage.id)));
  const summaries = await documentStore.listDocuments();
  const storedDocuments = (
    await Promise.all(summaries.map((summary) => documentStore.loadDocument(summary.id)))
  ).filter((document): document is AppDocument => document !== null);
  const storedIds = new Set(summaries.map((summary) => summary.id));
  const legacyOnlyDocuments = (
    await Promise.all(
      legacyStages
        .filter((stage) => !storedIds.has(stage.id))
        .map(async (stage): Promise<AppDocument | null> => {
          const snapshot = await getLegacyDocumentStore().read(stage.id);
          if (!snapshot) return null;
          const { stage: canonicalStage } = canonicalizeLegacyStage(snapshot.stage);
          const document: AppDocument = {
            stage: canonicalStage,
            scenes: snapshot.scenes.map(canonicalizeLegacyScene).sort((a, b) => a.order - b.order),
            dslVersion: DSL_VERSION,
          };
          if (snapshot.outline) document.outline = canonicalizeLegacyOutline(snapshot.outline);
          return document;
        }),
    )
  ).filter((document): document is AppDocument => document !== null);
  const documents = [...storedDocuments, ...legacyOnlyDocuments];
  const legacyChatMap = new Map<string, ChatSessionRecord>();
  for (const session of [
    ...(await db.chatSessions.toArray()),
    ...(await db.chatRestoreStaging.toArray()),
  ]) {
    legacyChatMap.set(JSON.stringify([session.stageId, session.id]), session);
  }
  const legacyChats = [...legacyChatMap.values()];
  const { loadChatSessions } = await import('./chat-storage');
  // Chat sessions live on the learner RuntimeStore — an independent seam from
  // document migration — so legacy-only documents can still own runtime chat
  // history. Enumerate chats for EVERY exported document, not just stored ones.
  // The legacy chat rows are collected separately above by direct table reads,
  // so the RuntimeStore load gets an empty legacy source and never needs the
  // cross-realm migration lock. It may still finalize a pending restore marker;
  // `observe: false` only keeps this export from changing partition memos.
  const runtimeChats = (
    await Promise.all(
      documents.map(async ({ stage }) =>
        (
          await loadChatSessions(stage.id, {
            legacyStore: { load: async () => [], clear: async () => {} },
            ...chatOptions,
            fallbackToLegacyOnError: false,
            observe: false,
          })
        ).map((session) => toChatSessionRecord(stage.id, session)),
      ),
    )
  ).flat();
  const runtimeChatKeys = new Set(
    runtimeChats.map((session) => JSON.stringify([session.stageId, session.id])),
  );

  return {
    documents,
    chatSessions: [
      ...runtimeChats,
      ...legacyChats.filter(
        (session) => !runtimeChatKeys.has(JSON.stringify([session.stageId, session.id])),
      ),
    ],
    playbackState: await db.playbackState.toArray(),
  };
}

/**
 * Import database contents (for restoring backups)
 */
export async function importDatabase(
  data: {
    documents?: AppDocument[];
    stages?: StageRecord[];
    scenes?: SceneRecord[];
    chatSessions?: ChatSessionRecord[];
    playbackState?: PlaybackStateRecord[];
  },
  chatOptions: ChatStorageOptions = {},
): Promise<void> {
  const { canonicalizeLegacyScene, canonicalizeLegacyStage, mutateDocument } =
    await import('@/lib/document-store');
  const legacyScenesByStage = new Map<string, SceneRecord[]>();
  for (const scene of data.scenes ?? []) {
    const scenes = legacyScenesByStage.get(scene.stageId) ?? [];
    scenes.push(scene);
    legacyScenesByStage.set(scene.stageId, scenes);
  }
  const documents =
    data.documents ??
    (data.stages ?? []).map((legacyStage): AppDocument => {
      const { stage } = canonicalizeLegacyStage(legacyStage);
      return {
        stage,
        scenes: (legacyScenesByStage.get(legacyStage.id) ?? [])
          .map(canonicalizeLegacyScene)
          .sort((a, b) => a.order - b.order),
      };
    });
  const importedDocuments: Array<{
    id: string;
    preImage: AppDocument | null;
    wasDeleted: boolean;
  }> = [];
  const importedCurrentScenes: Array<{ key: string; preImage: unknown | null }> = [];
  const kv = new BrowserKVStore();
  const { isStageDeleted, markStageDeleted, unmarkStageDeleted } = await import('./deleted-stages');

  try {
    for (const document of documents) {
      // Record the pre-import deletion state alongside the document pre-image:
      // a failed import rolls the document back, so it must roll this back too.
      const wasDeleted = isStageDeleted(document.stage.id);
      await mutateDocument(document.stage.id, async (_existing, store) => {
        const preImage = (await store.loadDocument(document.stage.id)) as AppDocument | null;
        await store.saveDocument(document);
        importedDocuments.push({ id: document.stage.id, preImage, wasDeleted });
      });
      // Explicit document (re)creation: a backup may restore a stage deleted
      // earlier this session under the same id. Lift the deleted flag so later
      // edits of the restored document persist instead of being dropped. (The
      // deletion epoch stays bumped, so pre-delete in-flight writes remain
      // fenced off the restored document.)
      unmarkStageDeleted(document.stage.id);
    }
    for (const legacyStage of data.stages ?? []) {
      if (legacyStage.currentSceneId !== undefined) {
        const key = `editor-current-scene:${legacyStage.id}`;
        const preImage = await kv.get<unknown>(key, 'device');
        await kv.set(
          key,
          { sceneId: legacyStage.currentSceneId, updatedAt: new Date().toISOString() },
          'device',
        );
        importedCurrentScenes.push({ key, preImage });
      }
    }

    await withRuntimeStorageSharedLock(async () => {
      const restoredChatStageIds =
        data.chatSessions === undefined
          ? []
          : [
              ...new Set([
                ...documents.map((document) => document.stage.id),
                ...data.chatSessions.map((session) => session.stageId),
              ]),
            ];
      const restoreRows = () =>
        db.transaction(
          'rw',
          [db.stages, db.scenes, db.chatSessions, db.chatRestoreStaging, db.playbackState],
          async () => {
            if (data.chatSessions) {
              for (const stageId of restoredChatStageIds) {
                await db.chatSessions.where('stageId').equals(stageId).delete();
                await db.chatRestoreStaging.where('stageId').equals(stageId).delete();
              }
              await db.chatRestoreStaging.bulkPut(data.chatSessions);
            }
            if (data.playbackState) await db.playbackState.bulkPut(data.playbackState);
          },
        );
      if (data.chatSessions !== undefined) {
        // RuntimeStore cannot join Dexie's transaction. Keep a rollback image
        // until durable restore markers have been created for every affected
        // runtime partition; marker creation failure must not commit half an
        // imported backup.
        const rollbackImage = await db.transaction(
          'r',
          [db.stages, db.scenes, db.chatSessions, db.chatRestoreStaging, db.playbackState],
          async () => ({
            chatSessions: (
              await Promise.all(
                restoredChatStageIds.map((stageId) =>
                  db.chatSessions.where('stageId').equals(stageId).toArray(),
                ),
              )
            ).flat(),
            chatRestoreStaging: (
              await Promise.all(
                restoredChatStageIds.map((stageId) =>
                  db.chatRestoreStaging.where('stageId').equals(stageId).toArray(),
                ),
              )
            ).flat(),
            playbackState: (
              await db.playbackState.bulkGet(
                (data.playbackState ?? []).map((playback) => playback.stageId),
              )
            ).filter((playback): playback is PlaybackStateRecord => playback !== undefined),
          }),
        );
        const rollbackRows = () =>
          db.transaction(
            'rw',
            [db.stages, db.scenes, db.chatSessions, db.chatRestoreStaging, db.playbackState],
            async () => {
              for (const stageId of restoredChatStageIds) {
                await db.chatSessions.where('stageId').equals(stageId).delete();
                await db.chatRestoreStaging.where('stageId').equals(stageId).delete();
              }
              await db.playbackState.bulkDelete(
                (data.playbackState ?? []).map((playback) => playback.stageId),
              );
              await db.chatSessions.bulkPut(rollbackImage.chatSessions);
              await db.chatRestoreStaging.bulkPut(rollbackImage.chatRestoreStaging);
              await db.playbackState.bulkPut(rollbackImage.playbackState);
            },
          );
        const { restoreChatSessionsFromBackup } = await import('./chat-storage');
        await restoreChatSessionsFromBackup(restoredChatStageIds, restoreRows, {
          ...chatOptions,
          globalLockHeld: true,
          rollbackLegacyRows: rollbackRows,
        });
      } else {
        await restoreRows();
      }
      log.info('Database imported successfully');
    });
  } catch (error) {
    for (const { key, preImage } of importedCurrentScenes.reverse()) {
      try {
        if (preImage === null) await kv.remove(key, 'device');
        else await kv.set(key, preImage, 'device');
      } catch (rollbackError) {
        log.error(`Failed to roll back imported current-scene key ${key}:`, rollbackError);
      }
    }
    for (const { id, preImage, wasDeleted } of importedDocuments.reverse()) {
      try {
        await mutateDocument(id, async (_document, store) => {
          if (preImage) await store.saveDocument(preImage);
          else await store.deleteDocument(id);
        });
        // The rollback reinstated the pre-import world; reinstate the deletion
        // state the import lifted, or the rolled-back (absent) document would
        // stay writable and an outstanding flush could recreate it. Re-marking
        // bumps the epoch again — consistent either way, since every pre-import
        // capture is already stale. Deliberately skipped when the rollback
        // write itself failed above: the imported document then still exists,
        // and re-marking would silently drop edits to a document that is
        // present (the exact bug the lift exists to prevent).
        if (wasDeleted) markStageDeleted(id);
      } catch (rollbackError) {
        log.error(`Failed to roll back imported document ${id}:`, rollbackError);
      }
    }
    throw error;
  }
}

// ==================== Convenience Query Functions ====================

/**
 * Get all scenes for a course
 */
export async function getScenesByStageId(stageId: string): Promise<Scene[]> {
  const { accessDocument } = await import('@/lib/document-store');
  return (await accessDocument(stageId)).document?.scenes ?? [];
}

/**
 * Delete a course and all its related data
 *
 * @deprecated No production caller remains. `deleteStageData` is the primary
 * stage-deletion path; this compatibility helper retains its broader author-side cascade.
 */
export async function deleteStageWithRelatedData(stageId: string): Promise<void> {
  const { mutateDocument } = await import('@/lib/document-store');
  // storageSharedLockHeld: the cascade holds the EXCLUSIVE epoch, which
  // subsumes shared — the generation-guarded store must not re-acquire shared
  // inside it (self-deadlock against our own exclusive hold).
  await mutateDocument(
    stageId,
    async (document, store) =>
      withRuntimeStorageExclusiveLockUntilSettled(async (releaseCaller) => {
        const {
          buildStageAssetReclamationPlan,
          executeStageAssetReclamation,
          loadStageAssetInventory,
        } = await import('@/lib/media/reclaim-stage-assets');
        const deletionDocument = document ?? {
          stage: { id: stageId, name: '', createdAt: 0, updatedAt: 0 },
          scenes: [],
        };
        const inventory = await loadStageAssetInventory(deletionDocument);
        const assetPlan = buildStageAssetReclamationPlan(
          stageId,
          inventory.refs,
          inventory.mediaRows,
          inventory.audioRows,
        );
        await store.deleteDocument(stageId);
        await executeStageAssetReclamation(assetPlan, null);
        await db.transaction(
          'rw',
          [
            db.stages,
            db.scenes,
            db.chatSessions,
            db.chatRestoreStaging,
            db.playbackState,
            db.stageOutlines,
            db.generatedAgents,
            db.agentEditSessions,
          ],
          async () => {
            await db.stages.delete(stageId);
            await db.scenes.where('stageId').equals(stageId).delete();
            await db.chatSessions.where('stageId').equals(stageId).delete();
            await db.chatRestoreStaging.where('stageId').equals(stageId).delete();
            await db.playbackState.delete(stageId);
            await db.stageOutlines.delete(stageId);
            await db.generatedAgents.where('stageId').equals(stageId).delete();
            await db.agentEditSessions.where('stageId').equals(stageId).delete();
          },
        );
        // Learner-runtime data lives in a separate IndexedDB database, so it is
        // cascaded after the Dexie transaction: it cannot join it, and a runtime
        // failure must not abort it (the helper warns instead of throwing).
        const runtimeDeletion = beginStageRuntimeDeletionSafely(stageId);
        await runtimeDeletion.completion;
        releaseCaller(undefined);
        await runtimeDeletion.settlement;
      }),
    { storageSharedLockHeld: true },
  );
}

/**
 * Get all generated agents for a course
 */
export async function getGeneratedAgentsByStageId(
  stageId: string,
): Promise<GeneratedAgentRecord[]> {
  return db.generatedAgents.where('stageId').equals(stageId).toArray();
}

/**
 * Get database statistics
 */
export async function getDatabaseStats() {
  const { getDocumentStore } = await import('@/lib/document-store');
  const documents = await getDocumentStore().listDocuments();
  return {
    documents: documents.length,
    documentScenes: documents.reduce((total, document) => total + document.sceneCount, 0),
    legacyStages: await db.stages.count(),
    legacyScenes: await db.scenes.count(),
    audioFiles: await db.audioFiles.count(),
    imageFiles: await db.imageFiles.count(),
    snapshots: await db.snapshots.count(),
    chatSessions: await db.chatSessions.count(),
    playbackState: await db.playbackState.count(),
    stageOutlines: await db.stageOutlines.count(),
    mediaFiles: await db.mediaFiles.count(),
    generatedAgents: await db.generatedAgents.count(),
  };
}