import z from "zod";

/**
 * Single source of truth for every wire type in `@suaveplan/mcp`.
 *
 * Every public TypeScript type for a wire shape is derived via `z.infer<>`
 * from one of the schemas declared below. Hand-rolled interfaces are NOT
 * permitted for wire-shape types — the schema IS the contract.
 */

// -----------------------------------------------------------------------------
// JSON-RPC 2.0 envelopes
// -----------------------------------------------------------------------------

export const JsonRpcIdSchema = z.union([z.string(), z.number()]);
export type JsonRpcId = z.infer<typeof JsonRpcIdSchema>;

export const JsonRpcRequestSchema = z.object({
    jsonrpc: z.literal("2.0"),
    id: JsonRpcIdSchema,
    method: z.string().min(1),
    params: z.unknown().optional(),
});
export type JsonRpcRequest = z.infer<typeof JsonRpcRequestSchema>;

export const JsonRpcErrorSchema = z.object({
    code: z.number().int(),
    message: z.string(),
    data: z.unknown().optional(),
});
export type JsonRpcErrorPayload = z.infer<typeof JsonRpcErrorSchema>;

export const JsonRpcResponseSchema = z.object({
    jsonrpc: z.literal("2.0"),
    id: JsonRpcIdSchema,
    result: z.unknown().optional(),
    error: JsonRpcErrorSchema.optional(),
});
export type JsonRpcResponse = z.infer<typeof JsonRpcResponseSchema>;

export const JsonRpcNotificationSchema = z.object({
    jsonrpc: z.literal("2.0"),
    method: z.string().min(1),
    params: z.unknown().optional(),
});
export type JsonRpcNotification = z.infer<typeof JsonRpcNotificationSchema>;

export const JsonRpcMessageSchema = z.union([
    JsonRpcRequestSchema,
    JsonRpcResponseSchema,
    JsonRpcNotificationSchema,
]);
export type JsonRpcMessage = z.infer<typeof JsonRpcMessageSchema>;

// -----------------------------------------------------------------------------
// Implementation / capabilities / initialize
// -----------------------------------------------------------------------------

export const ImplementationSchema = z.object({
    name: z.string().min(1),
    version: z.string().min(1),
});
export type Implementation = z.infer<typeof ImplementationSchema>;

export const ServerCapabilitiesSchema = z.object({
    tools: z.object({ listChanged: z.boolean().optional() }).optional(),
    resources: z
        .object({
            subscribe: z.boolean().optional(),
            listChanged: z.boolean().optional(),
        })
        .optional(),
    prompts: z.object({ listChanged: z.boolean().optional() }).optional(),
    logging: z.object({}).optional(),
    completions: z.object({}).optional(),
});
export type ServerCapabilities = z.infer<typeof ServerCapabilitiesSchema>;

export const ClientCapabilitiesSchema = z.object({
    sampling: z.object({}).optional(),
    roots: z.object({ listChanged: z.boolean().optional() }).optional(),
});
export type ClientCapabilities = z.infer<typeof ClientCapabilitiesSchema>;

export const InitializeParamsSchema = z.object({
    protocolVersion: z.string().min(1),
    capabilities: ClientCapabilitiesSchema,
    clientInfo: ImplementationSchema,
});
export type InitializeParams = z.infer<typeof InitializeParamsSchema>;

export const InitializeResultSchema = z.object({
    protocolVersion: z.string().min(1),
    capabilities: ServerCapabilitiesSchema,
    serverInfo: ImplementationSchema,
    instructions: z.string().optional(),
});
export type InitializeResult = z.infer<typeof InitializeResultSchema>;

// -----------------------------------------------------------------------------
// Content items (shared by tool results, prompt messages, resource reads)
// -----------------------------------------------------------------------------

export const TextContentSchema = z.object({
    type: z.literal("text"),
    text: z.string(),
});

export const ImageContentSchema = z.object({
    type: z.literal("image"),
    data: z.string(),
    mimeType: z.string(),
});

export const ResourceContentInlineSchema = z.object({
    type: z.literal("resource"),
    resource: z.object({
        uri: z.string(),
        mimeType: z.string().optional(),
        text: z.string().optional(),
        blob: z.string().optional(),
    }),
});

export const ContentItemSchema = z.discriminatedUnion("type", [
    TextContentSchema,
    ImageContentSchema,
    ResourceContentInlineSchema,
]);
export type ContentItem = z.infer<typeof ContentItemSchema>;

