import { Annotated } from "./annotated";
import { DecodeError, ShortPacket } from "./codec";
import { Tag } from "./constants";
import { Set, Dictionary } from "./dictionary";
import { DoubleFloat, SingleFloat } from "./float";
import { Record } from "./record";
import { Bytes, BytesLike, underlying } from "./bytes";
import { Value } from "./values";
import { is } from "./is";
import { embed, GenericEmbedded, Embedded, EmbeddedTypeDecode } from "./embedded";

export interface DecoderOptions {
    includeAnnotations?: boolean;
}

export interface DecoderEmbeddedOptions<T> extends DecoderOptions {
    embeddedDecode?: EmbeddedTypeDecode<T>;
}

export interface TypedDecoder<T> {
    atEnd(): boolean;

    mark(): any;
    restoreMark(m: any): void;

    skip(): void;
    next(): Value<T>;
    withEmbeddedDecode<S, R>(
        embeddedDecode: EmbeddedTypeDecode<S>,
        body: (d: TypedDecoder<S>) => R): R;

    nextBoolean(): boolean | undefined;
    nextFloat(): SingleFloat | undefined;
    nextDouble(): DoubleFloat | undefined;
    nextEmbedded(): Embedded<T> | undefined;
    nextSignedInteger(): number | undefined;
    nextString(): string | undefined;
    nextByteString(): Bytes | undefined;
    nextSymbol(): symbol | undefined;

    openRecord(): boolean;
    openSequence(): boolean;
    openSet(): boolean;
    openDictionary(): boolean;

    closeCompound(): boolean;
}

export function asLiteral<T, E extends Exclude<Value<T>, Annotated<T>>>(
    actual: Value<T>,
    expected: E): E | undefined
{
    return is(actual, expected) ? expected : void 0;
}

export class DecoderState {
    packet: Uint8Array;
    index = 0;
    options: DecoderOptions;

    constructor(packet: BytesLike, options: DecoderOptions) {
        this.packet = underlying(packet);
        this.options = options;
    }

    get includeAnnotations(): boolean {
        return this.options.includeAnnotations ?? false;
    }

    write(data: BytesLike) {
        if (this.index === this.packet.length) {
            this.packet = underlying(data);
        } else {
            this.packet = Bytes.concat([this.packet.slice(this.index), data])._view;
        }
        this.index = 0;
    }

    atEnd(): boolean {
        return this.index >= this.packet.length;
    }

    mark(): number {
        return this.index;
    }

    restoreMark(m: number): void {
        this.index = m;
    }

    shortGuard<R>(body: () => R, short: () => R): R {
        if (this.atEnd()) return short();
        // ^ important somewhat-common case optimization - avoid the exception

        const start = this.mark();
        try {
            return body();
        } catch (e) {
            if (ShortPacket.isShortPacket(e)) {
                this.restoreMark(start);
                return short();
            }
            throw e;
        }
    }

    nextbyte(): number {
        if (this.atEnd()) throw new ShortPacket("Short packet");
        return this.packet[this.index++];
    }

    nextbytes(n: number): DataView {
        const start = this.index;
        this.index += n;
        if (this.index > this.packet.length) throw new ShortPacket("Short packet");
        // ^ NOTE: greater-than, not greater-than-or-equal-to - this makes atEnd() inappropriate
        return new DataView(this.packet.buffer, this.packet.byteOffset + start, n);
    }

    varint(): number {
        // TODO: Bignums :-/
        const v = this.nextbyte();
        if (v < 128) return v;
        return (this.varint() << 7) + (v - 128);
    }

    peekend(): boolean {
        return (this.nextbyte() === Tag.End) || (this.index--, false);
    }

    nextint(n: number): number {
        // TODO: Bignums :-/
        if (n === 0) return 0;
        let acc = this.nextbyte();
        if (acc & 0x80) acc -= 256;
        for (let i = 1; i < n; i++) acc = (acc * 256) + this.nextbyte();
        return acc;
    }

    nextSmallOrMediumInteger(tag: number): number | undefined {
        if (tag >= Tag.SmallInteger_lo && tag <= Tag.SmallInteger_lo + 15) {
            const v = tag - Tag.SmallInteger_lo;
            return v > 12 ? v - 16 : v;
        }
        if (tag >= Tag.MediumInteger_lo && tag <= Tag.MediumInteger_lo + 15) {
            const n = tag - Tag.MediumInteger_lo;
            return this.nextint(n + 1);
        }
        return void 0;
    }

