export interface Query extends AsyncGenerator<SDKMessage, void> {
  // Control methods
  interrupt(): Promise<void>;
  setPermissionMode(mode: PermissionMode): Promise<void>;
  setModel(model?: string): Promise<void>;
  supportedCommands(): Promise<SlashCommand[]>;
  supportedModels(): Promise<ModelInfo[]>;
  mcpServerStatus(): Promise<McpServerStatus[]>;
}

export type Options = {
  forkSession?: boolean;           // ✅ Session forking flag
  resume?: string;                 // ✅ Resume from session ID
  resumeSessionAt?: string;        // ✅ Resume from specific message
  hooks?: Partial<Record<HookEvent, HookCallbackMatcher[]>>; // ✅ Hook system
  mcpServers?: Record<string, McpServerConfig>; // ✅ In-process MCP
  maxTurns?: number;
  permissionMode?: PermissionMode;
  canUseTool?: CanUseTool;
  // ... more options
};

export declare function tool<Schema extends ZodRawShape>(
  name: string,
  description: string,
  inputSchema: Schema,
  handler: (args, extra) => Promise<CallToolResult>
): SdkMcpToolDefinition<Schema>;

export declare function createSdkMcpServer(options: {
  name: string;
  version?: string;
  tools?: Array<SdkMcpToolDefinition<any>>;
}): McpSdkServerConfigWithInstance;

export const HOOK_EVENTS = [
  "PreToolUse", "PostToolUse", "Notification",
  "UserPromptSubmit", "SessionStart", "SessionEnd",
  "Stop", "SubagentStop", "PreCompact"
] as const;

export type HookCallback = (
  input: HookInput,
  toolUseID: string | undefined,
  options: { signal: AbortSignal }
) => Promise<HookJSONOutput>;

export type SDKMessage =
  | SDKAssistantMessage    // Model responses
  | SDKUserMessage         // User inputs
  | SDKResultMessage       // Execution results
  | SDKSystemMessage       // System info
  | SDKPartialAssistantMessage  // Streaming chunks
  | SDKCompactBoundaryMessage;  // Compaction events

// Each message has:
{
  uuid: UUID;
  session_id: string;
  // ... message-specific data
}

import { query, type Options, type SDKMessage } from '@anthropic-ai/claude-code';

class RealSessionForking {
  private sessions = new Map<string, SessionSnapshot>();

  async forkSession(baseSessionId: string): Promise<ForkedSession> {
    // 1. Capture current session state
    const snapshot = await this.captureSession(baseSessionId);

    // 2. Create forked query with SDK's forkSession
    const forkedQuery = query({
      prompt: 'Continue from fork',
      options: {
        forkSession: true,         // ✅ SDK creates new session ID
        resume: baseSessionId,      // ✅ SDK loads conversation history
        resumeSessionAt: snapshot.lastMessageId, // ✅ SDK resumes from point
      }
    });

    // 3. Extract new session ID from first message
    const firstMessage = await forkedQuery.next();
    const newSessionId = firstMessage.value?.session_id;

    // 4. Track fork relationship
    this.sessions.set(newSessionId!, {
      parentId: baseSessionId,
      forkedAt: Date.now(),
      messages: [firstMessage.value!],
    });

    return {
      sessionId: newSessionId!,
      query: forkedQuery,

      // Commit changes back to parent
      async commit() {
        const changes = await this.getChanges(newSessionId!);
        await this.applyToParent(baseSessionId, changes);
      },

      // Discard fork
      async rollback() {
        this.sessions.delete(newSessionId!);
        await forkedQuery.interrupt();
      }
    };
  }

  // Capture session state using SDK message history
  private async captureSession(sessionId: string): Promise<SessionSnapshot> {
    const snapshot = this.sessions.get(sessionId);
    if (!snapshot) throw new Error('Session not found');

    return {
      sessionId,
      lastMessageId: snapshot.messages[snapshot.messages.length - 1].uuid,
      messageCount: snapshot.messages.length,
      timestamp: Date.now(),
    };
  }

  // Track all messages for a session
  async trackSession(sessionId: string, query: AsyncGenerator<SDKMessage>) {
    const messages: SDKMessage[] = [];

    for await (const message of query) {
      messages.push(message);

      // Update snapshot after each message
      this.sessions.set(sessionId, {
        parentId: null,
        forkedAt: Date.now(),
        messages,
      });
    }
  }
}

// Usage
const forker = new RealSessionForking();

// Track original session
const originalQuery = query({ prompt: 'Start task', options: {} });
await forker.trackSession('original-session', originalQuery);

