import { z } from 'zod';

/**
 * Part 基础 Schema
 */
const PartBaseSchema = z.object({
  id: z.string(),
  createdAt: z.number(), // Unix timestamp in milliseconds
});

/**
 * 文本内容 Part
 */
export const TextPartSchema = PartBaseSchema.extend({
  type: z.literal('text'),
  text: z.string(),
});
export type TextPart = z.infer<typeof TextPartSchema>;

/**
 * 工具状态机 - Pending（等待执行）
 */
export const ToolStatePendingSchema = z.object({
  status: z.literal('pending'),
});
export type ToolStatePending = z.infer<typeof ToolStatePendingSchema>;

/**
 * 工具状态机 - Running（执行中）
 * 注意：input 使用 unknown 类型以兼容 AI SDK（可能是对象、字符串等）
 */
export const ToolStateRunningSchema = z.object({
  status: z.literal('running'),
  input: z.unknown(),
  time: z.object({ start: z.number() }),
});
export type ToolStateRunning = z.infer<typeof ToolStateRunningSchema>;

/**
 * 工具状态机 - Completed（执行完成）
 * 注意：input 使用 unknown 类型以兼容 AI SDK（可能是对象、字符串等）
 */
export const ToolStateCompletedSchema = z.object({
  status: z.literal('completed'),
  input: z.unknown(),
  output: z.unknown(),
  time: z.object({ start: z.number(), end: z.number() }),
});
export type ToolStateCompleted = z.infer<typeof ToolStateCompletedSchema>;

/**
 * 工具状态机 - Error（执行出错）
 * 注意：input 使用 unknown 类型以兼容 AI SDK（可能是对象、字符串等）
 */
export const ToolStateErrorSchema = z.object({
  status: z.literal('error'),
  input: z.unknown(),
  error: z.string(),
  time: z.object({ start: z.number(), end: z.number() }),
});
export type ToolStateError = z.infer<typeof ToolStateErrorSchema>;

/**
 * 工具状态联合类型
 */
export const ToolStateSchema = z.discriminatedUnion('status', [
  ToolStatePendingSchema,
  ToolStateRunningSchema,
  ToolStateCompletedSchema,
  ToolStateErrorSchema,
]);
export type ToolState = z.infer<typeof ToolStateSchema>;

/**
 * 工具状态字面量（用于类型检查）
 */
export type ToolStatus = ToolState['status'];

/**
 * 工具调用 Part（使用状态机模式）
 */
export const ToolPartSchema = PartBaseSchema.extend({
  type: z.literal('tool'),
  toolCallId: z.string(),
  toolName: z.string(),
  state: ToolStateSchema,
});
export type ToolPart = z.infer<typeof ToolPartSchema>;

/**
 * 推理/思考 Part
 */
export const ReasoningPartSchema = PartBaseSchema.extend({
  type: z.literal('reasoning'),
  text: z.string(),
});
export type ReasoningPart = z.infer<typeof ReasoningPartSchema>;

/**
 * 文件/图片附件 Part
 */
export const FilePartSchema = PartBaseSchema.extend({
  type: z.literal('file'),
  filename: z.string(),
  mimeType: z.string(),
  data: z.string(), // base64 encoded
  size: z.number().optional(),
});
export type FilePart = z.infer<typeof FilePartSchema>;

/**
 * 步骤开始标记 Part
 */
export const StepStartPartSchema = PartBaseSchema.extend({
  type: z.literal('step-start'),
  stepIndex: z.number(),
});
export type StepStartPart = z.infer<typeof StepStartPartSchema>;

/**
 * 步骤结束 Part（含 token 统计）
 */
export const StepFinishPartSchema = PartBaseSchema.extend({
  type: z.literal('step-finish'),
  stepIndex: z.number(),
  finishReason: z.string().optional(),
  usage: z
    .object({
      promptTokens: z.number(),
      completionTokens: z.number(),
      totalTokens: z.number(),
    })
    .optional(),
});
export type StepFinishPart = z.infer<typeof StepFinishPartSchema>;

/**
 * Git 快照引用 Part
 */
export const SnapshotPartSchema = PartBaseSchema.extend({
  type: z.literal('snapshot'),
  commitHash: z.string(),
  message: z.string().optional(),
});
export type SnapshotPart = z.infer<typeof SnapshotPartSchema>;

/**
 * Git 补丁引用 Part
 */
export const PatchPartSchema = PartBaseSchema.extend({
  type: z.literal('patch'),
  diff: z.string(),
  files: z.array(z.string()),
});
export type PatchPart = z.infer<typeof PatchPartSchema>;

/**
 * 子 Agent 调用 Part
 */
export const AgentPartSchema = PartBaseSchema.extend({
  type: z.literal('agent'),
  agentName: z.string(),
  sessionId: z.string().optional(), // 子会话 ID
  status: z.enum(['pending', 'running', 'completed', 'error']),
  result: z.string().optional(),
  error: z.string().optional(),
});
export type AgentPart = z.infer<typeof AgentPartSchema>;

/**
 * Task 工具委托 Part
 */
export const SubtaskPartSchema = PartBaseSchema.extend({
  type: z.literal('subtask'),
  taskId: z.string(),
  description: z.string(),
  status: z.enum(['pending', 'running', 'completed', 'error']),
  result: z.string().optional(),
  error: z.string().optional(),
});
export type SubtaskPart = z.infer<typeof SubtaskPartSchema>;

/**
 * 上下文压缩标记 Part
 */
export const CompactionPartSchema = PartBaseSchema.extend({
  type: z.literal('compaction'),
  summary: z.string(),
  originalTokens: z.number(),
  compressedTokens: z.number(),
});
export type CompactionPart = z.infer<typeof CompactionPartSchema>;

/**
 * 重试信息 Part
 */
export const RetryPartSchema = PartBaseSchema.extend({
  type: z.literal('retry'),
  reason: z.string(),
  attemptNumber: z.number(),
});
export type RetryPart = z.infer<typeof RetryPartSchema>;

/**
 * 所有 Part 类型的联合
 */
export const PartSchema = z.discriminatedUnion('type', [
  TextPartSchema,
  ToolPartSchema,
  ReasoningPartSchema,
  FilePartSchema,
  StepStartPartSchema,
  StepFinishPartSchema,
  SnapshotPartSchema,
  PatchPartSchema,
  AgentPartSchema,
  SubtaskPartSchema,
  CompactionPartSchema,
  RetryPartSchema,
]);
export type Part = z.infer<typeof PartSchema>;

/**
 * Part 类型字面量
 */
export type PartType = Part['type'];

/**
 * 创建 Part 的工厂函数
 */
export function createPart<T extends PartType>(
  type: T,
  data: Omit<Extract<Part, { type: T }>, 'id' | 'createdAt' | 'type'>
): Extract<Part, { type: T }> {
  return {
    id: generatePartId(),
    createdAt: Date.now(),
    type,
    ...data,
  } as Extract<Part, { type: T }>;
}

/**
 * 生成 Part ID
 */
function generatePartId(): string {
  const timestamp = Date.now();
  const random = Math.random().toString(36).substring(2, 8);
  return `part_${timestamp}_${random}`;
}
