import z from "zod";
import { Bus } from "../bus";

export namespace Message {
  export const ToolCall = z
    .object({
      state: z.literal("call"),
      step: z.number().optional(),
      toolCallId: z.string(),
      toolName: z.string(),
      args: z.custom<Required<unknown>>(),
    })
    .openapi({
      ref: "Message.ToolInvocation.ToolCall",
    });
  export type ToolCall = z.infer<typeof ToolCall>;

  export const ToolPartialCall = z
    .object({
      state: z.literal("partial-call"),
      step: z.number().optional(),
      toolCallId: z.string(),
      toolName: z.string(),
      args: z.custom<Required<unknown>>(),
    })
    .openapi({
      ref: "Message.ToolInvocation.ToolPartialCall",
    });
  export type ToolPartialCall = z.infer<typeof ToolPartialCall>;

  export const ToolResult = z
    .object({
      state: z.literal("result"),
      step: z.number().optional(),
      toolCallId: z.string(),
      toolName: z.string(),
      args: z.custom<Required<unknown>>(),
      result: z.string(),
    })
    .openapi({
      ref: "Message.ToolInvocation.ToolResult",
    });
  export type ToolResult = z.infer<typeof ToolResult>;

  export const ToolInvocation = z
    .discriminatedUnion("state", [ToolCall, ToolPartialCall, ToolResult])
    .openapi({
      ref: "Message.ToolInvocation",
    });
  export type ToolInvocation = z.infer<typeof ToolInvocation>;

  export const TextPart = z
    .object({
      type: z.literal("text"),
      text: z.string(),
    })
    .openapi({
      ref: "Message.Part.Text",
    });
  export type TextPart = z.infer<typeof TextPart>;

  export const ReasoningPart = z
    .object({
      type: z.literal("reasoning"),
      text: z.string(),
      providerMetadata: z.record(z.any()).optional(),
    })
    .openapi({
      ref: "Message.Part.Reasoning",
    });
  export type ReasoningPart = z.infer<typeof ReasoningPart>;

  export const ToolInvocationPart = z
    .object({
      type: z.literal("tool-invocation"),
      toolInvocation: ToolInvocation,
    })
    .openapi({
      ref: "Message.Part.ToolInvocation",
    });
  export type ToolInvocationPart = z.infer<typeof ToolInvocationPart>;

  export const SourceUrlPart = z
    .object({
      type: z.literal("source-url"),
      sourceId: z.string(),
      url: z.string(),
      title: z.string().optional(),
      providerMetadata: z.record(z.any()).optional(),
    })
    .openapi({
      ref: "Message.Part.SourceUrl",
    });
  export type SourceUrlPart = z.infer<typeof SourceUrlPart>;

  export const FilePart = z
    .object({
      type: z.literal("file"),
      mediaType: z.string(),
      filename: z.string().optional(),
      url: z.string(),
    })
    .openapi({
      ref: "Message.Part.File",
    });
  export type FilePart = z.infer<typeof FilePart>;

  export const StepStartPart = z
    .object({
      type: z.literal("step-start"),
    })
    .openapi({
      ref: "Message.Part.StepStart",
    });
  export type StepStartPart = z.infer<typeof StepStartPart>;

  export const Part = z
    .discriminatedUnion("type", [
      TextPart,
      ReasoningPart,
      ToolInvocationPart,
      SourceUrlPart,
      FilePart,
      StepStartPart,
    ])
    .openapi({
      ref: "Message.Part",
    });
  export type Part = z.infer<typeof Part>;

  export const Info = z
    .object({
      id: z.string(),
      role: z.enum(["system", "user", "assistant"]),
      parts: z.array(Part),
      metadata: z.object({
        time: z.object({
          created: z.number(),
          completed: z.number().optional(),
        }),
        sessionID: z.string(),
        tool: z.record(z.string(), z.any()),
        assistant: z
          .object({
            modelID: z.string(),
            providerID: z.string(),
            cost: z.number(),
            summary: z.boolean().optional(),
            tokens: z.object({
              input: z.number(),
              output: z.number(),
              reasoning: z.number(),
            }),
          })
          .optional(),
      }),
    })
    .openapi({
      ref: "Message.Info",
    });
  export type Info = z.infer<typeof Info>;

  export const Event = {
    Updated: Bus.event(
      "message.updated",
      z.object({
        info: Info,
      }),
    ),
  };
}