// Fork session (uses SDK's forkSession + resume)
const fork = await forker.forkSession('original-session');

// Work in fork
for await (const msg of fork.query) {
  console.log('Fork message:', msg);
}

// Commit or rollback
await fork.commit(); // Merge changes
// OR
await fork.rollback(); // Discard fork

import { query, type Query, type SDKMessage } from '@anthropic-ai/claude-code';

class RealQueryControl {
  private pausedQueries = new Map<string, PausedQuery>();

  async pauseQuery(activeQuery: Query, sessionId: string): Promise<string> {
    const messages: SDKMessage[] = [];

    // 1. Collect all messages up to pause point
    for await (const message of activeQuery) {
      messages.push(message);

      // Check if pause was requested
      if (this.shouldPause(sessionId)) {
        // 2. Stop iteration (SDK's interrupt happens on break)
        break;
      }
    }

    // 3. Save pause state
    const lastMessage = messages[messages.length - 1];
    const pausePoint: PausedQuery = {
      sessionId,
      messages,
      pausedAt: Date.now(),
      resumeFromMessageId: lastMessage.uuid,
    };

    this.pausedQueries.set(sessionId, pausePoint);

    // 4. Interrupt the query
    await activeQuery.interrupt();

    return lastMessage.uuid;
  }

  async resumeQuery(sessionId: string, continuePrompt: string): Promise<Query> {
    const paused = this.pausedQueries.get(sessionId);
    if (!paused) throw new Error('No paused query found');

    // 1. Resume using SDK's resumeSessionAt
    const resumedQuery = query({
      prompt: continuePrompt,
      options: {
        resume: sessionId,
        resumeSessionAt: paused.resumeFromMessageId, // ✅ SDK resumes from exact point!
      }
    });

    // 2. Clean up paused state
    this.pausedQueries.delete(sessionId);

    return resumedQuery;
  }

  // Pause control flag
  private pauseRequests = new Set<string>();

  requestPause(sessionId: string) {
    this.pauseRequests.add(sessionId);
  }

  private shouldPause(sessionId: string): boolean {
    return this.pauseRequests.has(sessionId);
  }
}

// Usage
const controller = new RealQueryControl();

// Start query
const activeQuery = query({
  prompt: 'Long running task',
  options: { sessionId: 'my-session' }
});

// Request pause (in another thread/callback)
setTimeout(() => controller.requestPause('my-session'), 5000);

// Pause will happen during iteration
const pausePoint = await controller.pauseQuery(activeQuery, 'my-session');

console.log('Paused at message:', pausePoint);

// Resume later (even after restart!)
const resumed = await controller.resumeQuery('my-session', 'Continue task');

for await (const msg of resumed) {
  console.log('Resumed:', msg);
}

import { query, type SDKMessage } from '@anthropic-ai/claude-code';

class RealCheckpointManager {
  private checkpoints = new Map<string, Checkpoint>();

  async createCheckpoint(
    sessionId: string,
    messages: SDKMessage[],
    description: string
  ): Promise<string> {
    // Get last message UUID as checkpoint ID
    const lastMessage = messages[messages.length - 1];
    const checkpointId = lastMessage.uuid;

    // Save checkpoint with full state
    const checkpoint: Checkpoint = {
      id: checkpointId,
      sessionId,
      description,
      timestamp: Date.now(),

      // Only store message IDs (efficient!)
      messageIds: messages.map(m => m.uuid),

      // Store any additional metadata
      metadata: {
        turnCount: messages.filter(m => m.type === 'assistant').length,
        totalTokens: this.calculateTokens(messages),
      }
    };

    this.checkpoints.set(checkpointId, checkpoint);

    // Optionally persist to disk
    await this.persist(checkpoint);

    return checkpointId;
  }

  async rollbackToCheckpoint(checkpointId: string, newPrompt: string): Promise<Query> {
    const checkpoint = this.checkpoints.get(checkpointId);
    if (!checkpoint) throw new Error('Checkpoint not found');

    // Use SDK's resumeSessionAt to jump to checkpoint!
    const rolledBackQuery = query({
      prompt: newPrompt,
      options: {
        resume: checkpoint.sessionId,
        resumeSessionAt: checkpointId, // ✅ SDK rewinds to this message
      }
    });

    return rolledBackQuery;
  }

