// Copyright (c) 2022, 2024 blackshirt. All rights reserved.
// Use of this source code is governed by a MIT License
// that can be found in the LICENSE file.
module asn1

// vfmt off
// bit masking values for ASN.1 tag header
const tag_class_mask     = 0xc0 // 192, bits 8-7
const constructed_mask     = 0x20 //  32, bits 6
const tag_number_mask     = 0x1f //  31, bits 1-5
// vfmt on

// Maximum number of bytes to represent tag number, includes the tag byte.
const max_tag_length = 3
const max_tag_number = 16383

// Maximum value for UNIVERSAL class tag number, see `TagType`,
// Tag number above this number should be considered to other class,
// PRIVATE, CONTEXT_SPECIFIC or APPLICATION class.
const max_universal_tagnumber = 255

// ASN.1 Tag identifier handling.
//
// Tag represents identifier of the ASN.1 element (object).
// ASN.1 Tag number can be represented in two form, the short form and the long form.
// The short form for tag number below <= 30 and stored enough in single byte.
// The long form for tag number > 30, and stored in two or more bytes.
//
// ASN.1 imposes no limit on the tag number, but the NIST Stable Implementation Agreements (1991)
// and its European and Asian counterparts limit the size of tags to 16383.
// see https://www.oss.com/asn1/resources/asn1-faq.html#tag-limitation
// We impose limit on the tag number to be in range 0..16383.
// Its big enough to accomodate and represent different of yours own tag number.
// Its represents 2 bytes length where maximum bytes arrays to represent tag number
// in multibyte (long) form is `[u8(0x1f), 0xff, 0x7f]` or 16383 in base 128.
pub struct Tag {
mut:
    class       TagClass = .universal
    constructed bool
    number      int
}

// `Tag.new` creates new ASN.1 tag identifier. Its accepts params of TagClass `cls`,
// the tag form in the constructed or primitive form in `constructed` flag, and the integer tag `number`.
pub fn Tag.new(cls TagClass, constructed bool, number int) !Tag {
    if number < 0 || number > max_tag_number {
        return error('Unallowed tag number, ${number} exceed limit')
    }
    // validates required criteria
    if cls == .universal {
        // UNIVERSAL tag number should below max_universal_tagnumber
        if number > max_universal_tagnumber {
            return error('Not a valid tag number for universal class=${number}')
        }
        // SEQUENCE (OF) or SET (OF) should be in constructed form
        if number == int(TagType.sequence) || number == int(TagType.set) {
            if !constructed {
                return error('For SEQUENCE(OF) or SET(OF) type, should be in constructed form')
            }
        }
    }

    // otherwise, ok
    tag := Tag{
        class:       cls
        constructed: constructed
        number:      number
    }
    return tag
}

// `tag_class` return the ASN.1 class of this tag
pub fn (t Tag) tag_class() TagClass {
    return t.class
}

// `is_constructed` tells us whether this tag is in constructed form
// or the primitive one, ie, not constructed.
pub fn (t Tag) is_constructed() bool {
    return t.constructed
}

// `tag_number` return the tag nunber of this tag
pub fn (t Tag) tag_number() int {
    return t.number
}

// `encode` serializes the tag into bytes array with default rule, ie, .der.
// Its serializes into dst bytes buffer or return error on fails.
pub fn (t Tag) encode(mut dst []u8) ! {
    t.encode_with_rule(mut dst, .der)!
}

// `Tag.from_bytes` creates a new Tag from bytes. Its return newly created tag and
// the rest of remaining bytes on success, or error on failures.
pub fn Tag.from_bytes(bytes []u8) !(Tag, []u8) {
    tag, next_pos := Tag.decode(bytes)!
    if next_pos < bytes.len {
        rest := unsafe { bytes[next_pos..] }
        return tag, rest
    }
    if next_pos == bytes.len {
        return tag, []u8{}
    }
    return error('Tag: too short data')
}

// `equal` checks whether this tag is equal with the other tag
pub fn (t Tag) equal(o Tag) bool {
    return t.class == o.class && t.constructed == o.constructed && t.number == o.number
}

// encode_with_rule serializes tag into bytes array
fn (t Tag) encode_with_rule(mut dst []u8, rule EncodingRule) ! {
    // we currently only support .der or (stricter) .ber
    if rule != .der && rule != .ber {
        return error('Tag: unsupported rule')
    }
    // makes sure tag number is valid
    if t.number > max_tag_number {
        return error('Tag: tag number exceed limit')
    }

    // get the class type and constructed bit and build the bytes tag.
    // if the tag number > 0x1f, represented in long form required two or more bytes,
    // otherwise, represented in short form, fit in single byte.
    mut b := (u8(t.class) << 6) & tag_class_mask
    if t.constructed {
        b |= constructed_mask
    }
    // The tag is in long form
    if t.number >= 0x1f {
        b |= tag_number_mask // 0x1f
        dst << b
        t.to_bytes_in_base128(mut dst)!
    } else {
        // short form
        b |= u8(t.number)
        dst << b
    }
}

