import z from "zod"
import { SessionID } from "./schema"
import { ModelID, ProviderID } from "../provider/schema"
import { NamedError } from "@woozlit/shared/util/error"

export const OutputLengthError = NamedError.create("MessageOutputLengthError", z.object({}))
export const AuthError = NamedError.create(
  "ProviderAuthError",
  z.object({
    providerID: z.string(),
    message: z.string(),
  }),
)

export const ToolCall = z
  .object({
    state: z.literal("call"),
    step: z.number().optional(),
    toolCallId: z.string(),
    toolName: z.string(),
    args: z.custom<Required<unknown>>(),
  })
  .meta({
    ref: "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>>(),
  })
  .meta({
    ref: "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(),
  })
  .meta({
    ref: "ToolResult",
  })
export type ToolResult = z.infer<typeof ToolResult>

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

export const TextPart = z
  .object({
    type: z.literal("text"),
    text: z.string(),
  })
  .meta({
    ref: "TextPart",
  })
export type TextPart = z.infer<typeof TextPart>

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

export const ToolInvocationPart = z
  .object({
    type: z.literal("tool-invocation"),
    toolInvocation: ToolInvocation,
  })
  .meta({
    ref: "ToolInvocationPart",
  })
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.string(), z.any()).optional(),
  })
  .meta({
    ref: "SourceUrlPart",
  })
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(),
  })
  .meta({
    ref: "FilePart",
  })
export type FilePart = z.infer<typeof FilePart>

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

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

export const Info = z
  .object({
    id: z.string(),
    role: z.enum(["user", "assistant"]),
    parts: z.array(MessagePart),
    metadata: z
      .object({
        time: z.object({
          created: z.number(),
          completed: z.number().optional(),
        }),
        error: z
          .discriminatedUnion("name", [AuthError.Schema, NamedError.Unknown.Schema, OutputLengthError.Schema])
          .optional(),
        sessionID: SessionID.zod,
        tool: z.record(
          z.string(),
          z
            .object({
              title: z.string(),
              snapshot: z.string().optional(),
              time: z.object({
                start: z.number(),
                end: z.number(),
              }),
            })
            .catchall(z.any()),
        ),
        assistant: z
          .object({
            system: z.string().array(),
            modelID: ModelID.zod,
            providerID: ProviderID.zod,
            path: z.object({
              cwd: z.string(),
              root: z.string(),
            }),
            cost: z.number(),
            summary: z.boolean().optional(),
            tokens: z.object({
              input: z.number(),
              output: z.number(),
              reasoning: z.number(),
              cache: z.object({
                read: z.number(),
                write: z.number(),
              }),
            }),
          })
          .optional(),
        snapshot: z.string().optional(),
      })
      .meta({ ref: "MessageMetadata" }),
  })
  .meta({
    ref: "Message",
  })
export type Info = z.infer<typeof Info>

export * as Message from "./message"
