import { PrimitiveType, schema, SchemaType } from "./annotations.js";
import { TypeContext } from "./types/TypeContext.js";
import { Metadata } from "./Metadata.js";
import { Iterator } from "./encoding/decode.js";
import { Encoder } from "./encoder/Encoder.js";
import { Decoder } from "./decoder/Decoder.js";
import { Schema } from "./Schema.js";
import { t, FieldBuilder } from "./types/builder.js";
import { ArraySchema } from "./types/custom/ArraySchema.js";
import { $encodeDescriptor, $numFields } from "./types/symbols.js";
import { isQuantizedType, resolveQuantize } from "./types/quantize.js";

/**
 * Static methods available on Reflection
 */
interface ReflectionStatic {
    /**
     * Encodes the TypeContext of an Encoder into a buffer.
     *
     * @param encoder Encoder instance
     * @param it
     * @returns
     */
    encode: (encoder: Encoder, it?: Iterator) => Uint8Array;

    /**
     * Decodes the TypeContext from a buffer into a Decoder instance.
     *
     * @param bytes Reflection.encode() output
     * @param it
     * @returns Decoder instance
     */
    decode: <T extends Schema = Schema>(bytes: Uint8Array, it?: Iterator) => Decoder<T>;

    /**
     * Upgrade a class produced by `Reflection.decode` so its instances
     * can be used as encode sources (for `InputEncoder` or `Encoder`).
     *
     * `Reflection.decode` reconstructs classes with decoder-only field
     * slots — `inst.x = 7` lands as a direct own property and bypasses
     * the change-tracking + `$values` plumbing that encoders rely on.
     * Calling `makeEncodable(ctor)` installs the same prototype accessor
     * descriptors and `metadata[$encoders]` lookup table that the
     * `schema(...)` / `@type` builders install at class-definition time.
     *
     * Idempotent. Pay-as-you-go: callers that only decode never invoke
     * this and pay nothing extra. Must be called BEFORE any instance of
     * the class is constructed and assigned to.
     */
    makeEncodable: (ctor: typeof Schema) => typeof Schema;
}

/**
 * Reflection
 */

/**
 * `t.quantized()` field descriptor as it rides the reflection handshake —
 * schema-typed (bit-exact float64 bounds), NOT a string grammar, so every
 * language port decodes it with the schema decoder it already has.
 */
export const QuantizedDescriptor = schema({
    min: t.float64(),
    max: t.float64(),
    bits: t.uint8(),
    mode: t.uint8(), // 0 = clamp, 1 = wrap
}, "QuantizedDescriptor");
export type QuantizedDescriptor = SchemaType<typeof QuantizedDescriptor>;

export const ReflectionField = schema({
    name: t.string(),
    type: t.string(),
    referencedType: t.number(),
    /** Primitive child of a collection (`array`/`map`/... of "string" etc.) —
     *  its own slot, replacing the legacy `"array:string"` colon packing. */
    childPrimitive: t.string(),
    /** Set only on `t.quantized()` fields (`.optional()` — no auto-instantiated
     *  default; its absence is the "not quantized" signal on decode). */
    quantized: t.ref(QuantizedDescriptor).optional(),
}, "ReflectionField");
export type ReflectionField = SchemaType<typeof ReflectionField>;

export const ReflectionType = schema({
    id: t.number(),
    extendsId: t.number(),
    fields: t.array(ReflectionField),
}, "ReflectionType");
export type ReflectionType = SchemaType<typeof ReflectionType>;

export const Reflection = schema({
    types: t.array(ReflectionType),
    rootType: t.number(),
}, "Reflection") as ReturnType<typeof schema<{
    types: FieldBuilder<ArraySchema<ReflectionType>, true, false>;
    rootType: FieldBuilder<number, false, false>;
}>> & ReflectionStatic;

export type Reflection = SchemaType<typeof Reflection>;

