import { z } from 'zod';
import * as base from './base.js';
import { generateSessionId } from '../id.js';

/**
 * Session Info Schema
 */
export const SessionInfoSchema = z.object({
  id: z.string(),
  projectId: z.string(),
  parentId: z.string().optional(),
  agentName: z.string().optional(),
  createdAt: z.number(),
  updatedAt: z.number(),
  workdir: z.string(),
  title: z.string().optional(),
  discoveredTools: z.array(z.string()).default([]),
  // 统计信息
  stats: z
    .object({
      messageCount: z.number(),
      inputTokens: z.number(),
      outputTokens: z.number(),
    })
    .optional(),
});
export type SessionInfo = z.infer<typeof SessionInfoSchema>;

/**
 * 创建会话
 */
export async function create(
  projectId: string,
  workdir: string,
  options?: {
    parentId?: string;
    agentName?: string;
    title?: string;
  }
): Promise<SessionInfo> {
  const session: SessionInfo = {
    id: generateSessionId(),
    projectId,
    workdir,
    createdAt: Date.now(),
    updatedAt: Date.now(),
    parentId: options?.parentId,
    agentName: options?.agentName,
    title: options?.title,
    discoveredTools: [],
  };

  await base.write(['session', projectId, session.id], session);
  return session;
}

/**
 * 读取会话
 */
export async function get(projectId: string, sessionId: string): Promise<SessionInfo | null> {
  try {
    return await base.read(['session', projectId, sessionId], SessionInfoSchema);
  } catch (e) {
    if (e instanceof base.StorageNotFoundError) {
      return null;
    }
    throw e;
  }
}

/**
 * 更新会话
 */
export async function update(
  projectId: string,
  sessionId: string,
  updates: Partial<Omit<SessionInfo, 'id' | 'projectId' | 'createdAt'>>
): Promise<SessionInfo> {
  return base.update(['session', projectId, sessionId], (session: SessionInfo) => {
    Object.assign(session, updates, { updatedAt: Date.now() });
  });
}

/**
 * 保存会话
 */
export async function save(session: SessionInfo): Promise<void> {
  session.updatedAt = Date.now();
  await base.write(['session', session.projectId, session.id], session);
}

/**
 * 删除会话
 */
export async function remove(projectId: string, sessionId: string): Promise<void> {
  await base.remove(['session', projectId, sessionId]);
}

/**
 * 列出项目的所有会话
 */
export async function listByProject(projectId: string): Promise<SessionInfo[]> {
  const keys = await base.list(['session', projectId]);
  const sessions: SessionInfo[] = [];

  for (const key of keys) {
    try {
      const session = await base.read<SessionInfo>(key, SessionInfoSchema);
      sessions.push(session);
    } catch {
      // 跳过无法解析的文件
    }
  }

  // 按更新时间降序排列
  return sessions.sort((a, b) => b.updatedAt - a.updatedAt);
}

/**
 * 列出所有会话（跨项目）
 */
export async function listAll(): Promise<SessionInfo[]> {
  const projectIds = await base.listDirect(['session']);
  const allSessions: SessionInfo[] = [];

  for (const projectId of projectIds) {
    const sessions = await listByProject(projectId);
    allSessions.push(...sessions);
  }

  // 按更新时间降序排列
  return allSessions.sort((a, b) => b.updatedAt - a.updatedAt);
}

/**
 * 更新会话标题
 */
export async function updateTitle(projectId: string, sessionId: string, title: string): Promise<void> {
  await update(projectId, sessionId, { title });
}

/**
 * 更新会话统计
 */
export async function updateStats(
  projectId: string,
  sessionId: string,
  stats: SessionInfo['stats']
): Promise<void> {
  await update(projectId, sessionId, { stats });
}

/**
 * 添加已发现的工具
 */
export async function addDiscoveredTool(
  projectId: string,
  sessionId: string,
  toolName: string
): Promise<void> {
  await base.update(['session', projectId, sessionId], (session: SessionInfo) => {
    if (!session.discoveredTools.includes(toolName)) {
      session.discoveredTools.push(toolName);
    }
    session.updatedAt = Date.now();
  });
}
