import * as BindingsLeaves from './leaves.js';
import { FieldsDecoder, ProvableSerializable } from './util.js';
import { versionBytes } from '../../crypto/constants.js';
import { Provable } from '../../../lib/provable/provable.js';
import { HashInput } from '../../../lib/provable/types/provable-derivers.js';
import { toBase58Check } from '../../../lib/util/base58.js';

const JsArray = Array;

abstract class ProvableBindingsType<T, Actual> {
  abstract Type(): ProvableSerializable<Actual>;

  sizeInFields(): number {
    return this.Type().sizeInFields();
  }

  toJSON(x: T): any {
    return this.Type().toJSON(x as never as Actual);
  }

  toInput(x: T): HashInput {
    return this.Type().toInput(x as never as Actual);
  }

  toFields(x: T): BindingsLeaves.Field[] {
    return this.Type().toFields(x as never as Actual);
  }

  toAuxiliary(x?: T): any[] {
    return this.Type().toAuxiliary(x as never as Actual);
  }

  fromFields(fields: BindingsLeaves.Field[], aux: any[]): T {
    return this.Type().fromFields(fields, aux) as never as T;
  }

  toValue(x: T): T {
    return x;
  }

  fromValue(x: T): T {
    return x;
  }

  check(x: T) {
    return this.Type().check(x as never as Actual);
  }
}

export type BindingsType<T> =
  | BindingsType.Leaf<T>
  | BindingsType.Object<T>
  | BindingsType.Option<T>
  | BindingsType.Array<T>;

function assertBindingsTypeImplementsProvable<
  T,
  B extends BindingsType<T> & ProvableSerializable<T>,
>(_x?: B) {}

assertBindingsTypeImplementsProvable<number, BindingsType<number>>();
assertBindingsTypeImplementsProvable<string, BindingsType<string>>();
assertBindingsTypeImplementsProvable<
  BindingsLeaves.AuthRequired,
  BindingsType<BindingsLeaves.AuthRequired>
>();
assertBindingsTypeImplementsProvable<BindingsLeaves.Bool, BindingsType<BindingsLeaves.Bool>>();
assertBindingsTypeImplementsProvable<BindingsLeaves.Field, BindingsType<BindingsLeaves.Field>>();
assertBindingsTypeImplementsProvable<
  BindingsLeaves.PublicKey,
  BindingsType<BindingsLeaves.PublicKey>
>();
assertBindingsTypeImplementsProvable<BindingsLeaves.Sign, BindingsType<BindingsLeaves.Sign>>();
assertBindingsTypeImplementsProvable<
  BindingsLeaves.TokenId,
  BindingsType<BindingsLeaves.TokenId>
>();
assertBindingsTypeImplementsProvable<
  BindingsLeaves.TokenSymbol,
  BindingsType<BindingsLeaves.TokenSymbol>
>();
assertBindingsTypeImplementsProvable<BindingsLeaves.UInt32, BindingsType<BindingsLeaves.UInt32>>();
assertBindingsTypeImplementsProvable<BindingsLeaves.UInt64, BindingsType<BindingsLeaves.UInt64>>();
assertBindingsTypeImplementsProvable<
  BindingsLeaves.ZkappUri,
  BindingsType<BindingsLeaves.ZkappUri>
>();
assertBindingsTypeImplementsProvable<{ x: number }, BindingsType<{ x: number }>>();
assertBindingsTypeImplementsProvable<number[], BindingsType<number[]>>();
assertBindingsTypeImplementsProvable<
  BindingsLeaves.Option<number>,
  BindingsType<BindingsLeaves.Option<number>>
>();
assertBindingsTypeImplementsProvable<
  BindingsLeaves.Option<BindingsLeaves.Range<number>>,
  BindingsType<BindingsLeaves.Option<BindingsLeaves.Range<number>>>
>();

export namespace BindingsType {
  export class Object<T> implements Provable<T> {
    readonly _T!: T extends { [key: string]: any } ? void : never;
    readonly name: string;
    readonly keys: (keyof T)[];
    readonly entries: T extends { [key: string]: any }
      ? { [key in keyof T]: BindingsType<T[key]> }
      : never;

    constructor({
      name,
      keys,
      entries,
    }: {
      name: Object<T>['name'];
      keys: Object<T>['keys'];
      entries: Object<T>['entries'];
    }) {
      this.name = name;
      this.keys = keys;
      this.entries = entries;
    }

    sizeInFields(): number {
      let sum = 0;
      for (const key of this.keys) {
        sum += this.entries[key].sizeInFields();
      }
      return sum;
    }

    toJSON(x: T): any {
      // TODO: type safety
      const x2 = x as { [key in keyof T]: any };
      const json: Partial<T> = {};
      for (const key of this.keys) {
        json[key] = this.entries[key].toJSON(x2[key]);
      }
      return json;
    }