Reflection.encode = function (encoder: Encoder, it: Iterator = { offset: 0 }) {
    const context = encoder.context;

    const reflection = new Reflection();
    const reflectionEncoder = new Encoder(reflection);

    // rootType is usually the first schema passed to the Encoder
    // (unless it inherits from another schema)
    const rootType = context.schemas.get(encoder.state.constructor);
    if (rootType > 0) { reflection.rootType = rootType; }

    const includedTypeIds = new Set<number>();
    const pendingReflectionTypes: { [typeid: number]: ReflectionType[] } = {};

    // add type to reflection in a way that respects inheritance
    // (parent types should be added before their children)
    const addType = (type: ReflectionType) => {
        if (type.extendsId === undefined || includedTypeIds.has(type.extendsId)) {
            includedTypeIds.add(type.id);

            reflection.types.push(type);

            const deps = pendingReflectionTypes[type.id];
            if (deps !== undefined) {
                delete pendingReflectionTypes[type.id];
                deps.forEach((childType) => addType(childType));
            }
        } else {
            if (pendingReflectionTypes[type.extendsId] === undefined) {
                pendingReflectionTypes[type.extendsId] = [];
            }
            pendingReflectionTypes[type.extendsId].push(type);
        }
    };

    context.schemas.forEach((typeid, klass) => {
        const type = new ReflectionType();
        type.id = Number(typeid);

        // support inheritance
        const inheritFrom = Object.getPrototypeOf(klass);
        if (inheritFrom !== Schema) {
            type.extendsId = context.schemas.get(inheritFrom);
        }

        const metadata = klass[Symbol.metadata];

        //
        // FIXME: this is a workaround for inherited types without additional fields
        // if metadata is the same reference as the parent class - it means the class has no own metadata
        //
        if (metadata !== inheritFrom[Symbol.metadata]) {
            // Walk by index rather than `for…in`: `@deprecated()` makes its
            // metadata slot non-enumerable, and dropping it from the payload
            // shifts every later field down one wire index on the peer.
            const numFields = (metadata[$numFields] ?? -1) as number;
            for (let index = 0; index <= numFields; index++) {
                const field = metadata[index];
                if (field === undefined) { continue; }

                const fieldName = field.name;

                // skip fields from parent classes
                if (!Object.prototype.hasOwnProperty.call(metadata, fieldName)) {
                    continue;
                }

                const reflectionField = new ReflectionField();
                reflectionField.name = fieldName;

                let fieldType: string;

                if (typeof (field.type) === "string") {
                    fieldType = field.type;

                } else if (isQuantizedType(field.type)) {
                    // Params ride as a schema-typed descriptor (bit-exact float64) —
                    // no string grammar for the peer (or a language port) to parse.
                    const d = field.type.quantized;
                    fieldType = "quantized";
                    const desc = new QuantizedDescriptor();
                    desc.min = d.min;
                    desc.max = d.max;
                    desc.bits = d.bits;
                    desc.mode = d.wrap ? 1 : 0;
                    reflectionField.quantized = desc;

                } else {
                    let childTypeSchema: typeof Schema;

                    //
                    // TODO: refactor below.
                    //
                    if (Schema.is(field.type)) {
                        fieldType = "ref";
                        childTypeSchema = field.type as typeof Schema;

                    } else {
                        fieldType = Object.keys(field.type)[0];

                        if (typeof (field.type[fieldType as keyof typeof field.type]) === "string") {
                            // primitive child gets its own slot (was packed as "array:string")
                            reflectionField.childPrimitive = field.type[fieldType as keyof typeof field.type] as string;

                        } else {
                            childTypeSchema = field.type[fieldType as keyof typeof field.type];
                        }
                    }

                    reflectionField.referencedType = (childTypeSchema)
                        ? context.getTypeId(childTypeSchema)
                        : -1;
                }

                reflectionField.type = fieldType;
                type.fields.push(reflectionField);
            }
        }

        addType(type);
    });

    // in case there are types that were not added due to inheritance
    for (const typeid in pendingReflectionTypes) {
        pendingReflectionTypes[typeid].forEach((type) =>
            reflection.types.push(type))
    }

    const buf = reflectionEncoder.encodeAll(it);
    return buf.slice(0, it.offset);
};

