import { SharedStruct } from './shared-struct.ts';
import { Bytes } from './bytes.ts';
import { SharedString } from './string.ts';
import { Tuple } from './tuple.ts';

const SHARED = Symbol.for('moroutine.shared');

const registry = new Map<string, new (buffer: SharedArrayBuffer, byteOffset: number) => unknown>();

export function registerSync(tag: string, ctor: new (buffer: SharedArrayBuffer, byteOffset: number) => unknown): void {
  registry.set(tag, ctor);
}

export function serializeArg(arg: unknown): unknown {
  if (typeof arg === 'object' && arg !== null && SHARED in arg) {
    const data = (arg as any)[SHARED]();
    if (data.tag === 'SharedStruct') {
      const serializedFields: Record<string, unknown> = {};
      for (const key in data.fields) {
        serializedFields[key] = serializeStructField(data.fields[key]);
      }
      return { __shared__: 'SharedStruct', fields: serializedFields };
    }
    if (data.tag === 'Tuple') {
      const serializedElements = (data.elements as any[]).map((el: any) => serializeStructField(el));
      return { __shared__: 'Tuple', elements: serializedElements };
    }
    return {
      __shared__: data.tag,
      buffer: data.buffer,
      byteOffset: data.byteOffset,
      ...(data.size !== undefined && { size: data.size }),
      ...(data.maxBytes !== undefined && { maxBytes: data.maxBytes }),
    };
  }
  return arg;
}

function serializeStructField(data: { tag: string; [key: string]: unknown }): unknown {
  if (data.tag === 'SharedStruct') {
    const serializedFields: Record<string, unknown> = {};
    for (const key in data.fields as Record<string, unknown>) {
      serializedFields[key] = serializeStructField((data.fields as any)[key]);
    }
    return { __shared__: 'SharedStruct', fields: serializedFields };
  }
  return {
    __shared__: data.tag,
    buffer: data.buffer,
    byteOffset: data.byteOffset,
    ...(data.size !== undefined && { size: data.size }),
    ...(data.maxBytes !== undefined && { maxBytes: data.maxBytes }),
  };
}

export function deserializeArg(arg: unknown): unknown {
  if (typeof arg === 'object' && arg !== null && '__shared__' in arg) {
    const data = arg as { __shared__: string; [key: string]: unknown };
    if (data.__shared__ === 'Bytes') {
      const typedData = data as { __shared__: string; buffer: SharedArrayBuffer; byteOffset: number; size: number };
      return new Bytes(typedData.size, typedData.buffer, typedData.byteOffset);
    }
    if (data.__shared__ === 'SharedString') {
      const typedData = data as { __shared__: string; buffer: SharedArrayBuffer; byteOffset: number; maxBytes: number };
      return new SharedString(typedData.maxBytes, typedData.buffer, typedData.byteOffset);
    }
    if (data.__shared__ === 'SharedStruct') {
      const fields = data.fields as Record<string, unknown>;
      const reconstructed: Record<string, unknown> = {};
      for (const key in fields) {
        reconstructed[key] = deserializeArg(fields[key]);
      }
      return new SharedStruct(reconstructed as any);
    }
    if (data.__shared__ === 'Tuple') {
      const elements = (data.elements as unknown[]).map((el) => deserializeArg(el));
      return new Tuple(elements as any);
    }
    const typedData = data as { __shared__: string; buffer: SharedArrayBuffer; byteOffset: number };
    const ctor = registry.get(typedData.__shared__);
    if (ctor) return new ctor(typedData.buffer, typedData.byteOffset);
  }
  return arg;
}