import { z } from "zod";
import {
  type AllowedFilter,
  type AllowedFilters,
  filterSchema,
  filtersParamSchema,
} from "./filters.js";
import {
  type AllowedSort,
  type AllowedSorts,
  sortSchema,
  sortsParamSchema,
} from "./sorts.js";

type AnyParamsSchema = z.ZodType;
type AnyObjectSchema = z.ZodObject<z.ZodRawShape>;
type ObjectShape = Record<string, z.ZodType | undefined>;

export interface LimitParamDefinition {
  default?: number;
  max?: number;
  min?: number;
  description?: string;
}

export interface ParamsDefinitionOptions<
  TSchema extends AnyParamsSchema,
  TF extends AllowedFilters | undefined = AllowedFilters | undefined,
  TS extends AllowedSorts | undefined = AllowedSorts | undefined,
  TL extends LimitParamDefinition | undefined = LimitParamDefinition | undefined,
> {
  schema: TSchema;
  filters?: TF;
  sorts?: TS;
  limit?: TL;
}

type AddFieldIfMissing<
  TShape extends z.ZodRawShape,
  TKey extends string,
  TField extends z.ZodType,
> = TKey extends keyof TShape ? {} : { [K in TKey]: TField };

type FiltersField<
  TShape extends z.ZodRawShape,
  TF extends AllowedFilters | undefined,
> = TF extends AllowedFilters
  ? AddFieldIfMissing<TShape, "filters", z.ZodOptional<typeof filtersParamSchema>>
  : {};

type SortsField<
  TShape extends z.ZodRawShape,
  TS extends AllowedSorts | undefined,
> = TS extends AllowedSorts
  ? AddFieldIfMissing<TShape, "sorts", z.ZodOptional<typeof sortsParamSchema>>
  : {};

type LimitField<
  TShape extends z.ZodRawShape,
  TL extends LimitParamDefinition | undefined,
> = TL extends LimitParamDefinition
  ? TL extends { default: number }
    ? AddFieldIfMissing<TShape, "limit", z.ZodDefault<z.ZodNumber>>
    : AddFieldIfMissing<TShape, "limit", z.ZodOptional<z.ZodNumber>>
  : {};

type ParamsShape<
  TSchema extends AnyParamsSchema,
  TF extends AllowedFilters | undefined,
  TS extends AllowedSorts | undefined,
  TL extends LimitParamDefinition | undefined,
> = TSchema extends z.ZodObject<infer TShape>
  ? z.ZodObject<TShape & FiltersField<TShape, TF> & SortsField<TShape, TS> & LimitField<TShape, TL>>
  : TSchema;

export interface ParamsDefinition<
  TSchema extends AnyParamsSchema = AnyParamsSchema,
  TF extends AllowedFilters | undefined = AllowedFilters | undefined,
  TS extends AllowedSorts | undefined = AllowedSorts | undefined,
  TL extends LimitParamDefinition | undefined = LimitParamDefinition | undefined,
> {
  readonly kind: "params_definition";
  readonly schema: ParamsShape<TSchema, TF, TS, TL>;
  readonly controls: {
    readonly filters?: TF;
    readonly sorts?: TS;
    readonly limit?: TL;
  };
}

export type ParamsLike = z.ZodType | ParamsDefinition;

export type ParamsSchema<P extends ParamsLike> = P extends ParamsDefinition
  ? P["schema"]
  : P extends z.ZodType
    ? P
    : never;

export type InferParams<P extends ParamsLike> = z.infer<ParamsSchema<P>>;

export function defineParams<
  TSchema extends AnyParamsSchema,
  TF extends AllowedFilters | undefined = undefined,
  TS extends AllowedSorts | undefined = undefined,
  TL extends LimitParamDefinition | undefined = undefined,
>(
  options: ParamsDefinitionOptions<TSchema, TF, TS, TL>,
): ParamsDefinition<TSchema, TF, TS, TL> {
  let schema: AnyParamsSchema = options.schema;
  if (options.filters) {
    schema = ensureBuiltInField(schema, "filters", filtersParamSchema.optional());
  }
  if (options.sorts) {
    schema = ensureBuiltInField(schema, "sorts", sortsParamSchema.optional());
  }
  if (options.limit) {
    schema = ensureBuiltInField(schema, "limit", buildLimitField(options.limit));
  }

  return Object.freeze({
    kind: "params_definition" as const,
    schema: schema as ParamsShape<TSchema, TF, TS, TL>,
    controls: Object.freeze({
      filters: options.filters,
      sorts: options.sorts,
      limit: options.limit,
    }),
  });
}

export function isParamsDefinition(value: ParamsLike): value is ParamsDefinition {
  return typeof value === "object" && value !== null && "kind" in value && value.kind === "params_definition";
}

export function getParamsSchema<P extends ParamsLike>(params: P): ParamsSchema<P> {
  return (isParamsDefinition(params) ? params.schema : params) as ParamsSchema<P>;
}

export function getParamsControls(params: ParamsLike): ParamsDefinition["controls"] | undefined {
  return isParamsDefinition(params) ? params.controls : undefined;
}

export function getTopLevelObjectShape(schema: z.ZodType): ObjectShape | undefined {
  return findTopLevelObjectSchema(schema)?.shape as ObjectShape | undefined;
}

function buildLimitField(definition: LimitParamDefinition): z.ZodType {
  let limitSchema = z.number().int();
  if (definition.min !== undefined) {
    limitSchema = limitSchema.min(definition.min);
  }
  if (definition.max !== undefined) {
    limitSchema = limitSchema.max(definition.max);
  }
  return definition.default !== undefined ? limitSchema.default(definition.default) : limitSchema.optional();
}

function ensureBuiltInField(
  schema: AnyParamsSchema,
  key: "filters" | "sorts" | "limit",
  fieldSchema: z.ZodType,
): AnyParamsSchema {
  if (hasTopLevelField(schema, key)) {
    return schema;
  }
  if (schema instanceof z.ZodObject) {
    switch (key) {
      case "filters":
        return schema.extend({ filters: fieldSchema });
      case "sorts":
        return schema.extend({ sorts: fieldSchema });
      case "limit":
        return schema.extend({ limit: fieldSchema });
    }
  }
  throw new Error(
    `defineParams() cannot add built-in ${key} to a wrapped params schema. ` +
      `Declare a top-level "${key}" field in the schema itself before calling defineParams().`,
  );
}

function hasTopLevelField(schema: z.ZodType, key: string): boolean {
  const shape = getTopLevelObjectShape(schema);
  return Boolean(shape?.[key]);
}

function findTopLevelObjectSchema(
  schema: z.ZodType,
  seen: Set<z.ZodType> = new Set(),
): AnyObjectSchema | undefined {
  if (seen.has(schema)) return undefined;
  seen.add(schema);

  if (schema instanceof z.ZodObject) {
    return schema;
  }
  if (schema instanceof z.ZodPipe) {
    return (
      findTopLevelObjectSchema(schema.out as unknown as z.ZodType, seen) ??
      findTopLevelObjectSchema(schema.in as unknown as z.ZodType, seen)
    );
  }
  if (hasUnwrapMethod(schema)) {
    return findTopLevelObjectSchema(schema.unwrap(), seen);
  }
  return undefined;
}

function hasUnwrapMethod(schema: z.ZodType): schema is z.ZodType & { unwrap(): z.ZodType } {
  return typeof (schema as { unwrap?: unknown }).unwrap === "function";
}

export type { AllowedFilter, AllowedFilters, AllowedSort, AllowedSorts };
export { filterSchema, filtersParamSchema, sortSchema, sortsParamSchema };