// -----------------------------------------------------------------------------
// Tools
// -----------------------------------------------------------------------------

export const ToolSchema = z.object({
    name: z.string().min(1),
    description: z.string().optional(),
    inputSchema: z.record(z.string(), z.unknown()),
});
export type Tool = z.infer<typeof ToolSchema>;

export const ToolsListResultSchema = z.object({
    tools: z.array(ToolSchema),
    nextCursor: z.string().optional(),
});
export type ToolsListResult = z.infer<typeof ToolsListResultSchema>;

export const ToolCallParamsSchema = z.object({
    name: z.string().min(1),
    arguments: z.record(z.string(), z.unknown()).optional(),
    _meta: z
        .object({ progressToken: z.union([z.string(), z.number()]).optional() })
        .optional(),
});
export type ToolCallParams = z.infer<typeof ToolCallParamsSchema>;

export const ToolCallResultSchema = z.object({
    content: z.array(ContentItemSchema),
    isError: z.boolean().optional(),
});
export type ToolCallResult = z.infer<typeof ToolCallResultSchema>;

// -----------------------------------------------------------------------------
// Resources
// -----------------------------------------------------------------------------

export const ResourceSchema = z.object({
    uri: z.string().min(1),
    name: z.string().min(1),
    description: z.string().optional(),
    mimeType: z.string().optional(),
});
export type Resource = z.infer<typeof ResourceSchema>;

export const ResourceTemplateSchema = z.object({
    uriTemplate: z.string().min(1),
    name: z.string().min(1),
    description: z.string().optional(),
    mimeType: z.string().optional(),
});
export type ResourceTemplate = z.infer<typeof ResourceTemplateSchema>;

export const ResourcesListResultSchema = z.object({
    resources: z.array(ResourceSchema),
    nextCursor: z.string().optional(),
});
export type ResourcesListResult = z.infer<typeof ResourcesListResultSchema>;

export const ResourceTemplatesListResultSchema = z.object({
    resourceTemplates: z.array(ResourceTemplateSchema),
});
export type ResourceTemplatesListResult = z.infer<
    typeof ResourceTemplatesListResultSchema
>;

export const ResourceContentSchema = z.object({
    uri: z.string().min(1),
    mimeType: z.string().optional(),
    text: z.string().optional(),
    blob: z.string().optional(),
});
export type ResourceContent = z.infer<typeof ResourceContentSchema>;

export const ResourceReadParamsSchema = z.object({
    uri: z.string().min(1),
});
export type ResourceReadParams = z.infer<typeof ResourceReadParamsSchema>;

export const ResourceReadResultSchema = z.object({
    contents: z.array(ResourceContentSchema),
});
export type ResourceReadResult = z.infer<typeof ResourceReadResultSchema>;

export const ResourceSubscribeParamsSchema = z.object({
    uri: z.string().min(1),
});
export type ResourceSubscribeParams = z.infer<
    typeof ResourceSubscribeParamsSchema
>;

// -----------------------------------------------------------------------------
// Prompts
// -----------------------------------------------------------------------------

export const PromptArgumentSchema = z.object({
    name: z.string().min(1),
    description: z.string().optional(),
    required: z.boolean().optional(),
});
export type PromptArgument = z.infer<typeof PromptArgumentSchema>;

export const PromptSchema = z.object({
    name: z.string().min(1),
    description: z.string().optional(),
    arguments: z.array(PromptArgumentSchema).optional(),
});
export type Prompt = z.infer<typeof PromptSchema>;

export const PromptsListResultSchema = z.object({
    prompts: z.array(PromptSchema),
    nextCursor: z.string().optional(),
});
export type PromptsListResult = z.infer<typeof PromptsListResultSchema>;

export const PromptMessageSchema = z.object({
    role: z.enum(["user", "assistant"]),
    content: ContentItemSchema,
});
export type PromptMessage = z.infer<typeof PromptMessageSchema>;

export const PromptGetParamsSchema = z.object({
    name: z.string().min(1),
    arguments: z.record(z.string(), z.string()).optional(),
});
export type PromptGetParams = z.infer<typeof PromptGetParamsSchema>;

export const PromptGetResultSchema = z.object({
    description: z.string().optional(),
    messages: z.array(PromptMessageSchema),
});
export type PromptGetResult = z.infer<typeof PromptGetResultSchema>;

