import type { JSONSchema } from '@squiz/json-schema-library';
import * as t from 'ts-brand';

type MaybePromise<T> = T | Promise<T>;

type JsonResolutionSchema<TITLE extends string> = JSONSchema & { title: TITLE };
/**
 * This type allows the TypeScript type to be encoded onto the JSON schema object
 */
type SchemaWithShape<TITLE extends string, SHAPE> = JsonResolutionSchema<TITLE> & { __shape__: SHAPE };

/**
 * A JSON schema which represents a primitive type which can be a resolve target
 *
 * The brand ensures that TypeScript can differentiate between a primitive type and a resolvable type
 */
export type PrimitiveType<TITLE extends string, SHAPE> = t.Brand<SchemaWithShape<TITLE, SHAPE>, 'primitive'>;
export function PrimitiveType<SHAPE, TITLE extends string>(
  jsonSchema: JsonResolutionSchema<TITLE>,
): PrimitiveType<TITLE, SHAPE> {
  return jsonSchema as PrimitiveType<TITLE, SHAPE>;
}
export type AnyPrimitiveType = PrimitiveType<string, any>;

/**
 * A JSON schema which represents a type which can be resolved into a primitive type
 *
 * The brand ensures that TypeScript can differentiate between a primitive type and a resolvable type
 */
export type ResolvableType<TITLE extends string, SHAPE> = t.Brand<SchemaWithShape<TITLE, SHAPE>, 'resolvable'>;
export function ResolvableType<SHAPE, TITLE extends string>(
  jsonSchema: JsonResolutionSchema<TITLE>,
): ResolvableType<TITLE, SHAPE> {
  return jsonSchema as ResolvableType<TITLE, SHAPE>;
}
export type AnyResolvableType = ResolvableType<string, any>;

export type ResolverContext = {
  baseUrl?: string;
  matrix?: {
    context?: string;
  };
};

export type Resolver<INPUT, OUTPUT> = (input: INPUT, ctx?: ResolverContext) => MaybePromise<OUTPUT>;

/**
 * A JSON Type Resolver class which stores the primitive and resolvable JSON Schema types and their resolvers
 *
 * No serious logic is required here. The class should only provide data access methods and type safety
 */
export class TypeResolver<P extends AnyPrimitiveType, R extends AnyResolvableType> {
  private primitives: Map<P['title'], P>;
  private resolvables: Map<R['title'], R>;
  constructor(
    primitives: P[],
    resolvables: R[] = [],
    public resolvers: {
      [PT in P as PT['title']]?: {
        [RT in R as RT['title']]?: Resolver<RT['__shape__'], PT['__shape__']>;
      };
    } = {},
  ) {
    this.primitives = new Map(primitives.map((p) => [p.title, p]));
    this.resolvables = new Map(resolvables.map((r) => [r.title, r]));
    for (const [primitiveKey, primitiveResolvers] of Object.entries(resolvers) as [string, Record<string, any>][]) {
      if (!this.primitives.has(primitiveKey)) {
        throw new Error('Resolver keys must match a primitive schema');
      }
      if (!Object.keys(primitiveResolvers).every((k) => this.resolvables.has(k))) {
        throw new Error('Primitive resolvers keys must match a resolvable schema');
      }
    }
  }

  get validationSchemaDefinitions() {
    return [...this.primitives.values(), ...this.resolvables.values()];
  }

  isPrimitiveType(type: string): type is P['title'] {
    return this.primitives.has(type);
  }

  isResolvableSchema(schema: JSONSchema): schema is R {
    return this.resolvables.has(schema.title);
  }

  getValidationSchemaForPrimitive(type: P['title']) {
    const primitiveSchema = this.primitives.get(type) as JsonResolutionSchema<string>;
    return [primitiveSchema, ...this.fetchResolvableSchemasForPrimitive(type)];
  }

  private *fetchResolvableSchemasForPrimitive(type: P['title']) {
    for (const resolverKey in this.resolvers[type]) {
      yield this.resolvables.get(resolverKey) as JsonResolutionSchema<string>;
    }
  }

  tryGetResolver(
    primitiveSchemaTitle: string,
    resolvableSchema: R,
  ): Resolver<R['__shape__'], P['__shape__']> | undefined {
    if (!(primitiveSchemaTitle in this.resolvers)) return;
    // Sometimes typescript can be insanely annoying
    return (this.resolvers[primitiveSchemaTitle as keyof typeof this.resolvers] as any)?.[resolvableSchema.title];
  }
}