import { z } from "zod";

export const RoleSchema = z.enum(["user", "model", "function"]);

const TextPartSchema = z
  .object({
    text: z.string(),
  })
  .describe("Inline text");

const InlineDataPartSchema = z
  .object({
    inlineData: z.object({
      mimeType: z
        .string()
        .describe(
          "The IANA standard MIME type of the source data. Examples: - image/png - image/jpeg If an unsupported MIME type is provided, an error will be returned",
        ),
      data: z.string().describe("Raw bytes for media formats. Base64 encoded"),
    }),
  })
  .describe(`https://ai.google.dev/api/caching#Blob`);

const FunctionCallPartSchema = z
  .object({
    functionCall: z.object({
      id: z
        .string()
        .optional()
        .describe(
          "Optional. The unique id of the function call. If populated, the client to execute the functionCall and return the response with the matching id.",
        ),
      name: z.string(
        "The name of the function to call. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 63",
      ),
      args: z
        .record(z.string(), z.any())
        .optional()
        .describe("The function parameters and values in JSON object format."),
    }),
  })
  .describe(`https://ai.google.dev/api/caching#FunctionCall`);

const FunctionResponsePartSchema = z
  .object({
    functionResponse: z.object({
      id: z
        .string()
        .optional()
        .describe(
          "The id of the function call this response is for. Populated by the client to match the corresponding function call id",
        ),
      name: z
        .string()
        .describe(
          "The name of the function to call. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 63.",
        ),
      response: z
        .record(z.string(), z.any())
        .describe("The function response in JSON object format."),
      willContinue: z
        .boolean()
        .optional()
        .describe(
          "Signals that function call continues, and more responses will be returned, turning the function call into a generator. Is only applicable to NON_BLOCKING function calls, is ignored otherwise. If set to false, future responses will not be considered. It is allowed to return empty response with willContinue=False to signal that the function call is finished. This may still trigger the model generation. To avoid triggering the generation and finish the function call, additionally set scheduling to SILENT",
        ),
      scheduling: z
        .enum(["SCHEDULING_UNSPECIFIED", "SILENT", "WHEN_IDLE", "INTERRUPT"])
        .optional()
        .describe(`
          Specifies how the response should be scheduled in the conversation. Only applicable to NON_BLOCKING function calls, is ignored otherwise. Defaults to WHEN_IDLE.

          https://ai.google.dev/api/caching#Scheduling
        `),
    }),
  })
  .describe(`https://ai.google.dev/api/caching#FunctionResponse`);

const FileDataPartSchema = z
  .object({
    fileData: z.object({
      mimeType: z
        .string()
        .optional()
        .describe("The IANA standard MIME type of the source data"),
      fileUri: z.string().describe("URI"),
    }),
  })
  .describe(`https://ai.google.dev/api/caching#FileData`);

const ExecutableCodePartSchema = z
  .object({
    language: z.enum(["LANGUAGE_UNSPECIFIED", "PYTHON"]).describe(`
      Programming language of the code

      https://ai.google.dev/api/caching#Language
    `),
    executableCode: z.object({
      code: z.string().describe("The code to be executed"),
    }),
  })
  .describe(`
Code generated by the model that is meant to be executed, and the result returned to the model.

Only generated when using the CodeExecution tool, in which the code will be automatically executed, and a corresponding CodeExecutionResult will also be generated

https://ai.google.dev/api/caching#ExecutableCode
`);

const CodeExecutionResultPartSchema = z
  .object({
    codeExecutionResult: z.object({
      outcome: z
        .enum([
          "OUTCOME_UNSPECIFIED",
          "OUTCOME_OK",
          "OUTCOME_FAILED",
          "OUTCOME_DEADLINE_EXCEEDED",
        ])
        .describe(`
          Outcome of the code execution.

          https://ai.google.dev/api/caching#Outcome
        `),
      output: z
        .string()
        .optional()
        .describe(
          "Contains stdout when code execution is successful, stderr or other description otherwise",
        ),
    }),
  })
  .describe(`
Result of executing the ExecutableCode.

Only generated when using the CodeExecution, and always follows a part containing the ExecutableCode.

https://ai.google.dev/api/caching#CodeExecutionResult
`);

