import { Tag } from "./constants";
import { Bytes } from "./bytes";
import { Value } from "./values";
import { EncodeError } from "./codec";
import { Record, Tuple } from "./record";
import { GenericEmbedded, EmbeddedTypeEncode } from "./embedded";
import type { Embedded } from "./embedded";

export type Encodable<T> =
    Value<T> | Preservable<T> | Iterable<Value<T>> | ArrayBufferView;

export interface Preservable<T> {
    __preserve_on__(encoder: Encoder<T>): void;
}

export function isPreservable<T>(v: any): v is Preservable<T> {
    return typeof v === 'object' && v !== null && '__preserve_on__' in v && typeof v.__preserve_on__ === 'function';
}

export interface EncoderOptions {
    canonical?: boolean;
    includeAnnotations?: boolean;
}

export interface EncoderEmbeddedOptions<T> extends EncoderOptions {
    embeddedEncode?: EmbeddedTypeEncode<T>;
}

export function asLatin1(bs: Uint8Array): string {
    return String.fromCharCode.apply(null, bs as any as number[]);
}

function isIterable<T>(v: any): v is Iterable<T> {
    return typeof v === 'object' && v !== null && typeof v[Symbol.iterator] === 'function';
}

let _nextId = 0;
const _registry = new WeakMap<object, number>();
export function embeddedId(v: any): number {
    let id = _registry.get(v);
    if (id === void 0) {
        id = _nextId++;
        _registry.set(v, id);
    }
    return id;
}

export const identityEmbeddedTypeEncode: EmbeddedTypeEncode<any> = {
    encode(s: EncoderState, v: any): void {
        new Encoder(s, this).push(embeddedId(v));
    }
};

export class EncoderState {
    chunks: Array<Uint8Array>;
    view: DataView;
    index: number;
    options: EncoderOptions;

    constructor(options: EncoderOptions) {
        this.chunks = [];
        this.view = new DataView(new ArrayBuffer(256));
        this.index = 0;
        this.options = options;
    }

    get canonical(): boolean {
        return this.options.canonical ?? true;
    }

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

    contents(): Bytes {
        if (this.chunks.length === 0) {
            const resultLength = this.index;
            this.index = 0;
            return new Bytes(this.view.buffer.slice(0, resultLength));
        } else {
            this.rotatebuffer(4096);
            const chunks = this.chunks;
            this.chunks = [];
            return Bytes.concat(chunks);
        }
    }

    /* Like contents(), but hands back a string containing binary data "encoded" via latin-1 */
    contentsString(): string {
        if (this.chunks.length === 0) {
            const s = asLatin1(new Uint8Array(this.view.buffer, 0, this.index));
            this.index = 0;
            return s;
        } else {
            this.rotatebuffer(4096);
            const chunks = this.chunks;
            this.chunks = [];
            return chunks.map(asLatin1).join('');
        }
    }

    rotatebuffer(size: number) {
        this.chunks.push(new Uint8Array(this.view.buffer, 0, this.index));
        this.view = new DataView(new ArrayBuffer(size));
        this.index = 0;
    }

    makeroom(amount: number) {
        if (this.index + amount > this.view.byteLength) {
            this.rotatebuffer(amount + 4096);
        }
    }

    emitbyte(b: number) {
        this.makeroom(1);
        this.view.setUint8(this.index++, b);
    }

    emitbytes(bs: Uint8Array) {
        this.makeroom(bs.length);
        (new Uint8Array(this.view.buffer)).set(bs, this.index);
        this.index += bs.length;
    }

    varint(v: number) {
        while (v >= 128) {
            this.emitbyte((v % 128) + 128);
            v = Math.floor(v / 128);
        }
        this.emitbyte(v);
    }

    encodeint(v: number) {
        // TODO: Bignums :-/
        const plain_bitcount = Math.floor(Math.log2(v > 0 ? v : -(1 + v))) + 1;
        const signed_bitcount = plain_bitcount + 1;
        const bytecount = (signed_bitcount + 7) >> 3;
        if (bytecount <= 16) {
            this.emitbyte(Tag.MediumInteger_lo + bytecount - 1);
        } else {
            this.emitbyte(Tag.SignedInteger);
            this.varint(bytecount);
        }
        const enc = (n: number, x: number) => {
            if (n > 0) {
                enc(n - 1, Math.floor(x / 256));
                this.emitbyte(x & 255);
            }
        };
        enc(bytecount, v);
    }

    encodebytes(tag: Tag, bs: Uint8Array) {
        this.emitbyte(tag);
        this.varint(bs.length);
        this.emitbytes(bs);
    }
}

export class Encoder<T = object> {
    state: EncoderState;
    embeddedEncode: EmbeddedTypeEncode<T>;

