import { z } from "zod";
import { JsonValueSchema } from "../jsonSchema";

const CacheControl = z.object({
  type: z.string(),
  ttl: z.string().optional(),
  scope: z.string().optional(),
});

const TextContentBlock = z.object({
  type: z.literal("text"),
  text: z.string(),
  cache_control: CacheControl.optional(),
});

const ThinkingContentBlock = z.object({
  type: z.literal("thinking"),
  thinking: z.string(),
  signature: z.string().optional(),
});

const ThinkContentBlock = z.object({
  type: z.literal("think"),
  thinking: z.string(),
  signature: z.string().optional(),
});

const ImageSourceBlock = z.object({
  type: z.literal("base64"),
  media_type: z.string(),
  data: z.string(),
});

const ImageContentBlock = z.object({
  type: z.literal("image"),
  source: ImageSourceBlock,
});

const ToolUseContentBlock = z.object({
  type: z.literal("tool_use"),
  id: z.string(),
  name: z.string(),
  input: z.record(z.string(), JsonValueSchema),
});

const ToolResultContentItem = z.discriminatedUnion("type", [TextContentBlock, ImageContentBlock]);

const ToolResultContentBlock = z.object({
  type: z.literal("tool_result"),
  tool_use_id: z.string().optional(),
  content: z.union([z.string(), z.array(ToolResultContentItem)]),
  is_error: z.boolean().optional(),
});

const ContentBlock = z.discriminatedUnion("type", [
  TextContentBlock,
  ThinkingContentBlock,
  ImageContentBlock,
  ToolUseContentBlock,
  ToolResultContentBlock,
]);

const MessageContent = z.union([z.string(), z.array(ContentBlock)]);

const Message = z.object({
  role: z.enum(["user", "assistant"]),
  content: MessageContent,
});

const SystemBlock = z.object({
  type: z.literal("text"),
  text: z.string(),
  cache_control: CacheControl.optional(),
});

const InputSchema = z.object({
  type: z.string(),
  properties: z.record(z.string(), z.record(z.string(), JsonValueSchema)).optional(),
  required: z.array(z.string()).optional(),
  additionalProperties: z.boolean().optional(),
  $schema: z.string().optional(),
});

const ToolDefinition = z.object({
  name: z.string(),
  description: z.string().optional(),
  input_schema: InputSchema.optional(),
  cache_control: CacheControl.optional(),
});

const ThinkingConfig = z.discriminatedUnion("type", [
  z.object({ type: z.literal("enabled"), budget_tokens: z.number() }),
  z.object({ type: z.literal("disabled") }),
  z.object({ type: z.literal("adaptive") }),
]);

export const AnthropicRequestSchema = z.object({
  model: z.string(),
  messages: z.array(Message),
  system: z.array(SystemBlock).optional(),
  tools: z.array(ToolDefinition).optional(),
  max_tokens: z.number().optional(),
  temperature: z.number().optional(),
  stream: z.boolean().optional(),
  thinking: ThinkingConfig.optional(),
  metadata: z
    .object({
      user_id: z.string().optional(),
    })
    .optional(),
});

const ResponseContentBlock = z.discriminatedUnion("type", [
  TextContentBlock,
  ThinkingContentBlock,
  ThinkContentBlock,
  ToolUseContentBlock,
]);

const ResponseUsageSchema = z
  .object({
    input_tokens: z.number().nullable().optional(),
    output_tokens: z.number().nullable().optional(),
    cache_creation_input_tokens: z.number().nullable().optional(),
    cache_read_input_tokens: z.number().nullable().optional(),
  })
  .passthrough(); // Allow extra fields from different providers

export const AnthropicResponseSchema = z
  .object({
    id: z.string(),
    type: z.literal("message"),
    model: z.string(),
    role: z.literal("assistant"),
    content: z.array(ResponseContentBlock),
    stop_reason: z.string().nullable(),
    stop_sequence: z.string().nullable().optional(),
    usage: ResponseUsageSchema,
  })
  .passthrough(); // Allow extra unknown fields

const SseMessageStartEvent = z.object({
  type: z.literal("message_start"),
  message: z.object({
    id: z.string(),
    type: z.literal("message"),
    model: z.string(),
    role: z.literal("assistant"),
    content: z.array(ResponseContentBlock),
    stop_reason: z.null(),
    stop_sequence: z.null(),
    usage: z
      .object({
        input_tokens: z.number(),
        cache_creation_input_tokens: z.number().optional(),
        cache_read_input_tokens: z.number().optional(),
      })
      .passthrough(),
  }),
});

const SseContentBlockStartEvent = z.object({
  type: z.literal("content_block_start"),
  index: z.number(),
  content_block: ResponseContentBlock,
});

const SseDeltaBlock = z.discriminatedUnion("type", [
  z.object({ type: z.literal("text_delta"), text: z.string() }),
  z.object({ type: z.literal("input_json_delta"), partial_json: z.string() }),
  z.object({ type: z.literal("thinking_delta"), thinking: z.string() }),
  z.object({ type: z.literal("signature_delta"), signature: z.string() }),
]);

const SseContentBlockDeltaEvent = z.object({
  type: z.literal("content_block_delta"),
  index: z.number(),
  delta: SseDeltaBlock,
});

const SseContentBlockStopEvent = z.object({
  type: z.literal("content_block_stop"),
  index: z.number(),
});

const SseMessageDeltaEvent = z.object({
  type: z.literal("message_delta"),
  delta: z.object({
    stop_reason: z.string().nullable(),
    stop_sequence: z.string().nullable().optional(),
  }),
  usage: z
    .object({
      output_tokens: z.number(),
      input_tokens: z.number().optional(),
      cache_creation_input_tokens: z.number().optional(),
      cache_read_input_tokens: z.number().optional(),
    })
    .passthrough(),
});

const SseMessageStopEvent = z.object({
  type: z.literal("message_stop"),
});

const SsePingEvent = z.object({
  type: z.literal("ping"),
});

const SseErrorEvent = z.object({
  type: z.literal("error"),
  error: z.object({
    type: z.string(),
    message: z.string(),
  }),
});

export const SseEventSchema = z.discriminatedUnion("type", [
  SseMessageStartEvent,
  SseContentBlockStartEvent,
  SseContentBlockDeltaEvent,
  SseContentBlockStopEvent,
  SseMessageDeltaEvent,
  SseMessageStopEvent,
  SsePingEvent,
  SseErrorEvent,
]);

export type AnthropicRequest = z.infer<typeof AnthropicRequestSchema>;
export type AnthropicResponse = z.infer<typeof AnthropicResponseSchema>;
export type ResponseContentBlockType = z.infer<typeof ResponseContentBlock>;

// Backward-compatible aliases (originally named Inspector*)
export const InspectorRequestSchema = AnthropicRequestSchema;
export const InspectorResponseSchema = AnthropicResponseSchema;
export type InspectorRequest = AnthropicRequest;
export type InspectorResponse = AnthropicResponse;