    wrap<T>(v: Value<T>): Value<T> {
        return this.includeAnnotations ? new Annotated(v) : v;
    }

    unshiftAnnotation<T>(a: Value<T>, v: Annotated<T>): Annotated<T> {
        if (this.includeAnnotations) {
            v.annotations.unshift(a);
        }
        return v;
    }
}

export const neverEmbeddedTypeDecode: EmbeddedTypeDecode<never> = {
    decode(_s: DecoderState): never {
        throw new Error("Embeddeds not permitted at this point in Preserves document");
    },

    fromValue(_v: Value<GenericEmbedded>): never {
        throw new Error("Embeddeds not permitted at this point in Preserves document");
    },
};

export class Decoder<T = never> implements TypedDecoder<T> {
    state: DecoderState;
    embeddedDecode: EmbeddedTypeDecode<T>;

    constructor(state: DecoderState, embeddedDecode?: EmbeddedTypeDecode<T>);
    constructor(packet?: BytesLike, options?: DecoderEmbeddedOptions<T>);
    constructor(
        packet_or_state: (DecoderState | BytesLike) = new Uint8Array(0),
        options_or_embeddedDecode?: (DecoderEmbeddedOptions<T> | EmbeddedTypeDecode<T>))
    {
        if (packet_or_state instanceof DecoderState) {
            this.state = packet_or_state;
            this.embeddedDecode = (options_or_embeddedDecode as EmbeddedTypeDecode<T>) ?? neverEmbeddedTypeDecode;
        } else {
            const options = (options_or_embeddedDecode as DecoderEmbeddedOptions<T>) ?? {};
            this.state = new DecoderState(packet_or_state, options);
            this.embeddedDecode = options.embeddedDecode ?? neverEmbeddedTypeDecode;
        }
    }

    write(data: BytesLike) {
        this.state.write(data);
    }

    nextvalues(): Value<T>[] {
        const result = [];
        while (!this.state.peekend()) result.push(this.next());
        return result;
    }

    static dictionaryFromArray<T>(vs: Value<T>[]): Dictionary<T> {
        const d = new Dictionary<T>();
        if (vs.length % 2) throw new DecodeError("Missing dictionary value");
        for (let i = 0; i < vs.length; i += 2) {
            d.set(vs[i], vs[i+1]);
        }
        return d;
    }

    next(): Value<T> {
        const tag = this.state.nextbyte();
        switch (tag) {
            case Tag.False: return this.state.wrap<T>(false);
            case Tag.True: return this.state.wrap<T>(true);
            case Tag.Float: return this.state.wrap<T>(new SingleFloat(this.state.nextbytes(4).getFloat32(0, false)));
            case Tag.Double: return this.state.wrap<T>(new DoubleFloat(this.state.nextbytes(8).getFloat64(0, false)));
            case Tag.End: throw new DecodeError("Unexpected Compound end marker");
            case Tag.Annotation: {
                const a = this.next();
                const v = this.next() as Annotated<T>;
                return this.state.unshiftAnnotation(a, v);
            }
            case Tag.Embedded: return this.state.wrap<T>(embed(this.embeddedDecode.decode(this.state)));
            case Tag.SignedInteger: return this.state.wrap<T>(this.state.nextint(this.state.varint()));
            case Tag.String: return this.state.wrap<T>(Bytes.from(this.state.nextbytes(this.state.varint())).fromUtf8());
            case Tag.ByteString: return this.state.wrap<T>(Bytes.from(this.state.nextbytes(this.state.varint())));
            case Tag.Symbol: return this.state.wrap<T>(Symbol.for(Bytes.from(this.state.nextbytes(this.state.varint())).fromUtf8()));
            case Tag.Record: {
                const vs = this.nextvalues();
                if (vs.length === 0) throw new DecodeError("Too few elements in encoded record");
                return this.state.wrap<T>(Record(vs[0], vs.slice(1)));
            }
            case Tag.Sequence: return this.state.wrap<T>(this.nextvalues());
            case Tag.Set: return this.state.wrap<T>(new Set(this.nextvalues()));
            case Tag.Dictionary: return this.state.wrap<T>(Decoder.dictionaryFromArray(this.nextvalues()));
            default: {
                const v = this.state.nextSmallOrMediumInteger(tag);
                if (v === void 0) {
                    throw new DecodeError("Unsupported Preserves tag: " + tag);
                }
                return this.state.wrap<T>(v);
            }
        }
    }