// -----------------------------------------------------------------------------
// Completion
// -----------------------------------------------------------------------------

export const CompletionRefSchema = z.union([
    z.object({ type: z.literal("ref/prompt"), name: z.string() }),
    z.object({ type: z.literal("ref/resource"), uri: z.string() }),
]);
export type CompletionRef = z.infer<typeof CompletionRefSchema>;

export const CompletionRequestSchema = z.object({
    ref: CompletionRefSchema,
    argument: z.object({ name: z.string(), value: z.string() }),
});
export type CompletionRequest = z.infer<typeof CompletionRequestSchema>;

export const CompletionResultSchema = z.object({
    completion: z.object({
        values: z.array(z.string()),
        total: z.number().int().nonnegative().optional(),
        hasMore: z.boolean().optional(),
    }),
});
export type CompletionResult = z.infer<typeof CompletionResultSchema>;

// -----------------------------------------------------------------------------
// Sampling
// -----------------------------------------------------------------------------

export const SamplingMessageSchema = z.object({
    role: z.enum(["user", "assistant"]),
    content: z.discriminatedUnion("type", [
        TextContentSchema,
        ImageContentSchema,
    ]),
});
export type SamplingMessage = z.infer<typeof SamplingMessageSchema>;

export const CreateMessageParamsSchema = z.object({
    messages: z.array(SamplingMessageSchema),
    maxTokens: z.number().int().positive(),
    temperature: z.number().optional(),
    systemPrompt: z.string().optional(),
    modelPreferences: z
        .object({
            hints: z.array(z.object({ name: z.string() })).optional(),
            costPriority: z.number().min(0).max(1).optional(),
            speedPriority: z.number().min(0).max(1).optional(),
            intelligencePriority: z.number().min(0).max(1).optional(),
        })
        .optional(),
});
export type CreateMessageParams = z.infer<typeof CreateMessageParamsSchema>;

export const CreateMessageResultSchema = z.object({
    role: z.enum(["user", "assistant"]),
    content: z.discriminatedUnion("type", [
        TextContentSchema,
        ImageContentSchema,
    ]),
    model: z.string(),
    stopReason: z.enum(["endTurn", "stopSequence", "maxTokens"]).optional(),
});
export type CreateMessageResult = z.infer<typeof CreateMessageResultSchema>;

// -----------------------------------------------------------------------------
// Roots
// -----------------------------------------------------------------------------

export const RootSchema = z.object({
    uri: z.string().min(1),
    name: z.string().optional(),
});
export type Root = z.infer<typeof RootSchema>;

export const RootsListResultSchema = z.object({
    roots: z.array(RootSchema),
});
export type RootsListResult = z.infer<typeof RootsListResultSchema>;

// -----------------------------------------------------------------------------
// Logging
// -----------------------------------------------------------------------------

export const LoggingLevelSchema = z.enum([
    "debug",
    "info",
    "notice",
    "warning",
    "error",
    "critical",
    "alert",
    "emergency",
]);
export type LoggingLevel = z.infer<typeof LoggingLevelSchema>;

export const LoggingSetLevelParamsSchema = z.object({
    level: LoggingLevelSchema,
});
export type LoggingSetLevelParams = z.infer<typeof LoggingSetLevelParamsSchema>;

export const LoggingMessageParamsSchema = z.object({
    level: LoggingLevelSchema,
    logger: z.string().optional(),
    data: z.unknown(),
});
export type LoggingMessageParams = z.infer<typeof LoggingMessageParamsSchema>;

// -----------------------------------------------------------------------------
// Notifications
// -----------------------------------------------------------------------------

export const ProgressNotificationParamsSchema = z.object({
    progressToken: z.union([z.string(), z.number()]),
    progress: z.number(),
    total: z.number().optional(),
});
export type ProgressNotificationParams = z.infer<
    typeof ProgressNotificationParamsSchema
>;

export const CancelledNotificationParamsSchema = z.object({
    requestId: z.union([z.string(), z.number()]),
    reason: z.string().optional(),
});
export type CancelledNotificationParams = z.infer<
    typeof CancelledNotificationParamsSchema
>;

// -----------------------------------------------------------------------------
// Ping (request returns empty object)
// -----------------------------------------------------------------------------

export const EmptyResultSchema = z.object({});
export type EmptyResult = z.infer<typeof EmptyResultSchema>;