import type { Loadable } from './loadable.ts';
import type { Descriptor, BytesDescriptor, StringDescriptor } from './descriptors.ts';
import { int32, int64, bool } from './descriptors.ts';
import { Int32 } from './int32.ts';
import { Int64 } from './int64.ts';
import { Bool } from './bool.ts';
import { Bytes } from './bytes.ts';
import { SharedString } from './string.ts';
import { SharedStruct } from './shared-struct.ts';
import { Tuple } from './tuple.ts';

function isDescriptor(value: unknown): value is Descriptor<unknown> {
  return typeof value === 'function' && 'byteSize' in value && '_class' in value;
}

function isBytesDescriptor(value: unknown): value is BytesDescriptor {
  return typeof value === 'object' && value !== null && '_size' in value && '_class' in value;
}

function isStringDescriptor(value: unknown): value is StringDescriptor {
  return typeof value === 'object' && value !== null && '_maxBytes' in value && '_class' in value;
}

function isStructSchema(value: unknown): value is Record<string, unknown> {
  return typeof value === 'object' && value !== null && !Array.isArray(value);
}

function align(offset: number, alignment: number): number {
  const remainder = offset % alignment;
  return remainder === 0 ? offset : offset + (alignment - remainder);
}

function resolveValue(value: unknown): { descriptor: Descriptor<unknown>; initialValue: unknown } | null {
  if (typeof value === 'number') {
    if (!Number.isInteger(value) || value > 2_147_483_647 || value < -2_147_483_648) {
      throw new RangeError(`Value ${value} out of range for int32. Use int64 with bigint instead.`);
    }
    return { descriptor: int32 as unknown as Descriptor<unknown>, initialValue: value };
  }
  if (typeof value === 'bigint') {
    return { descriptor: int64 as unknown as Descriptor<unknown>, initialValue: value };
  }
  if (typeof value === 'boolean') {
    return { descriptor: bool as unknown as Descriptor<unknown>, initialValue: value };
  }
  return null;
}

interface LeafEntry {
  path: string[];
  descriptor: Descriptor<unknown>;
  offset: number;
  initialValue?: unknown;
  bytesSize?: number;
  stringMaxBytes?: number;
}

interface ArrayLeafEntry {
  type: 'leaf';
  descriptor: Descriptor<unknown>;
  offset: number;
  initialValue?: unknown;
  bytesSize?: number;
  stringMaxBytes?: number;
}

interface ArrayStructEntry {
  type: 'struct';
  schema: Record<string, unknown>;
  leaves: LeafEntry[];
}

interface ArrayTupleEntry {
  type: 'tuple';
  entries: ArrayEntry[];
}

type ArrayEntry = ArrayLeafEntry | ArrayStructEntry | ArrayTupleEntry;

function collectLeaves(
  schema: Record<string, unknown>,
  path: string[],
  leaves: LeafEntry[],
  cursor: { offset: number },
): void {
  for (const key in schema) {
    const value = schema[key];
    if (isStringDescriptor(value)) {
      cursor.offset = align(cursor.offset, value.byteAlignment);
      leaves.push({
        path: [...path, key],
        descriptor: value as unknown as Descriptor<unknown>,
        offset: cursor.offset,
        stringMaxBytes: value._maxBytes,
      });
      cursor.offset += value.byteSize;
    } else if (isBytesDescriptor(value)) {
      cursor.offset = align(cursor.offset, value.byteAlignment);
      leaves.push({
        path: [...path, key],
        descriptor: value as unknown as Descriptor<unknown>,
        offset: cursor.offset,
        bytesSize: value._size,
      });
      cursor.offset += value.byteSize;
    } else if (isDescriptor(value)) {
      cursor.offset = align(cursor.offset, value.byteAlignment);
      leaves.push({ path: [...path, key], descriptor: value, offset: cursor.offset });
      cursor.offset += value.byteSize;
    } else if (Array.isArray(value)) {
      collectArrayLeaves(value, leaves, cursor);
    } else if (isStructSchema(value)) {
      collectLeaves(value as Record<string, unknown>, [...path, key], leaves, cursor);
    } else {
      const resolved = resolveValue(value);
      if (resolved !== null) {
        const { descriptor, initialValue } = resolved;
        cursor.offset = align(cursor.offset, descriptor.byteAlignment);
        leaves.push({ path: [...path, key], descriptor, offset: cursor.offset, initialValue });
        cursor.offset += descriptor.byteSize;
      }
    }
  }
}

