import { z } from 'zod';
import * as base from './base.js';
import { generateMessageId } from '../id.js';
import type { MessageInfo } from '../message.js';
import { MessageInfoSchema } from '../message.js';

// Re-export types
export type { MessageInfo } from '../message.js';

/**
 * 创建消息
 */
export async function create(
  sessionId: string,
  role: MessageInfo['role'],
  options?: {
    parentId?: string;
    partIds?: string[];
    metadata?: MessageInfo['metadata'];
  }
): Promise<MessageInfo> {
  const message: MessageInfo = {
    id: generateMessageId(),
    sessionId,
    role,
    parentId: options?.parentId,
    createdAt: Date.now(),
    partIds: options?.partIds || [],
    metadata: options?.metadata,
  };

  await base.write(['message', sessionId, message.id], message);
  return message;
}

/**
 * 读取消息
 */
export async function get(sessionId: string, messageId: string): Promise<MessageInfo | null> {
  try {
    return await base.read(['message', sessionId, messageId], MessageInfoSchema);
  } catch (e) {
    if (e instanceof base.StorageNotFoundError) {
      return null;
    }
    throw e;
  }
}

/**
 * 保存消息
 */
export async function save(message: MessageInfo): Promise<void> {
  await base.write(['message', message.sessionId, message.id], message);
}

/**
 * 更新消息
 */
export async function update(
  sessionId: string,
  messageId: string,
  updates: Partial<Omit<MessageInfo, 'id' | 'sessionId' | 'createdAt'>>
): Promise<MessageInfo> {
  return base.update(['message', sessionId, messageId], (message: MessageInfo) => {
    Object.assign(message, updates);
  });
}

/**
 * 添加 Part 到消息
 */
export async function addPart(sessionId: string, messageId: string, partId: string): Promise<void> {
  await base.update(['message', sessionId, messageId], (message: MessageInfo) => {
    message.partIds.push(partId);
  });
}

/**
 * 删除消息
 */
export async function remove(sessionId: string, messageId: string): Promise<void> {
  await base.remove(['message', sessionId, messageId]);
}

/**
 * 列出会话的所有消息
 */
export async function listBySession(sessionId: string): Promise<MessageInfo[]> {
  const keys = await base.list(['message', sessionId]);
  const messages: MessageInfo[] = [];

  for (const key of keys) {
    try {
      const message = await base.read<MessageInfo>(key, MessageInfoSchema);
      messages.push(message);
    } catch {
      // 跳过无法解析的文件
    }
  }

  // 按创建时间升序排列（最旧的在前）
  // 当时间相同时，按 ID 降序排列（ID 使用降序时间戳，所以 ID 大的更旧）
  return messages.sort((a, b) => {
    const timeDiff = a.createdAt - b.createdAt;
    if (timeDiff !== 0) return timeDiff;
    // ID 降序时间戳：数字大 = 时间旧，所以用 b - a 让旧的在前
    return b.id.localeCompare(a.id);
  });
}

/**
 * 删除会话的所有消息
 */
export async function removeBySession(sessionId: string): Promise<void> {
  await base.removeDir(['message', sessionId]);
}

/**
 * 获取会话的消息数量
 */
export async function countBySession(sessionId: string): Promise<number> {
  const keys = await base.list(['message', sessionId]);
  return keys.length;
}