const VideoMetadataSchema = z
  .object({
    startOffset: z
      .string()
      .optional()
      .describe(
        "The start offset of the video. A duration in seconds with up to nine fractional digits, ending with 's'. Example: '3.5s'",
      ),
    endOffset: z
      .string()
      .optional()
      .describe(
        "The end offset of the video. A duration in seconds with up to nine fractional digits, ending with 's'. Example: '3.5s'",
      ),
    fps: z
      .number()
      .optional()
      .describe(
        "The frame rate of the video sent to the model. If not specified, the default value will be 1.0. The fps range is (0.0, 24.0]",
      ),
  })
  .describe(`https://ai.google.dev/api/caching#VideoMetadata`);

const _PartDataSchema = z
  .union([
    TextPartSchema,
    InlineDataPartSchema,
    FunctionCallPartSchema,
    FunctionResponsePartSchema,
    FileDataPartSchema,
    ExecutableCodePartSchema,
    CodeExecutionResultPartSchema,
  ])
  .describe(`https://ai.google.dev/api/caching#Part`);

const PartMetadataSchema = z
  .union([VideoMetadataSchema])
  .optional()
  .describe(`https://ai.google.dev/api/caching#PartMetadata`);

const BasePartFields = {
  thought: z
    .boolean()
    .optional()
    .describe(`Indicates if the part is thought from the model`),
  thoughtSignature: z
    .string()
    .optional()
    .describe(
      "An opaque signature for the thought so it can be reused in subsequent requests. A base64-encoded string",
    ),
};

const TextPartTopLevel = z.object({
  ...BasePartFields,
  text: z.string(),
  metadata: PartMetadataSchema,
});

const InlineDataPartTopLevel = z.object({
  ...BasePartFields,
  inlineData: z.object({
    mimeType: z.string().optional(),
    data: z.string(),
  }),
  metadata: PartMetadataSchema,
});

const FunctionCallPartTopLevel = z.object({
  ...BasePartFields,
  functionCall: z.object({
    id: z.string().optional(),
    name: z.string(),
    args: z.record(z.string(), z.any()).optional(),
  }),
  metadata: PartMetadataSchema,
});

const FunctionResponsePartTopLevel = z.object({
  ...BasePartFields,
  functionResponse: z.object({
    id: z.string().optional(),
    name: z.string(),
    response: z.record(z.string(), z.any()),
    willContinue: z.boolean().optional(),
    scheduling: z
      .enum(["SCHEDULING_UNSPECIFIED", "SILENT", "WHEN_IDLE", "INTERRUPT"])
      .optional(),
  }),
  metadata: PartMetadataSchema,
});

const FileDataPartTopLevel = z.object({
  ...BasePartFields,
  fileData: z.object({
    mimeType: z.string().optional(),
    fileUri: z.string(),
  }),
  metadata: PartMetadataSchema,
});

const ExecutableCodePartTopLevel = z.object({
  ...BasePartFields,
  language: z.enum(["LANGUAGE_UNSPECIFIED", "PYTHON"]),
  executableCode: z.object({ code: z.string() }),
  metadata: PartMetadataSchema,
});

const CodeExecutionResultPartTopLevel = z.object({
  ...BasePartFields,
  codeExecutionResult: z.object({
    outcome: z
      .enum([
        "OUTCOME_UNSPECIFIED",
        "OUTCOME_OK",
        "OUTCOME_FAILED",
        "OUTCOME_DEADLINE_EXCEEDED",
      ])
      .describe("Outcome of the code execution."),
    output: z.string().optional(),
  }),
  metadata: PartMetadataSchema,
});

const ShorthandPartSchema = z.union([
  TextPartTopLevel,
  InlineDataPartTopLevel,
  FunctionCallPartTopLevel,
  FunctionResponsePartTopLevel,
  FileDataPartTopLevel,
  ExecutableCodePartTopLevel,
  CodeExecutionResultPartTopLevel,
]);

export const PartSchema = z
  .union([ShorthandPartSchema])
  .describe(`https://ai.google.dev/api/caching#Part`);

export const ContentSchema = z
  .object({
    role: z.string().describe("The role of the author of this content."),
    parts: z.array(PartSchema),
  })
  .describe(`
The base structured datatype containing multi-part content of a message.

A Content includes a role field designating the producer of the Content and a parts field containing multi-part data that contains the content of the message turn.

https://ai.google.dev/api/caching#Content
`);