function collectArrayLeaves(schema: unknown[], leaves: LeafEntry[], cursor: { offset: number }): void {
  for (const element of schema) {
    if (isStringDescriptor(element)) {
      cursor.offset = align(cursor.offset, element.byteAlignment);
      leaves.push({
        path: [],
        descriptor: element as unknown as Descriptor<unknown>,
        offset: cursor.offset,
        stringMaxBytes: element._maxBytes,
      });
      cursor.offset += element.byteSize;
    } else if (isBytesDescriptor(element)) {
      cursor.offset = align(cursor.offset, element.byteAlignment);
      leaves.push({
        path: [],
        descriptor: element as unknown as Descriptor<unknown>,
        offset: cursor.offset,
        bytesSize: element._size,
      });
      cursor.offset += element.byteSize;
    } else if (isDescriptor(element)) {
      cursor.offset = align(cursor.offset, element.byteAlignment);
      leaves.push({ path: [], descriptor: element, offset: cursor.offset });
      cursor.offset += element.byteSize;
    } else if (Array.isArray(element)) {
      collectArrayLeaves(element, leaves, cursor);
    } else if (isStructSchema(element)) {
      collectLeaves(element as Record<string, unknown>, [], leaves, cursor);
    } else {
      const resolved = resolveValue(element);
      if (resolved !== null) {
        const { descriptor, initialValue } = resolved;
        cursor.offset = align(cursor.offset, descriptor.byteAlignment);
        leaves.push({ path: [], descriptor, offset: cursor.offset, initialValue });
        cursor.offset += descriptor.byteSize;
      }
    }
  }
}

function processArraySchema(schema: unknown[], cursor: { offset: number }): { entries: ArrayEntry[] } {
  const entries: ArrayEntry[] = [];
  for (const element of schema) {
    if (isStringDescriptor(element)) {
      cursor.offset = align(cursor.offset, element.byteAlignment);
      entries.push({
        type: 'leaf',
        descriptor: element as unknown as Descriptor<unknown>,
        offset: cursor.offset,
        stringMaxBytes: element._maxBytes,
      });
      cursor.offset += element.byteSize;
    } else if (isBytesDescriptor(element)) {
      cursor.offset = align(cursor.offset, element.byteAlignment);
      entries.push({
        type: 'leaf',
        descriptor: element as unknown as Descriptor<unknown>,
        offset: cursor.offset,
        bytesSize: element._size,
      });
      cursor.offset += element.byteSize;
    } else if (isDescriptor(element)) {
      cursor.offset = align(cursor.offset, element.byteAlignment);
      entries.push({ type: 'leaf', descriptor: element, offset: cursor.offset });
      cursor.offset += element.byteSize;
    } else if (Array.isArray(element)) {
      const nested = processArraySchema(element, cursor);
      entries.push({ type: 'tuple', entries: nested.entries });
    } else if (isStructSchema(element)) {
      const leaves: LeafEntry[] = [];
      collectLeaves(element as Record<string, unknown>, [], leaves, cursor);
      entries.push({ type: 'struct', schema: element as Record<string, unknown>, leaves });
    } else {
      const resolved = resolveValue(element);
      if (resolved !== null) {
        const { descriptor, initialValue } = resolved;
        cursor.offset = align(cursor.offset, descriptor.byteAlignment);
        entries.push({ type: 'leaf', descriptor, offset: cursor.offset, initialValue });
        cursor.offset += descriptor.byteSize;
      }
    }
  }
  return { entries };
}

function buildTupleFromEntries(entries: ArrayEntry[], buffer: SharedArrayBuffer): Tuple<any> {
  const elements = entries.map((entry) => {
    if (entry.type === 'leaf') {
      if (entry.stringMaxBytes !== undefined) {
        return new SharedString(entry.stringMaxBytes, buffer, entry.offset);
      }
      if (entry.bytesSize !== undefined) {
        return new Bytes(entry.bytesSize, buffer, entry.offset);
      }
      const instance = new entry.descriptor._class(buffer, entry.offset);
      if ('initialValue' in entry) {
        (instance as any).store(entry.initialValue);
      }
      return instance;
    }
    if (entry.type === 'struct') {
      const leafIndex = { i: 0 };
      return buildStructTree(entry.schema, entry.leaves, buffer, leafIndex);
    }
    if (entry.type === 'tuple') {
      return buildTupleFromEntries(entry.entries, buffer);
    }
  });
  return new Tuple(elements as any);
}

