import type { JSONSchema7 } from "@ai-sdk/provider"
import type { JSONSchema } from "zod/v4/core"
import type * as Provider from "./provider"
import { OUTPUT_TOKEN_MAX } from "./transform-params"

export function maxOutputTokens(model: Provider.Model): number {
  return Math.min(model.limit.output, OUTPUT_TOKEN_MAX) || OUTPUT_TOKEN_MAX
}

export function schema(model: Provider.Model, schema: JSONSchema.BaseSchema | JSONSchema7): JSONSchema7 {
  if (model.providerID === "moonshotai" || model.api.id.toLowerCase().includes("kimi")) {
    const sanitizeMoonshot = (obj: unknown): unknown => {
      if (obj === null || typeof obj !== "object") return obj
      if (Array.isArray(obj)) return obj.map(sanitizeMoonshot)
      if ("$ref" in obj && typeof (obj as any).$ref === "string") return { $ref: (obj as any).$ref }
      const result = Object.fromEntries(Object.entries(obj).map(([key, value]) => [key, sanitizeMoonshot(value)]))
      if (Array.isArray((result as any).items)) (result as any).items = (result as any).items[0] ?? {}
      return result
    }
    schema = sanitizeMoonshot(schema) as JSONSchema.BaseSchema | JSONSchema7
  }

  if (model.providerID === "google" || model.api.id.includes("gemini")) {
    const isPlainObject = (node: unknown): node is Record<string, any> => typeof node === "object" && node !== null && !Array.isArray(node)
    const hasCombiner = (node: unknown) => isPlainObject(node) && (Array.isArray(node.anyOf) || Array.isArray(node.oneOf) || Array.isArray(node.allOf))
    const hasSchemaIntent = (node: unknown) => {
      if (!isPlainObject(node)) return false
      if (hasCombiner(node)) return true
      return ["type", "properties", "items", "prefixItems", "enum", "const", "$ref", "additionalProperties", "patternProperties", "required", "not", "if", "then", "else"].some((key) => key in node)
    }
    const sanitizeGemini = (obj: any): any => {
      if (obj === null || typeof obj !== "object") return obj
      if (Array.isArray(obj)) return obj.map(sanitizeGemini)
      const result: any = {}
      for (const [key, value] of Object.entries(obj)) {
        if (key === "enum" && Array.isArray(value)) {
          result[key] = value.map((v: any) => String(v))
          if (result.type === "integer" || result.type === "number") result.type = "string"
        } else if (typeof value === "object" && value !== null) result[key] = sanitizeGemini(value)
        else result[key] = value
      }
      if (result.type === "object" && result.properties && Array.isArray(result.required)) result.required = result.required.filter((field: any) => field in result.properties)
      if (result.type === "array" && !hasCombiner(result)) {
        if (result.items == null) result.items = {}
        if (isPlainObject(result.items) && !hasSchemaIntent(result.items)) result.items.type = "string"
      }
      if (result.type && result.type !== "object" && !hasCombiner(result)) { delete result.properties; delete result.required }
      return result
    }
    schema = sanitizeGemini(schema)
  }

  return schema as JSONSchema7
}