// Tag.decode tries to deserializes bytes into Tag.
// Its return error on fails.
fn Tag.decode(bytes []u8) !(Tag, int) {
    tag, next := Tag.decode_from_offset(bytes, 0)!
    return tag, next
}

// Tag.decode tries to deserializes bytes into Tag. its return error on fails.
fn Tag.decode_from_offset(bytes []u8, pos int) !(Tag, int) {
    // default rule
    tag, next := Tag.decode_with_rule(bytes, pos, .der)!
    return tag, next
}

// Tag.decode_with_rule deserializes bytes back into Tag structure start from `loc` offset.
// By default, its decodes in .der encoding rule, if you want more control, pass your `Params`.
// Its return Tag and next offset to operate on, and return error if it fails to decode.
fn Tag.decode_with_rule(bytes []u8, loc int, rule EncodingRule) !(Tag, int) {
    // preliminary check
    if bytes.len < 1 {
        return error('Tag: bytes underflow')
    }
    if rule != .der && rule != .ber {
        return error('Tag: unsupported rule')
    }
    // when accessing byte at ofset `loc` within bytes, ie, `b := bytes[loc]`,
    // its maybe can lead to panic when the loc is not be checked.
    if loc >= bytes.len {
        return error('Tag: invalid pos')
    }
    mut pos := loc
    // first byte of tag bytes
    b := bytes[pos]
    pos += 1

    // First we get the first byte from the bytes, check and gets the class and constructed bits
    // and the tag number marker. If this marker == 0x1f, it tells whether the tag number is represented
    // in multibyte (long form), or short form otherwise.
    class := int((b & tag_class_mask) >> 6)
    constructed := b & constructed_mask == constructed_mask
    mut number := int(b & tag_number_mask)

    // check if this `number` is in long (multibyte) form, and interpretes more bytes as a tag number.
    if number == 0x1f {
        // we only allowed `max_tag_length` bytes following to represent tag number.
        number, pos = Tag.read_tagnum(bytes, pos)!

        // pos is the next position to read next bytes, so check tag bytes length
        if pos >= max_tag_length + loc + 1 {
            return error('Tag: tag bytes is too long')
        }
        if number < 0x1f && rule == .der {
            // requirement for DER encoding.
            // TODO: the other encoding may remove this restriction
            return error('Tag: non-minimal tag')
        }
    }
    // build the tag
    tag := Tag.new(TagClass.from_int(class)!, constructed, number)!

    return tag, pos
}

fn (t Tag) str() string {
    cls := t.class.str()
    form := if t.constructed { 'true' } else { 'false' }
    number := t.number.str()
    return '${cls}-${form}-${number}'
}

// bytes_len tells amount of bytes needed to store tag in base 128
fn (t Tag) bytes_len() int {
    if t.number == 0 {
        return 1
    }
    mut n := t.number
    mut ret := 0

    for n > 0 {
        ret += 1
        n >>= 7
    }

    return ret
}

// `tag_size` informs us how many bytes needed to store this tag includes one byte marker if in the long form.
pub fn (t Tag) tag_size() int {
    // when number is greater than 31 (0x1f), its need more bytes
    // to represent this number, includes one byte marker for long form tag
    len := if t.number < 0x1f { 1 } else { t.bytes_len() + 1 }
    return len
}

// to_bytes_in_base128 serializes tag number into bytes in base 128
fn (t Tag) to_bytes_in_base128(mut dst []u8) ! {
    n := t.bytes_len()
    for i := n - 1; i >= 0; i-- {
        mut o := u8(t.number >> u32(i * 7))
        o &= 0x7f
        if i != 0 {
            o |= 0x80
        }
        dst << o
    }
}

// Tag.read_tagnum read the tag number from bytes from offset pos in base 128.
// Its return deserialized Tag number and next offset to process on.
fn Tag.read_tagnum(bytes []u8, pos int) !(int, int) {
    tnum, next := Tag.read_tagnum_with_rule(bytes, pos, .der)!
    return tnum, next
}

// read_tagnum_with_rule is the main routine to read the tag number part in the bytes source,
// start from offset loc in base 128. Its return the tag number and next offset to process on, or error on fails.
fn Tag.read_tagnum_with_rule(bytes []u8, loc int, rule EncodingRule) !(int, int) {
    if loc > bytes.len {
        return error('Tag number: invalid pos')
    }
    mut pos := loc
    mut ret := 0
    for s := 0; pos < bytes.len; s++ {
        ret <<= 7
        b := bytes[pos]

        if s == 0 && b == 0x80 {
            if rule == .der {
                // requirement for DER encoding
                return error('Tag number: integer is not minimally encoded')
            }
        }
        ret |= b & 0x7f
        pos += 1

        if b & 0x80 == 0 {
            if ret > max_tag_number {
                return error('Tag number: base 128 integer too large')
            }
            if ret < 0 {
                return error('Negative tag number')
            }

            return ret, pos
        }
    }
    return error('Tag: truncated base 128 integer')
}