function buildStructTree(
  schema: Record<string, unknown>,
  leaves: LeafEntry[],
  buffer: SharedArrayBuffer,
  leafIndex: { i: number },
): SharedStruct<any> {
  const fields: Record<string, any> = {};
  for (const key in schema) {
    const value = schema[key];
    if (isStringDescriptor(value)) {
      const leaf = leaves[leafIndex.i++];
      fields[key] = new SharedString(leaf.stringMaxBytes!, buffer, leaf.offset);
    } else if (isBytesDescriptor(value)) {
      const leaf = leaves[leafIndex.i++];
      fields[key] = new Bytes(leaf.bytesSize!, buffer, leaf.offset);
    } else if (isDescriptor(value)) {
      const leaf = leaves[leafIndex.i++];
      fields[key] = new leaf.descriptor._class(buffer, leaf.offset);
    } else if (Array.isArray(value)) {
      const arrayEntries = processArraySchemaFromLeaves(value, leaves, leafIndex);
      fields[key] = buildTupleFromEntries(arrayEntries, buffer);
    } else if (isStructSchema(value)) {
      fields[key] = buildStructTree(value as Record<string, unknown>, leaves, buffer, leafIndex);
    } else if (resolveValue(value) !== null) {
      const leaf = leaves[leafIndex.i++];
      const instance = new leaf.descriptor._class(buffer, leaf.offset) as any;
      if ('initialValue' in leaf) {
        instance.store(leaf.initialValue);
      }
      fields[key] = instance;
    }
  }
  return new SharedStruct(fields);
}

function processArraySchemaFromLeaves(schema: unknown[], leaves: LeafEntry[], leafIndex: { i: number }): ArrayEntry[] {
  const entries: ArrayEntry[] = [];
  for (const element of schema) {
    if (isStringDescriptor(element)) {
      const leaf = leaves[leafIndex.i++];
      entries.push({
        type: 'leaf',
        descriptor: leaf.descriptor,
        offset: leaf.offset,
        stringMaxBytes: leaf.stringMaxBytes,
      });
    } else if (isBytesDescriptor(element)) {
      const leaf = leaves[leafIndex.i++];
      entries.push({ type: 'leaf', descriptor: leaf.descriptor, offset: leaf.offset, bytesSize: leaf.bytesSize });
    } else if (isDescriptor(element)) {
      const leaf = leaves[leafIndex.i++];
      entries.push({ type: 'leaf', descriptor: leaf.descriptor, offset: leaf.offset });
    } else if (Array.isArray(element)) {
      const nested = processArraySchemaFromLeaves(element, leaves, leafIndex);
      entries.push({ type: 'tuple', entries: nested });
    } else if (isStructSchema(element)) {
      const startIndex = leafIndex.i;
      // Count leaves for this struct
      countStructLeaves(element as Record<string, unknown>, leaves, leafIndex);
      const endIndex = leafIndex.i;
      const subLeaves = leaves.slice(startIndex, endIndex);
      entries.push({
        type: 'struct',
        schema: element as Record<string, unknown>,
        leaves: subLeaves,
      });
    } else if (resolveValue(element) !== null) {
      const leaf = leaves[leafIndex.i++];
      entries.push({ type: 'leaf', descriptor: leaf.descriptor, offset: leaf.offset, initialValue: leaf.initialValue });
    }
  }
  return entries;
}

function countStructLeaves(schema: Record<string, unknown>, leaves: LeafEntry[], leafIndex: { i: number }): void {
  for (const key in schema) {
    const value = schema[key];
    if (isStringDescriptor(value) || isBytesDescriptor(value) || isDescriptor(value)) {
      leafIndex.i++;
    } else if (Array.isArray(value)) {
      countArrayLeaves(value, leaves, leafIndex);
    } else if (isStructSchema(value)) {
      countStructLeaves(value as Record<string, unknown>, leaves, leafIndex);
    } else if (resolveValue(value) !== null) {
      leafIndex.i++;
    }
  }
}