    toInput(x: T): HashInput {
      // TODO: type safety
      const x2 = x as { [key in keyof T]: any };
      const acc: HashInput = { fields: [], packed: [] };
      for (const key of this.keys) {
        // surely there is an optimization here to avoid allocating so many temporary arrays
        const { fields, packed } = this.entries[key].toInput(x2[key]);
        acc.fields!.push(...(fields ?? []));
        acc.packed!.push(...(packed ?? []));
      }
      return acc;
    }

    toFields(x: T): BindingsLeaves.Field[] {
      // TODO: type safety
      const x2 = x as { [key in keyof T]: any };
      return this.keys.map((key) => this.entries[key].toFields(x2[key])).flat();
    }

    toAuxiliary(x?: T): any[] {
      // TODO: type safety
      const x2 = x as { [key in keyof T]: any } | undefined;
      const entries2 = this.entries as { [key in keyof T]: BindingsType<any> };
      return this.keys.map((key) =>
        entries2[key].toAuxiliary(x2 !== undefined ? x2[key] : undefined)
      );
    }

    fromFields(fields: BindingsLeaves.Field[], aux: any[]): T {
      const decoder = new FieldsDecoder(fields);
      // TODO: make this type-safe
      // const obj: Partial<T> = {};
      const obj: any = {};

      for (const i in this.keys) {
        const key = this.keys[i];
        const entryType = this.entries[key];
        const entryAux = aux[i];
        // console.log(`${this.name}[${JSON.stringify(key)}] :: aux = ${JSON.stringify(entryAux)}`);
        obj[key] = decoder.decode(entryType.sizeInFields(), (entryFields) =>
          entryType.fromFields(entryFields, entryAux)
        );
      }

      return obj;
    }

    toValue(x: T): T {
      return x;
    }

    fromValue(x: T): T {
      return x;
    }

    check(_x: T) {
      throw new Error('TODO');
    }
  }

  export class Array<T> implements Provable<T> {
    readonly _T!: T extends any[] ? void : never;
    readonly staticLength: number | null;
    readonly inner: T extends (infer U)[] ? BindingsType<U> : never;

    constructor({
      staticLength,
      inner,
    }: {
      staticLength: Array<T>['staticLength'];
      inner: Array<T>['inner'];
    }) {
      this.staticLength = staticLength;
      this.inner = inner;
    }

    sizeInFields(): number {
      if (this.staticLength !== null) {
        return this.staticLength * this.inner.sizeInFields();
      } else {
        return 0;
      }
    }

    toJSON(x: T extends any[] ? T : never): any {
      // TODO: type safety
      const inner: BindingsType<any> = this.inner;
      return x.map((el) => inner.toJSON(el));
    }

    toInput(x: T): HashInput {
      if (!(x instanceof JsArray)) throw new Error('impossible');

      // TODO: type safety
      const inner: BindingsType<any> = this.inner;
      const acc: HashInput = { fields: [], packed: [] };
      x.forEach((el) => {
        const { fields, packed } = inner.toInput(el);
        acc.fields!.push(...(fields ?? []));
        acc.packed!.push(...(packed ?? []));
      });
      return acc;
    }

    toFields(x: T): BindingsLeaves.Field[] {
      if (!(x instanceof JsArray)) throw new Error('impossible');

      // TODO: type safety
      const inner: BindingsType<any> = this.inner;
      return x.map((el) => inner.toFields(el)).flat();
    }

    toAuxiliary(x?: T): any[] {
      if (this.staticLength !== null) {
        if (x !== undefined) {
          // TODO: type safety
          const x2 = x as any[];
          if (x2.length !== this.staticLength) throw new Error('invalid array length');
          return x2.map((v) => this.inner.toAuxiliary(v));
        } else {
          return new JsArray(this.staticLength).fill(this.inner.toAuxiliary());
        }
      } else {
        // TODO: type safety
        return x as any[];
      }
    }

    fromFields(fields: BindingsLeaves.Field[], aux: any[]): T {
      if (this.staticLength !== null) {
        const decoder = new FieldsDecoder(fields);
        const x = new JsArray();
        for (let i = 0; i < this.staticLength; i++)
          x[i] = decoder.decode(this.inner.sizeInFields(), (f) => this.inner.fromFields(f, aux[i]));
        // TODO: type safety
        return x as T;
      } else {
        // TODO: type safety
        return aux as T;
      }
    }

    toValue(x: T): T {
      return x;
    }

    fromValue(x: T): T {
      return x;
    }

    check(_x: T) {
      throw new Error('TODO');
    }
  }

  export type Option<T> = Option.OrUndefined<T> | Option.Flagged<T> | Option.ClosedInterval<T>;

  export namespace Option {
    export class OrUndefined<T> implements Provable<T> {
      readonly _T!: T extends infer _U | undefined ? void : never;

