import type { StorableClass, StorableInstance } from "./storable-protocol.ts";
import type { SerializationContext } from "./serialization-context.ts";
import type {
  JsonWireValue,
  SerializedForm,
} from "./json-serialization-context.ts";
import { ExplicitTagStorable } from "./explicit-tag-storable.ts";
import {
  StorableError,
  StorableMap,
  StorableRegExp,
  StorableSet,
  StorableUint8Array,
} from "./storable-native-instances.ts";
import { TAGS } from "./type-tags.ts";

/**
 * JSON serialization context implementing the `/<Type>@<Version>` wire format
 * from the formal spec (Section 5). Manages a static type registry for the
 * types in scope and handles encoding/decoding of tagged values.
 * See Section 5.2 of the formal spec.
 */
export class JsonEncodingContext
  implements SerializationContext<JsonWireValue> {
  /** Tag -> class registry for known types. */
  private readonly registry = new Map<
    string,
    StorableClass<StorableInstance>
  >();

  /** Whether failed reconstructions produce `ProblematicStorable` instead of
   *  throwing. */
  readonly lenient: boolean;

  constructor(options?: { lenient?: boolean }) {
    this.lenient = options?.lenient ?? false;

    // Register native wrapper classes for deserialization. Each wrapper's
    // static [RECONSTRUCT] method is used by the class registry fallback
    // path in deserialize(). This replaces the old ErrorHandler approach.
    this.registry.set(TAGS.Error, StorableError);
    this.registry.set(TAGS.Map, StorableMap);
    this.registry.set(TAGS.Set, StorableSet);
    // Note: TAGS.EpochNsec and TAGS.EpochDays are NOT registered here --
    // they have dedicated TypeHandlers (EpochNsecHandler, EpochDaysHandler)
    // that handle both serialization and deserialization directly.
    // Note: TAGS.BigInt is NOT registered here -- bigint is a primitive in
    // StorableDatum and is handled by a TypeHandler (like UndefinedHandler),
    // not a StorableInstance wrapper.
    this.registry.set(TAGS.Bytes, StorableUint8Array);
    this.registry.set(TAGS.RegExp, StorableRegExp);
  }

  /** Get the wire format tag for a storable instance's type. */
  getTagFor(value: StorableInstance): string {
    if (value instanceof ExplicitTagStorable) {
      return value.typeTag;
    }
    // Check for typeTag property (used by native-wrapping StorableInstance classes).
    const typeTag = (value as { typeTag?: unknown }).typeTag;
    if (typeof typeTag === "string") {
      return typeTag;
    }
    // Future rounds will add Cell/Stream/etc. here.
    throw new Error(
      `JsonEncodingContext: no tag registered for value: ${value}`,
    );
  }

  /** Get the class that can reconstruct instances for a given tag. */
  getClassFor(
    tag: string,
  ): StorableClass<StorableInstance> | undefined {
    return this.registry.get(tag);
  }

  /**
   * Encode a tag and state into the `/<tag>` wire format. Prepends `/` to the
   * tag to produce the JSON key. See Section 5.2 of the formal spec.
   */
  encode(tag: string, state: SerializedForm): SerializedForm {
    return { [`/${tag}`]: state } as SerializedForm;
  }

  /**
   * Decode a wire representation. Detects single-key objects with `/`-prefixed
   * keys. Returns `{ tag, state }` or `null` if not a tagged value.
   * See Section 5.4 of the formal spec.
   */
  decode(
    data: SerializedForm,
  ): { tag: string; state: SerializedForm } | null {
    if (
      data === null || typeof data !== "object" || Array.isArray(data)
    ) {
      return null;
    }

    const keys = Object.keys(data);
    if (keys.length !== 1) {
      return null;
    }

    const key = keys[0];
    if (!key.startsWith("/")) {
      return null;
    }

    const tag = key.slice(1);
    const state = (data as Record<string, SerializedForm>)[key];
    return { tag, state };
  }

  /** Convert a JsonWireValue tree to UTF-8-encoded JSON bytes. */
  finalize(data: SerializedForm): Uint8Array {
    return new TextEncoder().encode(JSON.stringify(data));
  }

  /** Parse UTF-8-encoded JSON bytes back into a JsonWireValue tree. */
  parse(bytes: Uint8Array): SerializedForm {
    return JSON.parse(new TextDecoder().decode(bytes)) as SerializedForm;
  }
}