fn (t Tag) valid_supported_universal_tagnum() bool {
    return t.class == .universal && t.number < max_universal_tagnumber
}

// `universal_tag_type` turns this Tag into available UNIVERSAL class of TagType,
// or return error if it is unknown number.
pub fn (t Tag) universal_tag_type() !TagType {
    // currently, only support Standard universal tag number
    if t.number > max_universal_tagnumber {
        return error('Tag number: unknown TagType number=${t.number}')
    }
    match t.class {
        .universal {
            match t.number {
                // vfmt off
                0 { return .reserved }
                1 { return .boolean }
                2 { return .integer }
                3 { return .bitstring }
                4 { return .octetstring }
                5 { return .null }
                6 { return .oid }
                7 { return .objdesc }
                8 { return .external }
                9 { return .real }
                10 { return .enumerated }
                11 { return .embedded }
                12 { return .utf8string }
                13 { return .relativeoid }
                14 { return .time }
                16 { return .sequence }
                17 { return .set }
                18 { return .numericstring }
                19 { return .printablestring }
                20 { return .t61string }
                21 { return .videotexstring }
                22 { return .ia5string }
                23 { return .utctime }
                24 { return .generalizedtime }
                25 { return .graphicstring }
                26 { return .visiblestring }
                27 { return .generalstring }
                28 { return .universalstring }
                29 { return .characterstring }
                30 { return .bmpstring }
                31 { return .date }
                32 { return .time_of_day }
                33 { return .date_time }
                34 { return .duration }
                35 { return .i18_oid }
                36 { return .relative_i18_oid }
                else {
                    return error('reserved or unknonw number')
                }
                // vfmt on
            }
        }
        else {
            return error('Not universal class type')
        }
    }
}

// TagClass is an enum of ASN.1 tag class.
//
// To make sure ASN.1 encodings are not ambiguous, every ASN.1 type is associated with a tag.
// A tag consists of three parts: the tag class, tag form and the tag number.
// The following classes are defined in the ASN.1 standard.
pub enum TagClass {
    universal        = 0x00 // 0b00
    application      = 0x01 // 0b01
    context_specific = 0x02 // 0b10
    private          = 0x03 // 0b11
}

// `TagClass.from_int` creates TagClass from integer v.
// Its return a new TagClass on success or error on fails.
fn TagClass.from_int(v int) !TagClass {
    match v {
        // vfmt off
        0x00 { return .universal }
        0x01 { return .application }
        0x02 { return .context_specific }
        0x03 { return .private }
        else {
            return error('Bad class number')
        }
        // vfmt on
    }
}

// `TagClass.from_string` creates a TagClass from string s.
// Its return a new TagClass on success or error on fails.
fn TagClass.from_string(s string) !TagClass {
    match s {
        // vfmt off
        'universal' { return .universal }
        'private' { return .private }
        'application' { return .application }
        'context_specific' { return .context_specific }
        else { 
            return error('bad class string') 
        }
        // vfmt on
    }
}

fn (c TagClass) str() string {
    match c {
        .universal { return 'universal' }
        .application { return 'application' }
        .context_specific { return 'context_specific' }
        .private { return 'private' }
    }
}

// TagType is standard UNIVERSAL tag number. Some of them was deprecated,
// so its not going to be supported on this module.
pub enum TagType {
    // vfmt off
    reserved             = 0     // reserved for BER
    boolean             = 1     // BOOLEAN type
    integer             = 2     // INTEGER type
    bitstring             = 3     // BIT STRING
    octetstring         = 4     // OCTET STRING
    null                 = 5     // NULL
    oid                 = 6        // OBJECT IDENTIFIER
    objdesc             = 7     // OBJECT DESCRIPTOR
    external             = 8     // INSTANCE OF, EXTERNAL
    real                 = 9     // REAL
    enumerated             = 10     // ENUMERATED
    embedded             = 11     // EMBEDDED PDV
    utf8string             = 12     // UTF8STRING
    relativeoid         = 13     // RELATIVE-OID
    // deprecated
    time                 = 14
    // 0x0f is reserved
    sequence             = 16     // SEQUENCE, SEQUENCE OF, Constructed
    set                 = 17     // SET, SET OF, Constructed
    numericstring         = 18     // NUMERICSTRING
    printablestring     = 19     // PRINTABLESTRING
    t61string             = 20     // TELETEXSTRING, T61STRING
    videotexstring         = 21     // VIDEOTEXSTRING
    ia5string             = 22     // IA5STRING
    utctime             = 23     // UTCTIME
    generalizedtime     = 24     // GENERALIZEDTIME
    graphicstring         = 25     // GRAPHICSTRING
    visiblestring         = 26     // VISIBLESTRING, ISO646STRING
    generalstring         = 27     // GENERALSTRING
    universalstring     = 28     // UNIVERSALSTRING
    characterstring     = 29     // CHARACTER STRING
    bmpstring           = 30     // BMPSTRING
    // unsupported 
    date                = 0x1f
    time_of_day         = 0x20
    date_time           = 0x21
    duration            = 0x22
    // Internationalized OID
    i18_oid             = 0x23
    // Internationalized Relative OID
    // Reserved 0x25 and above
    relative_i18_oid     = 0x24
    // vfmt on
}