      constructor(public readonly inner: T extends infer U | undefined ? BindingsType<U> : never) {}

      sizeInFields(): number {
        return 0;
      }

      toJSON(x: T): any {
        // TODO: type safety
        const x2 = x as any | undefined;
        const inner = this.inner as BindingsType<any>;
        return x2 !== undefined ? inner.toJSON(x2) : null;
      }

      toInput(_x: T): any {
        return {};
      }

      toFields(_x: T): BindingsLeaves.Field[] {
        return [];
      }

      toAuxiliary(x?: T): any[] {
        return x === undefined ? [false] : [true, this.inner.toAuxiliary(x)];
      }

      fromFields(fields: BindingsLeaves.Field[], aux: any[]): T {
        // TODO: type safety
        return (aux[0] ? this.inner.fromFields(fields, aux[1]) : undefined) as T;
      }

      toValue(x: T): T {
        return x;
      }

      fromValue(x: T): T {
        return x;
      }

      check(_x: T) {
        throw new Error('TODO');
      }
    }

    export class Flagged<T> extends ProvableBindingsType<T, BindingsLeaves.Option<any>> {
      readonly _T!: T extends BindingsLeaves.Option<any> ? void : never;

      constructor(
        public readonly inner: T extends BindingsLeaves.Option<infer U> ? BindingsType<U> : never
      ) {
        super();
      }

      Type() {
        return BindingsLeaves.Option(this.inner as ProvableSerializable<any>);
      }
    }

    export class ClosedInterval<T> extends ProvableBindingsType<
      T,
      BindingsLeaves.Option<BindingsLeaves.Range<any>>
    > {
      readonly _T!: T extends BindingsLeaves.Option<BindingsLeaves.Range<any>> ? void : never;

      constructor(
        public readonly inner: T extends BindingsLeaves.Option<BindingsLeaves.Range<infer U>>
          ? BindingsType<U>
          : never
      ) {
        super();
      }

      Type() {
        return BindingsLeaves.Option(BindingsLeaves.Range(this.inner as ProvableSerializable<any>));
      }
    }
  }

  export type Leaf<T> =
    | Leaf.Number<T>
    | Leaf.String<T>
    | Leaf.Actions<T>
    | Leaf.AuthRequired<T>
    | Leaf.Bool<T>
    | Leaf.Events<T>
    | Leaf.Field<T>
    | Leaf.Int64<T>
    | Leaf.PublicKey<T>
    | Leaf.Sign<T>
    | Leaf.StateHash<T>
    | Leaf.TokenId<T>
    | Leaf.TokenSymbol<T>
    | Leaf.UInt32<T>
    | Leaf.UInt64<T>
    | Leaf.ZkappUri<T>;

  export namespace Leaf {
    abstract class AuxiliaryLeaf<T> {
      constructor() {}

      sizeInFields(): number {
        return 0;
      }

      toJSON(x: T): any {
        return x;
      }

      toInput(_x: T): HashInput {
        return {};
      }

      toFields(_x: T): BindingsLeaves.Field[] {
        return [];
      }

      toAuxiliary(x?: T): any[] {
        return [x];
      }

      fromFields(_fields: BindingsLeaves.Field[], aux: any[]): T {
        return aux[0];
      }

      toValue(x: T): T {
        return x;
      }

      fromValue(x: T): T {
        return x;
      }

      check(_x: T) {
        throw new Error('TODO');
      }
    }

    export class Number<T = number> extends AuxiliaryLeaf<T> {
      readonly _T!: T extends number ? void : never;
      readonly type: 'number' = 'number';
    }

    export class String<T = string> extends AuxiliaryLeaf<T> {
      readonly _T!: T extends string ? void : never;
      readonly type: 'string' = 'string';
    }

    export class Actions<T = BindingsLeaves.Actions> extends ProvableBindingsType<
      T,
      BindingsLeaves.Actions
    > {
      readonly _T!: T extends number ? void : never;
      readonly type: 'number' = 'number';

      Type() {
        return BindingsLeaves.Actions;
      }
    }

    export class AuthRequired<T = BindingsLeaves.AuthRequired> extends ProvableBindingsType<
      T,
      BindingsLeaves.AuthRequired
    > {
      readonly _T!: T extends BindingsLeaves.AuthRequired ? void : never;
      readonly type: 'AuthRequired' = 'AuthRequired';

      Type() {
        return BindingsLeaves.AuthRequired;
      }
    }

    export class Bool<T = BindingsLeaves.Bool> extends ProvableBindingsType<
      T,
      BindingsLeaves.Bool
    > {
      readonly _T!: T extends BindingsLeaves.Bool ? void : never;
      readonly type: 'Bool' = 'Bool';

      Type() {
        return BindingsLeaves.Bool;
      }
    }