Reflection.decode = function <T extends Schema = Schema>(bytes: Uint8Array, it?: Iterator): Decoder<T> {
    const reflection = new Reflection();

    const reflectionDecoder = new Decoder(reflection);
    reflectionDecoder.decode(bytes, it);

    const typeContext = new TypeContext();

    // 1st pass, initialize metadata + inheritance
    reflection.types.forEach((reflectionType) => {
        const parentClass: typeof Schema = typeContext.get(reflectionType.extendsId) ?? Schema;
        const schema: typeof Schema = class _ extends parentClass { };

        // register for inheritance support
        TypeContext.register(schema);

        typeContext.add(schema, reflectionType.id);
    }, {});

    // define fields
    const addFields = (metadata: Metadata, reflectionType: ReflectionType, parentFieldIndex: number) => {
        reflectionType.fields.forEach((field, i) => {
            const fieldIndex = parentFieldIndex + i;

            if (field.quantized !== undefined) {
                // Schema-typed descriptor → resolved codec (validation stays in
                // resolveQuantize, same as the builder path).
                const q = field.quantized;
                Metadata.addField(metadata, fieldIndex, field.name, {
                    quantized: resolveQuantize({ min: q.min, max: q.max, bits: q.bits as 8 | 16 | 32, mode: q.mode === 1 ? "wrap" : "clamp" }),
                } as any);

            } else if (field.referencedType !== undefined) {
                const fieldType = field.type;
                // Schema child by type id; a primitive child (referencedType -1)
                // rides its own childPrimitive slot.
                const refType: PrimitiveType = typeContext.get(field.referencedType)
                    ?? field.childPrimitive as PrimitiveType;

                if (fieldType === "ref") {
                    Metadata.addField(metadata, fieldIndex, field.name, refType);

                } else {
                    Metadata.addField(metadata, fieldIndex, field.name, { [fieldType]: refType });
                }

            } else {
                Metadata.addField(metadata, fieldIndex, field.name, field.type as PrimitiveType);
            }
        });
    };

    // 2nd pass, set fields
    reflection.types.forEach((reflectionType) => {
        const schema = typeContext.get(reflectionType.id);

        // for inheritance support
        const metadata = Metadata.initialize(schema);

        const inheritedTypes: ReflectionType[] = [];

        let parentType: ReflectionType = reflectionType;
        do {
            inheritedTypes.push(parentType);
            parentType = reflection.types.find((t) => t.id === parentType.extendsId);
        } while (parentType);

        let parentFieldIndex = 0;

        inheritedTypes.reverse().forEach((reflectionType) => {
            // add fields from all inherited classes
            // TODO: refactor this to avoid adding fields from parent classes
            addFields(metadata, reflectionType, parentFieldIndex);
            parentFieldIndex += reflectionType.fields.length;
        });
    });

    const state: T = new (typeContext.get(reflection.rootType || 0) as unknown as any)();

    return new Decoder<T>(state, typeContext);
}

Reflection.makeEncodable = function (ctor: typeof Schema): typeof Schema {
    const metadata: any = (ctor as any)[Symbol.metadata];
    if (!metadata) return ctor;

    const numFields = metadata[$numFields];
    if (numFields === undefined) return ctor;

    // Walk every field index across the inheritance chain. Repeat calls
    // are cheap: defineField overwrites the same descriptor and re-stamps
    // the same `metadata[$encoders]` slot (idempotent).
    for (let i = 0; i <= numFields; i++) {
        const field = metadata[i];
        if (!field) continue;
        Metadata.defineField(ctor, metadata, i, field.name, field.type);
    }

    // Invalidate any cached encode descriptor — `getEncodeDescriptor`
    // memoizes on the constructor. If something already constructed it
    // (e.g. a prior `InputEncoder(...)` call that threw), drop the stale
    // entry so the next read sees the upgraded metadata.
    if (Object.prototype.hasOwnProperty.call(ctor, $encodeDescriptor)) {
        delete (ctor as any)[$encodeDescriptor];
    }

    return ctor;
};