    constructor(options: EncoderEmbeddedOptions<T>);
    constructor(state: EncoderState, embeddedEncode?: EmbeddedTypeEncode<T>);
    constructor(
        state_or_options: (EncoderState | EncoderEmbeddedOptions<T>) = {},
        embeddedEncode?: EmbeddedTypeEncode<T>)
    {
        if (state_or_options instanceof EncoderState) {
            this.state = state_or_options;
            this.embeddedEncode = embeddedEncode ?? identityEmbeddedTypeEncode;
        } else {
            this.state = new EncoderState(state_or_options);
            this.embeddedEncode = state_or_options.embeddedEncode ?? identityEmbeddedTypeEncode;
        }
    }

    withEmbeddedEncode<S>(
        embeddedEncode: EmbeddedTypeEncode<S>,
        body: (e: Encoder<S>) => void): this
    {
        body(new Encoder(this.state, embeddedEncode));
        return this;
    }

    get canonical(): boolean {
        return this.state.canonical;
    }

    get includeAnnotations(): boolean {
        return this.state.includeAnnotations;
    }

    contents(): Bytes {
        return this.state.contents();
    }

    contentsString(): string {
        return this.state.contentsString();
    }

    encodevalues(tag: Tag, items: Iterable<Value<T>>) {
        this.state.emitbyte(tag);
        for (let i of items) { this.push(i); }
        this.state.emitbyte(Tag.End);
    }

    push(v: Encodable<T>) {
        if (isPreservable<unknown>(v)) {
            v.__preserve_on__(this);
        }
        else if (isPreservable<T>(v)) {
            v.__preserve_on__(this);
        }
        else if (typeof v === 'boolean') {
            this.state.emitbyte(v ? Tag.True : Tag.False);
        }
        else if (typeof v === 'number') {
            if (v >= -3 && v <= 12) {
                this.state.emitbyte(Tag.SmallInteger_lo + ((v + 16) & 0xf));
            } else {
                this.state.encodeint(v);
            }
        }
        else if (typeof v === 'string') {
            this.state.encodebytes(Tag.String, new Bytes(v)._view);
        }
        else if (typeof v === 'symbol') {
            const key = Symbol.keyFor(v);
            if (key === void 0) throw new EncodeError("Cannot preserve non-global Symbol", v);
            this.state.encodebytes(Tag.Symbol, new Bytes(key)._view);
        }
        else if (ArrayBuffer.isView(v)) {
            if (v instanceof Uint8Array) {
                this.state.encodebytes(Tag.ByteString, v);
            } else {
                const bs = new Uint8Array(v.buffer, v.byteOffset, v.byteLength);
                this.state.encodebytes(Tag.ByteString, bs);
            }
        }
        else if (Record.isRecord<Value<T>, Tuple<Value<T>>, T>(v)) {
            this.state.emitbyte(Tag.Record);
            this.push(v.label);
            for (let i of v) { this.push(i); }
            this.state.emitbyte(Tag.End);
        }
        else if (isIterable<Value<T>>(v)) {
            this.encodevalues(Tag.Sequence, v);
        }
        else {
            ((v: Embedded<T>) => {
                this.state.emitbyte(Tag.Embedded);
                this.embeddedEncode.encode(this.state, v.embeddedValue);
            })(v);
        }
        return this; // for chaining
    }
}

export function encode<T>(
    v: Encodable<T>,
    options: EncoderEmbeddedOptions<T> = {}): Bytes
{
    return new Encoder(options).push(v).contents();
}

const _canonicalEncoder = new Encoder({ canonical: true });
let _usingCanonicalEncoder = false;

export function canonicalEncode(v: Encodable<never>, options?: EncoderEmbeddedOptions<never>): Bytes;
export function canonicalEncode(v: Encodable<any>, options?: EncoderEmbeddedOptions<any>): Bytes;
export function canonicalEncode(v: any, options?: EncoderEmbeddedOptions<any>): Bytes {
    if (options === void 0 && !_usingCanonicalEncoder) {
        _usingCanonicalEncoder = true;
        const bs = _canonicalEncoder.push(v).contents();
        _usingCanonicalEncoder = false;
        return bs;
    } else {
        return encode(v, { ... options, canonical: true });
    }
}

export function canonicalString(v: Encodable<any>): string {
    if (!_usingCanonicalEncoder) {
        _usingCanonicalEncoder = true;
        const s = _canonicalEncoder.push(v).contentsString();
        _usingCanonicalEncoder = false;
        return s;
    } else {
        return new Encoder({ canonical: true }).push(v).contentsString();
    }
}

export function encodeWithAnnotations<T>(v: Encodable<T>,
                                         options: EncoderEmbeddedOptions<T> = {}): Bytes {
    return encode(v, { ... options, includeAnnotations: true });
}