    export class Events<T = BindingsLeaves.Events> extends ProvableBindingsType<
      T,
      BindingsLeaves.Events
    > {
      readonly _T!: T extends number ? void : never;
      readonly type: 'number' = 'number';

      Type() {
        return BindingsLeaves.Events;
      }
    }

    export class Field<T = BindingsLeaves.Field> extends ProvableBindingsType<
      T,
      BindingsLeaves.Field
    > {
      readonly _T!: T extends BindingsLeaves.Field ? void : never;
      readonly type: 'Field' = 'Field';

      Type() {
        return BindingsLeaves.Field;
      }
    }

    export class Int64<T = BindingsLeaves.Int64> extends ProvableBindingsType<
      T,
      BindingsLeaves.Int64
    > {
      readonly _T!: T extends BindingsLeaves.Int64 ? void : never;
      readonly type: 'Int64' = 'Int64';

      Type() {
        return BindingsLeaves.Int64;
      }
    }

    export class PublicKey<T = BindingsLeaves.PublicKey> extends ProvableBindingsType<
      T,
      BindingsLeaves.PublicKey
    > {
      readonly _T!: T extends BindingsLeaves.PublicKey ? void : never;
      readonly type: 'PublicKey' = 'PublicKey';

      Type() {
        return BindingsLeaves.PublicKey;
      }
    }

    export class Sign<T = BindingsLeaves.Sign> extends ProvableBindingsType<
      T,
      BindingsLeaves.Sign
    > {
      readonly _T!: T extends BindingsLeaves.Sign ? void : never;
      readonly type: 'Sign' = 'Sign';

      Type() {
        return BindingsLeaves.Sign;
      }
    }

    export class StateHash<T = BindingsLeaves.StateHash> extends ProvableBindingsType<
      T,
      BindingsLeaves.StateHash
    > {
      readonly _T!: T extends BindingsLeaves.StateHash ? void : never;
      readonly type: 'StateHash' = 'StateHash';

      Type() {
        return BindingsLeaves.StateHash;
      }
    }

    // TODO NOW
    export class TokenId<T = BindingsLeaves.TokenId> implements Provable<T> {
      readonly _T!: T extends BindingsLeaves.TokenId ? void : never;
      readonly type: 'TokenId' = 'TokenId';

      constructor() {}

      sizeInFields(): number {
        return BindingsLeaves.Field.sizeInFields();
      }

      toJSON(x: T): any {
        // TODO: type safety
        return toBase58Check(
          BindingsLeaves.Field.toBytes(x as BindingsLeaves.Field),
          versionBytes.tokenIdKey
        );
      }

      toInput(x: T): HashInput {
        // TODO: type safety
        return BindingsLeaves.Field.toInput(x as BindingsLeaves.Field);
      }

      toFields(x: T): BindingsLeaves.Field[] {
        // TODO: type safety
        return BindingsLeaves.Field.toFields(x as BindingsLeaves.Field);
      }

      toAuxiliary(_x?: T): any[] {
        return [];
      }

      fromFields(fields: BindingsLeaves.Field[], _aux: any[]): T {
        // TODO: type safety
        return BindingsLeaves.Field.fromFields(fields) as T;
      }

      toValue(x: T): T {
        return x;
      }

      fromValue(x: T): T {
        return x;
      }

      check(_x: T) {
        throw new Error('TODO');
      }
    }

    export class TokenSymbol<T = BindingsLeaves.TokenSymbol> extends ProvableBindingsType<
      T,
      BindingsLeaves.TokenSymbol
    > {
      readonly _T!: T extends BindingsLeaves.TokenId ? void : never;
      readonly type: 'TokenId' = 'TokenId';

      Type() {
        return BindingsLeaves.TokenSymbol;
      }
    }

    export class UInt32<T = BindingsLeaves.UInt32> extends ProvableBindingsType<
      T,
      BindingsLeaves.UInt32
    > {
      readonly _T!: T extends BindingsLeaves.UInt32 ? void : never;
      readonly type: 'UInt32' = 'UInt32';

      Type() {
        return BindingsLeaves.UInt32;
      }
    }

    export class UInt64<T = BindingsLeaves.UInt64> extends ProvableBindingsType<
      T,
      BindingsLeaves.UInt64
    > {
      readonly _T!: T extends BindingsLeaves.UInt64 ? void : never;
      readonly type: 'UInt64' = 'UInt64';

      Type() {
        return BindingsLeaves.UInt64;
      }
    }

    export class ZkappUri<T = BindingsLeaves.ZkappUri> extends ProvableBindingsType<
      T,
      BindingsLeaves.ZkappUri
    > {
      readonly _T!: T extends BindingsLeaves.ZkappUri ? void : never;
      readonly type: 'ZkappUri';

      Type() {
        return BindingsLeaves.ZkappUri;
      }
    }
  }
}