    try_next(): Value<T> | undefined {
        return this.state.shortGuard(() => this.next(), () => void 0);
    }

    atEnd(): boolean {
        return this.state.atEnd();
    }

    mark(): any {
        return this.state.mark();
    }

    restoreMark(m: any): void {
        this.state.restoreMark(m);
    }

    skip(): void {
        // TODO: be more efficient
        this.next();
    }

    withEmbeddedDecode<S, R>(
        embeddedDecode: EmbeddedTypeDecode<S>,
        body: (d: TypedDecoder<S>) => R): R
    {
        return body(new Decoder(this.state, embeddedDecode));
    }

    skipAnnotations(): void {
        if (!this.state.atEnd() && this.state.packet[this.state.index] === Tag.Annotation) {
            this.state.index++;
            this.skip();
        }
    }

    nextBoolean(): boolean | undefined {
        this.skipAnnotations();
        switch (this.state.nextbyte()) {
            case Tag.False: return false;
            case Tag.True: return true;
            default: return void 0;
        }
    }

    nextFloat(): SingleFloat | undefined {
        this.skipAnnotations();
        switch (this.state.nextbyte()) {
            case Tag.Float: return new SingleFloat(this.state.nextbytes(4).getFloat32(0, false));
            default: return void 0;
        }
    }

    nextDouble(): DoubleFloat | undefined {
        this.skipAnnotations();
        switch (this.state.nextbyte()) {
            case Tag.Double: return new DoubleFloat(this.state.nextbytes(8).getFloat64(0, false));
            default: return void 0;
        }
    }

    nextEmbedded(): Embedded<T> | undefined {
        this.skipAnnotations();
        switch (this.state.nextbyte()) {
            case Tag.Embedded: return embed(this.embeddedDecode.decode(this.state));
            default: return void 0;
        }
    }

    nextSignedInteger(): number | undefined {
        this.skipAnnotations();
        const b = this.state.nextbyte();
        switch (b) {
            case Tag.SignedInteger: return this.state.nextint(this.state.varint());
            default: return this.state.nextSmallOrMediumInteger(b);
        }
    }

    nextString(): string | undefined {
        this.skipAnnotations();
        switch (this.state.nextbyte()) {
            case Tag.String: return Bytes.from(this.state.nextbytes(this.state.varint())).fromUtf8();
            default: return void 0;
        }
    }

    nextByteString(): Bytes | undefined {
        this.skipAnnotations();
        switch (this.state.nextbyte()) {
            case Tag.ByteString: return Bytes.from(this.state.nextbytes(this.state.varint()));
            default: return void 0;
        }
    }

    nextSymbol(): symbol | undefined {
        this.skipAnnotations();
        switch (this.state.nextbyte()) {
            case Tag.Symbol:
                return Symbol.for(Bytes.from(this.state.nextbytes(this.state.varint())).fromUtf8());
            default:
                return void 0;
        }
    }

    openRecord(): boolean {
        this.skipAnnotations();
        return (this.state.nextbyte() === Tag.Record) || (this.state.index--, false);
    }

    openSequence(): boolean {
        this.skipAnnotations();
        return (this.state.nextbyte() === Tag.Sequence) || (this.state.index--, false);
    }

    openSet(): boolean {
        this.skipAnnotations();
        return (this.state.nextbyte() === Tag.Set) || (this.state.index--, false);
    }

    openDictionary(): boolean {
        this.skipAnnotations();
        return (this.state.nextbyte() === Tag.Dictionary) || (this.state.index--, false);
    }

    closeCompound(): boolean {
        return this.state.peekend();
    }
}

export function decode<T>(bs: BytesLike, options: DecoderEmbeddedOptions<T> = {}): Value<T> {
    return new Decoder(bs, options).next();
}

export function decodeWithAnnotations<T>(bs: BytesLike,
                                         options: DecoderEmbeddedOptions<T> = {}): Annotated<T> {
    return decode(bs, { ... options, includeAnnotations: true }) as Annotated<T>;
}