function countArrayLeaves(schema: unknown[], leaves: LeafEntry[], leafIndex: { i: number }): void {
  for (const element of schema) {
    if (isStringDescriptor(element) || isBytesDescriptor(element) || isDescriptor(element)) {
      leafIndex.i++;
    } else if (Array.isArray(element)) {
      countArrayLeaves(element, leaves, leafIndex);
    } else if (isStructSchema(element)) {
      countStructLeaves(element as Record<string, unknown>, leaves, leafIndex);
    } else if (resolveValue(element) !== null) {
      leafIndex.i++;
    }
  }
}

// Maps a schema value to its resolved instance type
type ResolveField<T> =
  T extends Descriptor<infer R>
    ? R
    : T extends BytesDescriptor
      ? Bytes
      : T extends StringDescriptor
        ? SharedString
        : T extends readonly unknown[]
          ? ResolveTuple<T>
          : T extends Record<string, unknown>
            ? ResolveStruct<T>
            : T extends number
              ? Int32
              : T extends bigint
                ? Int64
                : T extends boolean
                  ? Bool
                  : never;

type ResolveStruct<T extends Record<string, unknown>> = SharedStruct<{ [K in keyof T]: ResolveField<T[K]> }>;
type ResolveTuple<T extends readonly unknown[]> = Tuple<ResolveTupleElements<T>>;
type ResolveTupleElements<T extends readonly unknown[]> = { [K in keyof T]: ResolveField<T[K]> } & Loadable<any>[];

/**
 * Allocates shared memory from a schema. Accepts descriptors, plain objects (structs),
 * arrays (tuples), or primitive values (shorthand). Compound schemas pack all fields
 * into a single SharedArrayBuffer.
 * @param schema - A descriptor, plain object, array, or primitive value defining the shape.
 * @returns A {@link Loadable} instance (or struct/tuple/lock) backed by shared memory.
 */
export function shared<T extends Descriptor<any>>(schema: T): ReturnType<T>;
export function shared(schema: BytesDescriptor): Bytes;
export function shared(schema: StringDescriptor): SharedString;
export function shared(schema: number): Int32;
export function shared(schema: bigint): Int64;
export function shared(schema: boolean): Bool;
export function shared<T extends readonly unknown[]>(schema: [...T]): ResolveTuple<T>;
export function shared<T extends Record<string, unknown>>(schema: T): ResolveStruct<T>;
export function shared(schema: unknown): any {
  if (typeof schema === 'number') {
    if (!Number.isInteger(schema) || schema > 2_147_483_647 || schema < -2_147_483_648) {
      throw new RangeError(`Value ${schema} out of range for int32. Use int64 with bigint instead.`);
    }
    const instance = new Int32(new SharedArrayBuffer(Int32.byteSize), 0);
    instance.store(schema);
    return instance;
  }

  if (typeof schema === 'bigint') {
    const instance = new Int64(new SharedArrayBuffer(Int64.byteSize), 0);
    instance.store(schema);
    return instance;
  }

  if (typeof schema === 'boolean') {
    const instance = new Bool(new SharedArrayBuffer(Bool.byteSize), 0);
    instance.store(schema);
    return instance;
  }

  if (isStringDescriptor(schema)) {
    return schema; // already a standalone instance
  }

  if (isBytesDescriptor(schema)) {
    return schema; // already a standalone instance
  }

  if (isDescriptor(schema)) {
    return schema();
  }

  if (Array.isArray(schema)) {
    const cursor = { offset: 0 };
    const { entries } = processArraySchema(schema, cursor);
    const buffer = new SharedArrayBuffer(cursor.offset);
    return buildTupleFromEntries(entries, buffer);
  }

  if (isStructSchema(schema)) {
    const leaves: LeafEntry[] = [];
    const cursor = { offset: 0 };
    collectLeaves(schema as Record<string, unknown>, [], leaves, cursor);
    const buffer = new SharedArrayBuffer(cursor.offset);
    const leafIndex = { i: 0 };
    return buildStructTree(schema as Record<string, unknown>, leaves, buffer, leafIndex);
  }

  throw new Error('Invalid schema');
}