// Cloud Code Assist API rejects a subset of JSON Schema keywords.
// This module scrubs/normalizes tool schemas to keep Gemini happy.

// Keywords that Cloud Code Assist API rejects (not compliant with their JSON Schema subset)
const UNSUPPORTED_SCHEMA_KEYWORDS = new Set([
  "patternProperties",
  "additionalProperties",
  "$schema",
  "$id",
  "$ref",
  "$defs",
  "definitions",
]);

// Check if an anyOf/oneOf array contains only literal values that can be flattened.
// TypeBox Type.Literal generates { const: "value", type: "string" }.
// Some schemas may use { enum: ["value"], type: "string" }.
// Both patterns are flattened to { type: "string", enum: ["a", "b", ...] }.
function tryFlattenLiteralAnyOf(
  variants: unknown[],
): { type: string; enum: unknown[] } | null {
  if (variants.length === 0) return null;

  const allValues: unknown[] = [];
  let commonType: string | null = null;

  for (const variant of variants) {
    if (!variant || typeof variant !== "object") return null;
    const v = variant as Record<string, unknown>;

    let literalValue: unknown;
    if ("const" in v) {
      literalValue = v.const;
    } else if (Array.isArray(v.enum) && v.enum.length === 1) {
      literalValue = v.enum[0];
    } else {
      return null;
    }

    const variantType = typeof v.type === "string" ? v.type : null;
    if (!variantType) return null;
    if (commonType === null) commonType = variantType;
    else if (commonType !== variantType) return null;

    allValues.push(literalValue);
  }

  if (commonType && allValues.length > 0)
    return { type: commonType, enum: allValues };
  return null;
}

type SchemaDefs = Map<string, unknown>;

function extendSchemaDefs(
  defs: SchemaDefs | undefined,
  schema: Record<string, unknown>,
): SchemaDefs | undefined {
  const defsEntry =
    schema.$defs &&
    typeof schema.$defs === "object" &&
    !Array.isArray(schema.$defs)
      ? (schema.$defs as Record<string, unknown>)
      : undefined;
  const legacyDefsEntry =
    schema.definitions &&
    typeof schema.definitions === "object" &&
    !Array.isArray(schema.definitions)
      ? (schema.definitions as Record<string, unknown>)
      : undefined;

  if (!defsEntry && !legacyDefsEntry) return defs;

  const next = defs ? new Map(defs) : new Map<string, unknown>();
  if (defsEntry) {
    for (const [key, value] of Object.entries(defsEntry)) next.set(key, value);
  }
  if (legacyDefsEntry) {
    for (const [key, value] of Object.entries(legacyDefsEntry))
      next.set(key, value);
  }
  return next;
}

function decodeJsonPointerSegment(segment: string): string {
  return segment.replaceAll("~1", "/").replaceAll("~0", "~");
}

function tryResolveLocalRef(
  ref: string,
  defs: SchemaDefs | undefined,
): unknown {
  if (!defs) return undefined;
  const match = ref.match(/^#\/(?:\$defs|definitions)\/(.+)$/);
  if (!match) return undefined;
  const name = decodeJsonPointerSegment(match[1] ?? "");
  if (!name) return undefined;
  return defs.get(name);
}

function cleanSchemaForGeminiWithDefs(
  schema: unknown,
  defs: SchemaDefs | undefined,
  refStack: Set<string> | undefined,
): unknown {
  if (!schema || typeof schema !== "object") return schema;
  if (Array.isArray(schema)) {
    return schema.map((item) =>
      cleanSchemaForGeminiWithDefs(item, defs, refStack),
    );
  }

  const obj = schema as Record<string, unknown>;
  const nextDefs = extendSchemaDefs(defs, obj);

  const refValue = typeof obj.$ref === "string" ? obj.$ref : undefined;
  if (refValue) {
    if (refStack?.has(refValue)) return {};

    const resolved = tryResolveLocalRef(refValue, nextDefs);
    if (resolved) {
      const nextRefStack = refStack ? new Set(refStack) : new Set<string>();
      nextRefStack.add(refValue);

      const cleaned = cleanSchemaForGeminiWithDefs(
        resolved,
        nextDefs,
        nextRefStack,
      );
      if (!cleaned || typeof cleaned !== "object" || Array.isArray(cleaned)) {
        return cleaned;
      }

      const result: Record<string, unknown> = {
        ...(cleaned as Record<string, unknown>),
      };
      for (const key of ["description", "title", "default", "examples"]) {
        if (key in obj && obj[key] !== undefined) result[key] = obj[key];
      }
      return result;
    }

    const result: Record<string, unknown> = {};
    for (const key of ["description", "title", "default", "examples"]) {
      if (key in obj && obj[key] !== undefined) result[key] = obj[key];
    }
    return result;
  }

  const hasAnyOf = "anyOf" in obj && Array.isArray(obj.anyOf);
  const hasOneOf = "oneOf" in obj && Array.isArray(obj.oneOf);

  if (hasAnyOf) {
    const flattened = tryFlattenLiteralAnyOf(obj.anyOf as unknown[]);
    if (flattened) {
      const result: Record<string, unknown> = {
        type: flattened.type,
        enum: flattened.enum,
      };
      for (const key of ["description", "title", "default", "examples"]) {
        if (key in obj && obj[key] !== undefined) result[key] = obj[key];
      }
      return result;
    }
  }

  if (hasOneOf) {
    const flattened = tryFlattenLiteralAnyOf(obj.oneOf as unknown[]);
    if (flattened) {
      const result: Record<string, unknown> = {
        type: flattened.type,
        enum: flattened.enum,
      };
      for (const key of ["description", "title", "default", "examples"]) {
        if (key in obj && obj[key] !== undefined) result[key] = obj[key];
      }
      return result;
    }
  }

  const cleaned: Record<string, unknown> = {};

  for (const [key, value] of Object.entries(obj)) {
    if (UNSUPPORTED_SCHEMA_KEYWORDS.has(key)) continue;

    if (key === "const") {
      cleaned.enum = [value];
      continue;
    }

    if (key === "type" && (hasAnyOf || hasOneOf)) continue;

    if (key === "properties" && value && typeof value === "object") {
      const props = value as Record<string, unknown>;
      cleaned[key] = Object.fromEntries(
        Object.entries(props).map(([k, v]) => [
          k,
          cleanSchemaForGeminiWithDefs(v, nextDefs, refStack),
        ]),
      );
    } else if (key === "items" && value && typeof value === "object") {
      cleaned[key] = cleanSchemaForGeminiWithDefs(value, nextDefs, refStack);
    } else if (key === "anyOf" && Array.isArray(value)) {
      cleaned[key] = value.map((variant) =>
        cleanSchemaForGeminiWithDefs(variant, nextDefs, refStack),
      );
    } else if (key === "oneOf" && Array.isArray(value)) {
      cleaned[key] = value.map((variant) =>
        cleanSchemaForGeminiWithDefs(variant, nextDefs, refStack),
      );
    } else if (key === "allOf" && Array.isArray(value)) {
      cleaned[key] = value.map((variant) =>
        cleanSchemaForGeminiWithDefs(variant, nextDefs, refStack),
      );
    } else {
      cleaned[key] = value;
    }
  }

  return cleaned;
}

export function cleanSchemaForGemini(schema: unknown): unknown {
  if (!schema || typeof schema !== "object") return schema;
  if (Array.isArray(schema)) return schema.map(cleanSchemaForGemini);

  const defs = extendSchemaDefs(undefined, schema as Record<string, unknown>);
  return cleanSchemaForGeminiWithDefs(schema, defs, undefined);
}