import {z} from 'zod'

export type BedrockContentBlock = {text?: string | undefined, image?: ({format?: ('jpeg' | 'png' | 'gif' | 'webp') | undefined, source?: ({bytes?: string | undefined}) | undefined}) | undefined, document?: ({format?: ('pdf' | 'csv' | 'doc' | 'docx' | 'xls' | 'xlsx' | 'html' | 'txt' | 'md') | undefined, name?: string | undefined, source?: ({bytes?: string | undefined, text?: string | undefined}) | undefined}) | undefined, toolUse?: ({toolUseId?: string | undefined, name?: string | undefined, input?: Record<string, never> | undefined}) | undefined, toolResult?: ({toolUseId?: string | undefined, content?: Array<BedrockContentBlock> | undefined, status?: ('success' | 'error') | undefined}) | undefined, guardContent?: ({text?: ({text?: string | undefined, qualifiers?: Array<string> | undefined}) | undefined}) | undefined, reasoningContent?: ({reasoningText?: ({text?: string | undefined, signature?: string | undefined}) | undefined}) | undefined, json?: Record<string, never> | undefined, cachePoint?: ({type?: 'default' | undefined}) | undefined};

export type BedrockMessage = {role: 'user' | 'assistant', content: Array<BedrockContentBlock>};

export type BedrockConverseRequest = {messages?: Array<BedrockMessage> | undefined, system?: (Array<{text?: string | undefined, guardContent?: ({text?: ({text?: string | undefined, qualifiers?: Array<string> | undefined}) | undefined}) | undefined, cachePoint?: ({type?: 'default' | undefined}) | undefined}>) | undefined, inferenceConfig?: ({maxTokens?: number | undefined, temperature?: number | undefined, topP?: number | undefined, stopSequences?: Array<string> | undefined}) | undefined, toolConfig?: ({tools?: (Array<{toolSpec?: ({name?: string | undefined, description?: string | undefined, inputSchema?: ({json?: Record<string, never> | undefined}) | undefined}) | undefined, cachePoint?: ({type?: 'default' | undefined}) | undefined}>) | undefined, toolChoice?: ({auto?: Record<string, never> | undefined, any?: Record<string, never> | undefined, tool?: ({name?: string | undefined}) | undefined}) | undefined}) | undefined, guardrailConfig?: ({guardrailIdentifier?: string | undefined, guardrailVersion?: string | undefined, trace?: ('enabled' | 'disabled') | undefined}) | undefined, additionalModelRequestFields?: Record<string, never> | undefined, additionalModelResponseFieldPaths?: Array<string> | undefined, performanceConfig?: ({latency?: ('standard' | 'optimized') | undefined}) | undefined, promptVariables?: ({[key: string]: {text?: string | undefined}}) | undefined, requestMetadata?: {[key: string]: string} | undefined, serviceTier?: ({type?: ('reserved' | 'priority' | 'default' | 'flex') | undefined}) | undefined, fallbacks?: Array<string> | undefined};

export type CreateApiBedrockModelModelIdConverseArgs = {modelId: string, body: BedrockConverseRequest};

export type CreateApiBedrockModelModelIdConverseOptions = {baseUrl?: string | undefined, BearerAuth?: string | undefined, fetch?: typeof fetch | undefined};

export const bedrockContentBlock: z.ZodType<BedrockContentBlock> = z.object({text: z.string().optional(), image: z.object({format: z.enum(["jpeg", "png", "gif", "webp"]).optional(), source: z.object({bytes: z.string().optional()}).optional()}).optional(), document: z.object({format: z.enum(["pdf", "csv", "doc", "docx", "xls", "xlsx", "html", "txt", "md"]).optional(), name: z.string().optional(), source: z.object({bytes: z.string().optional(), text: z.string().optional()}).optional()}).optional(), toolUse: z.object({toolUseId: z.string().optional(), name: z.string().optional(), input: z.object({}).optional()}).optional(), toolResult: z.object({toolUseId: z.string().optional(), content: z.array(z.lazy(() => bedrockContentBlock)).optional(), status: z.enum(["success", "error"]).optional()}).optional(), guardContent: z.object({text: z.object({text: z.string().optional(), qualifiers: z.array(z.string()).optional()}).optional()}).optional(), reasoningContent: z.object({reasoningText: z.object({text: z.string().optional(), signature: z.string().optional()}).optional()}).optional(), json: z.object({}).optional(), cachePoint: z.object({type: z.literal("default").optional()}).optional()});