  // Auto-checkpoint on important events using hooks
  async enableAutoCheckpoint(sessionId: string) {
    return query({
      prompt: 'Task',
      options: {
        hooks: {
          PostToolUse: [{
            async hooks(input) {
              // Create checkpoint after each tool use
              if (input.tool_name === 'Edit' || input.tool_name === 'Write') {
                await this.createCheckpoint(
                  sessionId,
                  this.getMessages(sessionId),
                  `After ${input.tool_name}: ${input.tool_input.file_path}`
                );
              }
              return { continue: true };
            }
          }]
        }
      }
    });
  }

  private calculateTokens(messages: SDKMessage[]): number {
    return messages
      .filter(m => m.type === 'assistant')
      .reduce((sum, m) => {
        if ('usage' in m) {
          return sum + m.usage.input_tokens + m.usage.output_tokens;
        }
        return sum;
      }, 0);
  }

  private async persist(checkpoint: Checkpoint) {
    const fs = await import('fs/promises');
    await fs.writeFile(
      `.checkpoints/${checkpoint.id}.json`,
      JSON.stringify(checkpoint, null, 2)
    );
  }
}

// Usage
const checkpointMgr = new RealCheckpointManager();

// Start task and track messages
const messages: SDKMessage[] = [];
const taskQuery = query({ prompt: 'Complex task', options: {} });

for await (const msg of taskQuery) {
  messages.push(msg);

  // Create checkpoint after important messages
  if (msg.type === 'assistant') {
    await checkpointMgr.createCheckpoint('session-1', messages, 'After step');
  }
}

// Later: rollback to checkpoint
const checkpointId = messages[5].uuid; // 6th message
const rolledBack = await checkpointMgr.rollbackToCheckpoint(
  checkpointId,
  'Try different approach'
);

// Conversation continues from message 6!
for await (const msg of rolledBack) {
  console.log('After rollback:', msg);
}

import {
  createSdkMcpServer,
  tool,
  query
} from '@anthropic-ai/claude-code';
import { z } from 'zod';

// Create in-process MCP server
const myServer = createSdkMcpServer({
  name: 'my-tools',
  version: '1.0.0',
  tools: [
    // Tool with schema validation
    tool(
      'calculate',
      'Perform calculation',
      {
        operation: z.enum(['add', 'subtract', 'multiply', 'divide']),
        a: z.number(),
        b: z.number(),
      },
      async (args) => {
        // Direct function call - no IPC, no serialization!
        const { operation, a, b } = args;
        let result: number;

        switch (operation) {
          case 'add': result = a + b; break;
          case 'subtract': result = a - b; break;
          case 'multiply': result = a * b; break;
          case 'divide': result = a / b; break;
        }

        return {
          content: [{ type: 'text', text: `Result: ${result}` }]
        };
      }
    ),

    // Complex tool with object schema
    tool(
      'process_data',
      'Process complex data structure',
      {
        data: z.array(z.object({
          id: z.string(),
          value: z.number(),
        })),
        transformType: z.enum(['sum', 'average', 'max']),
      },
      async (args) => {
        // Direct access to JavaScript objects - no serialization!
        const values = args.data.map(d => d.value);
        let result: number;

        switch (args.transformType) {
          case 'sum':
            result = values.reduce((a, b) => a + b, 0);
            break;
          case 'average':
            result = values.reduce((a, b) => a + b, 0) / values.length;
            break;
          case 'max':
            result = Math.max(...values);
            break;
        }

        return {
          content: [{ type: 'text', text: `Result: ${result}` }]
        };
      }
    )
  ]
});

// Use in-process server
const response = query({
  prompt: 'Calculate 5 + 3',
  options: {
    mcpServers: {
      'my-tools': myServer  // ✅ SDK uses in-process, no IPC!
    }
  }
});

for await (const msg of response) {
  console.log(msg);
}

// Benchmark real speedup
async function benchmarkInProcess() {
  const iterations = 1000;
  const complexData = Array.from({ length: 1000 }, (_, i) => ({
    id: `item-${i}`,
    value: Math.random()
  }));

  // In-process
  const inProcessStart = Date.now();
  for (let i = 0; i < iterations; i++) {
    await myServer.instance.callTool('process_data', {
      data: complexData,
      transformType: 'sum'
    });
  }
  const inProcessTime = Date.now() - inProcessStart;

  // stdio MCP (for comparison)
  // Would require: JSON.stringify(complexData) -> IPC -> JSON.parse -> handler
  // Estimated: 50-100ms per call vs <1ms for in-process

  console.log(`In-process: ${inProcessTime}ms (${inProcessTime/iterations}ms per call)`);
  console.log(`Expected stdio: ~${iterations * 50}ms (50ms per call)`);
  console.log(`Real speedup: ~${(iterations * 50) / inProcessTime}x`);
}