fn (t TagType) str() string {
    match t {
        .boolean { return 'BOOLEAN' }
        .integer { return 'INTEGER' }
        .bitstring { return 'BITSTRING' }
        .octetstring { return 'OCTETSTRING' }
        .null { return 'NULL' }
        .oid { return 'OID' }
        .objdesc { return 'OBJECT_DESCRIPTOR' }
        .external { return 'EXTERNAL' }
        .real { return 'REAL' }
        .enumerated { return 'ENUMERATED' }
        .embedded { return 'EMBEDDED' }
        .utf8string { return 'UTF8STRING' }
        .relativeoid { return 'RELATIVEOID' }
        .time { return 'TIME' }
        .sequence { return 'SEQUENCE_OR_SEQUENCEOF' }
        .set { return 'SET_OR_SET_OF' }
        .numericstring { return 'NUMERICSTRING' }
        .printablestring { return 'PRINTABLESTRING' }
        .t61string { return 'T61STRING' }
        .videotexstring { return 'VIDEOTEXSTRING' }
        .ia5string { return 'IA5STRING' }
        .utctime { return 'UTCTIME' }
        .generalizedtime { return 'GENERALIZEDTIME' }
        .graphicstring { return 'GRAPHICSTRING' }
        .visiblestring { return 'VISIBLESTRING' }
        .generalstring { return 'GENERALSTRING' }
        .universalstring { return 'UNIVERSALSTRING' }
        .characterstring { return 'CHARACTERSTRING' }
        .bmpstring { return 'BMPSTRING' }
        else { return 'UNSUPPORTED_TAG_TYPE' }
    }
}

// universal_tag_from_int gets default Tag for universal class type from integer value v.
fn universal_tag_from_int(v int) !Tag {
    if v < 0 || v > max_universal_tagnumber {
        return error('Get unexpected tag number for universal type')
    }
    match v {
        // vfmt off
        1 { return default_boolean_tag }
        2 { return default_integer_tag }
        3 { return default_bitstring_tag }
        4 { return default_octetstring_tag }
        5 { return default_null_tag }
        6 { return default_oid_tag }
        10 { return default_enumerated_tag }
        12 { return default_utf8string_tag }
        16 { return default_sequence_tag }
        17 { return default_set_tag }
        18 { return default_numericstring_tag }
        19 { return default_printablestring_tag }
        22 { return default_ia5string_tag }
        23 { return default_utctime_tag }
        24 { return default_generalizedtime_tag }
        26 { return default_visiblestring_tag }
        else {
            // should in primitive form 
            return Tag{.universal, false, v}
        }
        // vfmt on
    }
}

// EncodingRule is standard of rule thats drives how some ASN.1
// element was encoded or deserialized.
pub enum EncodingRule {
    // Distinguished Encoding Rules (DER)
    der = 0
    // Basic Encoding Rules (BER)
    ber = 1
    // Octet Encoding Rules (OER)
    oer = 2
    // Packed Encoding Rules (PER)
    per = 3
    // XML Encoding Rules (XER)
    xer = 4
}

// EXPLICIT and IMPLICIT MODE
//
// TaggedMode is way of rule how some element is being wrapped (unwrapped).
// Explicit rule add new (outer) tag to the existing element,
// where implicit rule replaces the tag of existing element.
pub enum TaggedMode {
    explicit
    implicit
}

// `TaggedMode.from_string` creates TaggedMode from string s.
fn TaggedMode.from_string(s string) !TaggedMode {
    match s {
        'explicit' { return .explicit }
        'implicit' { return .implicit }
        // empty string marked as .explicit
        '' { return .explicit }
        else { return error('Bad string for tagged mode') }
    }
}

fn (m TaggedMode) str() string {
    match m {
        .explicit { return 'explicit' }
        .implicit { return 'implicit' }
    }
}