import * as base from './base.js';
import { generatePartId } from '../id.js';
import type { Part, PartType, ToolPart, ToolStatus, ToolState } from '../parts.js';
import { PartSchema } from '../parts.js';

// Re-export types
export type { Part, PartType, ToolPart, ToolStatus, ToolState } from '../parts.js';

/**
 * 创建 Part
 */
export async function create<T extends Part>(
  messageId: string,
  type: T['type'],
  data: Omit<T, 'id' | 'createdAt' | 'type'>
): Promise<T> {
  const part = {
    id: generatePartId(),
    createdAt: Date.now(),
    type,
    ...data,
  } as T;

  await base.write(['part', messageId, part.id], part);
  return part;
}

/**
 * 读取 Part
 */
export async function get(messageId: string, partId: string): Promise<Part | null> {
  try {
    return await base.read(['part', messageId, partId], PartSchema);
  } catch (e) {
    if (e instanceof base.StorageNotFoundError) {
      return null;
    }
    throw e;
  }
}

/**
 * 保存 Part
 */
export async function save(messageId: string, part: Part): Promise<void> {
  await base.write(['part', messageId, part.id], part);
}

/**
 * 更新 Part
 */
export async function update<T extends Part>(
  messageId: string,
  partId: string,
  updates: Partial<Omit<T, 'id' | 'createdAt' | 'type'>>
): Promise<Part> {
  return base.update(['part', messageId, partId], (part: Part) => {
    Object.assign(part, updates);
  });
}

/**
 * 更新 ToolPart 状态（新的状态机模式）
 */
export async function updateToolState(
  messageId: string,
  partId: string,
  state: ToolState
): Promise<ToolPart> {
  return base.update(['part', messageId, partId], (part: ToolPart) => {
    part.state = state;
  }) as Promise<ToolPart>;
}

/**
 * 更新 ToolPart 为 running 状态
 */
export async function setToolRunning(
  messageId: string,
  partId: string,
  input: Record<string, unknown>
): Promise<ToolPart> {
  return updateToolState(messageId, partId, {
    status: 'running',
    input,
    time: { start: Date.now() },
  });
}

/**
 * 更新 ToolPart 为 completed 状态
 */
export async function setToolCompleted(
  messageId: string,
  partId: string,
  output: unknown,
  startTime: number
): Promise<ToolPart> {
  // 先获取当前 part 以获取 input
  const part = await get(messageId, partId) as ToolPart | null;
  // 从 running 状态获取 input（需要类型断言因为 discriminated union）
  const state = part?.state;
  const input = state && state.status !== 'pending'
    ? (state as { input: Record<string, unknown> }).input
    : {};

  return updateToolState(messageId, partId, {
    status: 'completed',
    input: input ?? {},
    output,
    time: { start: startTime, end: Date.now() },
  });
}

/**
 * 更新 ToolPart 为 error 状态
 */
export async function setToolError(
  messageId: string,
  partId: string,
  error: string,
  startTime: number
): Promise<ToolPart> {
  // 先获取当前 part 以获取 input
  const part = await get(messageId, partId) as ToolPart | null;
  // 从 running 状态获取 input（需要类型断言因为 discriminated union）
  const state = part?.state;
  const input = state && state.status !== 'pending'
    ? (state as { input: Record<string, unknown> }).input
    : {};

  return updateToolState(messageId, partId, {
    status: 'error',
    input: input ?? {},
    error,
    time: { start: startTime, end: Date.now() },
  });
}

/**
 * 删除 Part
 */
export async function remove(messageId: string, partId: string): Promise<void> {
  await base.remove(['part', messageId, partId]);
}

/**
 * 列出消息的所有 Parts
 */
export async function listByMessage(messageId: string): Promise<Part[]> {
  const keys = await base.list(['part', messageId]);
  const parts: Part[] = [];

  for (const key of keys) {
    try {
      const part = await base.read<Part>(key, PartSchema);
      parts.push(part);
    } catch {
      // 跳过无法解析的文件
    }
  }

  // 按创建时间升序排列（Part ID 是升序的）
  return parts.sort((a, b) => a.createdAt - b.createdAt);
}

/**
 * 根据 Part IDs 获取 Parts（保持顺序）
 */
export async function getByIds(messageId: string, partIds: string[]): Promise<Part[]> {
  const parts: Part[] = [];

  for (const partId of partIds) {
    const part = await get(messageId, partId);
    if (part) {
      parts.push(part);
    }
  }

  return parts;
}

/**
 * 删除消息的所有 Parts
 */
export async function removeByMessage(messageId: string): Promise<void> {
  await base.removeDir(['part', messageId]);
}

/**
 * 创建文本 Part
 */
export async function createText(messageId: string, text: string): Promise<Part> {
  return create(messageId, 'text', { text });
}

/**
 * 创建工具调用 Part（pending 状态）
 */
export async function createTool(
  messageId: string,
  toolCallId: string,
  toolName: string
): Promise<ToolPart> {
  return create<ToolPart>(messageId, 'tool', {
    toolCallId,
    toolName,
    state: { status: 'pending' },
  });
}

/**
 * 创建工具调用 Part（直接 running 状态）
 */
export async function createToolRunning(
  messageId: string,
  toolCallId: string,
  toolName: string,
  input: Record<string, unknown>
): Promise<ToolPart> {
  return create<ToolPart>(messageId, 'tool', {
    toolCallId,
    toolName,
    state: {
      status: 'running',
      input,
      time: { start: Date.now() },
    },
  });
}

/**
 * 创建推理 Part
 */
export async function createReasoning(messageId: string, text: string): Promise<Part> {
  return create(messageId, 'reasoning', { text });
}
