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

// OpenAI Chat Completions Schemas

export const OpenAIMessageContent = z.union([
  z.string(),
  z.array(
    z.discriminatedUnion("type", [
      z.object({ type: z.literal("text"), text: z.string() }),
      z.object({
        type: z.literal("image_url"),
        image_url: z.object({ url: z.string(), detail: z.string().optional() }),
      }),
    ]),
  ),
]);

export const OpenAIMessage = z.object({
  role: z.enum(["system", "user", "assistant", "tool"]),
  content: OpenAIMessageContent,
  name: z.string().optional(),
  reasoning_content: z.string().optional(),
  thinking: z.string().optional(), // Some providers use 'thinking' field
  think: z.string().optional(), // MiniMax uses 'think' field
});

export const OpenAIFunctionCall = z.object({
  name: z.string(),
  arguments: z.string(),
});

/** A single OpenAI tool_call (the modern multi-tool form, returned by most
 *  providers in `choices[].message.tool_calls`). `function.arguments` is
 *  a JSON-encoded string per the OpenAI spec — we keep it as a string at
 *  parse time and let the renderer decide whether to JSON-parse it for
 *  display. */
export const OpenAIToolCallSchema = z.object({
  index: z.number().optional(),
  id: z.string().optional(),
  type: z.literal("function").optional(),
  function: z.object({
    name: z.string().optional(),
    arguments: z.string().optional(),
  }),
});

export type OpenAIToolCall = z.infer<typeof OpenAIToolCallSchema>;

const OpenAIMessageWithFunctionCall = OpenAIMessage.extend({
  content: z
    .union([z.string(), z.array(z.object({ type: z.literal("text"), text: z.string() }))])
    .optional(),
  function_call: OpenAIFunctionCall.optional(),
  tool_calls: z.array(OpenAIToolCallSchema).optional(),
});

export const OpenAIToolDefinition = z.object({
  type: z.literal("function"),
  function: z.object({
    name: z.string(),
    description: z.string().optional(),
    parameters: z.record(z.string(), JsonValueSchema),
  }),
});

export const OpenAIRequestSchema = z.object({
  model: z.string(),
  messages: z.array(OpenAIMessage),
  temperature: z.number().optional(),
  max_tokens: z.number().optional(),
  stream: z.boolean().optional(),
  tools: z.array(OpenAIToolDefinition).optional(),
  tool_choice: z
    .union([
      z.enum(["auto", "none", "required"]),
      z.object({ type: z.literal("auto") }),
      z.object({ type: z.literal("none") }),
      z.object({ type: z.literal("function"), function: z.object({ name: z.string() }) }),
    ])
    .optional(),
  user: z.string().optional(),
});

export const OpenAIChoiceDelta = z.object({
  role: z.enum(["assistant"]).optional(),
  content: z.string().nullable().optional(),
  reasoning_content: z.string().nullable().optional(),
  thinking: z.string().nullable().optional(), // Some providers use 'thinking' field
  think: z.string().nullable().optional(), // MiniMax uses 'think' field
  function_call: z
    .object({ name: z.string().optional(), arguments: z.string().optional() })
    .nullable()
    .optional(),
  tool_calls: z
    .array(
      z.object({
        index: z.number(),
        id: z.string().optional(),
        type: z.literal("function").optional(),
        function: z.object({
          name: z.string().optional(),
          arguments: z.string().optional(),
        }),
      }),
    )
    .nullable()
    .optional(),
});

export const OpenAIChoice = z.object({
  index: z.number(),
  message: z
    .object({
      role: z.enum(["assistant"]),
      content: z.string().nullable(),
      reasoning_content: z.string().optional(),
      thinking: z.string().optional(), // Some providers use 'thinking' field in message
      think: z.string().optional(), // MiniMax uses 'think' field in message
      function_call: z.object({ name: z.string(), arguments: z.string() }).nullable().optional(),
      tool_calls: z.array(OpenAIToolCallSchema).optional(),
    })
    .optional(),
  delta: OpenAIChoiceDelta.optional(),
  finish_reason: z.string().nullable(),
});

export const OpenAIResponseSchema = z
  .object({
    id: z.string(),
    object: z.literal("chat.completion"),
    created: z.number(),
    model: z.string(),
    choices: z.array(OpenAIChoice),
    usage: z
      .object({
        prompt_tokens: z.number().nullable().optional(),
        completion_tokens: z.number().nullable().optional(),
        total_tokens: z.number().nullable().optional(),
      })
      .passthrough(), // Allow extra fields like prompt_tokens_details, completion_tokens_details
    service_tier: z.string().nullable().optional(),
    system_fingerprint: z.string().nullable().optional(),
  })
  .passthrough(); // Allow extra fields like refusal, annotations, audio

// Schema for actual OpenAI SSE chunk format (not wrapped in event/data)
export const OpenAISSERawChunkSchema = z
  .object({
    id: z.string(),
    object: z.literal("chat.completion.chunk"),
    created: z.number(),
    model: z.string(),
    choices: z.array(
      z.object({
        index: z.number(),
        delta: OpenAIChoiceDelta,
        finish_reason: z.string().nullable().optional(),
      }),
    ),
    usage: z
      .object({
        prompt_tokens: z.number().nullable().optional(),
        completion_tokens: z.number().nullable().optional(),
        total_tokens: z.number().nullable().optional(),
      })
      .passthrough()
      .nullable()
      .optional(),
  })
  .passthrough(); // Allow extra fields in chunks

export type OpenAIResponse = z.infer<typeof OpenAIResponseSchema>;

/**
 * Parse an OpenAI response body.
 */
export function parseOpenAIResponse(rawBody: string): OpenAIResponse | null {
  try {
    const json: unknown = JSON.parse(rawBody);
    const result = OpenAIResponseSchema.safeParse(json);
    if (result.success) return result.data;
    return null;
  } catch {
    return null;
  }
}