export const bedrockMessage = z.object({role: z.enum(["user", "assistant"]), content: z.array(bedrockContentBlock)});

export const bedrockConverseResponse = z.object({output: z.object({message: bedrockMessage.optional()}).optional(), stopReason: z.enum(["end_turn", "tool_use", "max_tokens", "stop_sequence", "guardrail_intervened", "content_filtered"]).optional(), usage: z.object({inputTokens: z.number().int().optional(), outputTokens: z.number().int().optional(), totalTokens: z.number().int().optional(), cacheReadInputTokens: z.number().int().optional(), cacheWriteInputTokens: z.number().int().optional()}).optional(), metrics: z.object({latencyMs: z.number().int().optional()}).optional(), additionalModelResponseFields: z.object({}).optional(), trace: z.object({}).optional(), performanceConfig: z.object({latency: z.enum(["standard", "optimized"]).optional()}).optional(), serviceTier: z.object({type: z.enum(["reserved", "priority", "default", "flex"]).optional()}).optional()});

export const validateResponse = async <Output>(
  schema: {
    readonly '~standard': {
      readonly validate: (value: unknown) => ({ readonly value: Output; readonly issues?: undefined } | { readonly issues: ReadonlyArray<{ readonly message: string; readonly path?: ReadonlyArray<PropertyKey | { readonly key: PropertyKey }> | undefined }> }) | Promise<{ readonly value: Output; readonly issues?: undefined } | { readonly issues: ReadonlyArray<{ readonly message: string; readonly path?: ReadonlyArray<PropertyKey | { readonly key: PropertyKey }> | undefined }> }>
    }
  },
  value: unknown
): Promise<Output> => {
  const result = await schema['~standard'].validate(value)

  if (result.issues !== undefined) {
    const detail = result.issues
      .map(({ message, path }) =>
        path?.length
          ? `${path.map((segment) => String(typeof segment === 'object' ? segment.key : segment)).join('.')}: ${message}`
          : message
      )
      .join('; ')

    throw new Error(`Response does not match the schema: ${detail}`, { cause: { issues: result.issues } })
  }

  return result.value
};

/**
 * Converse with model (Bedrock format)
 * 
 * POST /bedrock/model/{modelId}/converse
 * 
 * Also accepts auth: BasicAuth, VirtualKeyAuth, ApiKeyAuth
 */
export const createApiBedrockModelModelIdConverse = async (args: CreateApiBedrockModelModelIdConverseArgs, options?: CreateApiBedrockModelModelIdConverseOptions) => {
const {baseUrl = "http://localhost:8080", fetch: fetchFn = fetch, BearerAuth} = options ?? {}

const query = new URLSearchParams()
const headers = new Headers()
headers.set('content-type', "application/json")
headers.set('accept', "application/json")

if (BearerAuth !== undefined) headers.set('authorization', `Bearer ${BearerAuth}`)

const url = `${baseUrl.replace(/\/+$/, '')}/bedrock/model/${encodeURIComponent(String(args.modelId))}/converse${query.size > 0 ? `?${query}` : ''}`

const res = await fetchFn(url, {method: "POST", headers, body: JSON.stringify(args.body)})

if (!res.ok) throw new Error(`Request failed with status ${res.status}`, {cause: {status: res.status, body: await res.text()}})

return validateResponse(bedrockConverseResponse, await res.json())
};