use crate::core::constants::*;
use crate::core::prefix::{
    FLOAT_POSITIVE_NEGATIVE_MASK, PREFIX_BYTES, PREFIX_CONTAINER, PREFIX_FLOAT,
    PREFIX_NEGATIVE_INT, PREFIX_POSITIVE_INT, PREFIX_SIMPLE, PREFIX_STRING, PREFIX_TAG,
};
use crate::core::simple_value::SimpleValue;
use crate::ir::ast::{Array, Node, Object, Value};

pub struct Encoder {
    buf: Vec<u8>,
    offset: usize,
}

impl Encoder {
    pub fn new() -> Self {
        Self {
            buf: vec![0u8; DEFAULT_BUF_SIZE],
            offset: 0,
        }
    }

    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            buf: vec![0u8; capacity],
            offset: 0,
        }
    }

    fn ensure_capacity(&mut self, needed: usize) {
        while self.offset + needed > self.buf.len() {
            self.buf.resize(self.buf.len() * 2, 0);
        }
    }

    fn write_byte(&mut self, b: u8) {
        self.ensure_capacity(1);
        self.buf[self.offset] = b;
        self.offset += 1;
    }

    fn write_bytes(&mut self, bytes: &[u8]) {
        self.ensure_capacity(bytes.len());
        self.buf[self.offset..self.offset + bytes.len()].copy_from_slice(bytes);
        self.offset += bytes.len();
    }

    pub fn encode(&mut self, node: &Node) -> Vec<u8> {
        match node {
            Node::Object(obj) => {
                self.encode_object(obj);
            }
            Node::Array(arr) => {
                self.encode_array(arr);
            }
            Node::Value(val) => {
                self.encode_value(val);
            }
        }
        self.buf[..self.offset].to_vec()
    }

    fn encode_object(&mut self, obj: &Object) {
        let mut key_buf = Vec::new();
        let mut val_buf = Vec::new();

        for field in &obj.fields {
            let field_offset = self.offset;
            self.offset = 0;

            match &field.value {
                Node::Object(o) => self.encode_object(o),
                Node::Array(a) => self.encode_array(a),
                Node::Value(v) => self.encode_value(v),
            }

            let encoded = self.buf[..self.offset].to_vec();
            val_buf.extend_from_slice(&encoded);

            self.offset = 0;
            self.encode_string(&field.key);
            let encoded_key = self.buf[..self.offset].to_vec();
            key_buf.extend_from_slice(&encoded_key);

            self.offset = field_offset;
        }

        let total_len = key_buf.len() + val_buf.len();

        let mut prefix = PREFIX_CONTAINER | CONTAINER_MAP;

        if total_len < 254 {
            self.write_byte(prefix);
            self.write_bytes(&val_buf);
            self.write_bytes(&key_buf);
        } else if total_len < 65536 {
            prefix |= CONTAINER_LEN_1;
            self.write_byte(prefix);
            self.write_byte((total_len & 0xFF) as u8);
            self.write_bytes(&val_buf);
            self.write_bytes(&key_buf);
        } else {
            prefix |= CONTAINER_LEN_2;
            self.write_byte(prefix);
            self.write_byte((total_len >> 8 & 0xFF) as u8);
            self.write_byte((total_len & 0xFF) as u8);
            self.write_bytes(&val_buf);
            self.write_bytes(&key_buf);
        }

        if let Some(t) = &obj.tag {
            let tag_bytes = t.to_bytes();
            if !tag_bytes.is_empty() {
                self.encode_tag(&tag_bytes);
            }
        }
    }

    fn encode_array(&mut self, arr: &Array) {
        let mut val_buf = Vec::new();

        for item in &arr.items {
            let field_offset = self.offset;
            self.offset = 0;

            match item {
                Node::Object(o) => self.encode_object(o),
                Node::Array(a) => self.encode_array(a),
                Node::Value(v) => self.encode_value(v),
            }

            let encoded = self.buf[..self.offset].to_vec();
            val_buf.extend_from_slice(&encoded);

            self.offset = field_offset;
        }

        let total_len = val_buf.len();

        let mut prefix = PREFIX_CONTAINER | CONTAINER_ARRAY;

        if total_len < 254 {
            self.write_byte(prefix);
            self.write_bytes(&val_buf);
        } else if total_len < 65536 {
            prefix |= CONTAINER_LEN_1;
            self.write_byte(prefix);
            self.write_byte((total_len & 0xFF) as u8);
            self.write_bytes(&val_buf);
        } else {
            prefix |= CONTAINER_LEN_2;
            self.write_byte(prefix);
            self.write_byte((total_len >> 8 & 0xFF) as u8);
            self.write_byte((total_len & 0xFF) as u8);
            self.write_bytes(&val_buf);
        }

        if let Some(t) = &arr.tag {
            let tag_bytes = t.to_bytes();
            if !tag_bytes.is_empty() {
                self.encode_tag(&tag_bytes);
            }
        }
    }

    fn encode_value(&mut self, val: &Value) {
        match &val.data {
            crate::ir::ast::ValueData::Bool(b) => {
                self.encode_bool(*b);
            }
            crate::ir::ast::ValueData::String(s) => {
                self.encode_string(s);
            }
            crate::ir::ast::ValueData::Int(i) => {
                self.encode_int64(*i);
            }
            crate::ir::ast::ValueData::Float(f) => {
                self.encode_float(*f);
            }
            crate::ir::ast::ValueData::Bytes(b) => {
                self.encode_bytes(b);
            }
            crate::ir::ast::ValueData::Null => {
                self.encode_simple(SimpleValue::NullInt);
            }
            _ => {}
        }

        if let Some(t) = &val.tag {
            let tag_bytes = t.to_bytes();
            if !tag_bytes.is_empty() {
                self.encode_tag(&tag_bytes);
            }
        }
    }

    pub fn encode_bool(&mut self, v: bool) {
        let value = if v {
            SimpleValue::True
        } else {
            SimpleValue::False
        };
        self.encode_simple(value);
    }

    fn encode_simple(&mut self, value: SimpleValue) {
        self.write_byte(PREFIX_SIMPLE | value.to_byte());
    }

    pub fn encode_int64(&mut self, v: i64) {
        if v >= 0 {
            self.encode_uint64(v as u64);
        } else {
            let uv = if v == i64::MIN {
                9223372036854775808u64
            } else {
                (-v) as u64
            };
            self.encode_uint64_with_sign(PREFIX_NEGATIVE_INT, uv);
        }
    }

    pub fn encode_uint64(&mut self, v: u64) {
        self.encode_uint64_with_sign(PREFIX_POSITIVE_INT, v);
    }

    fn encode_uint64_with_sign(&mut self, sign: u8, v: u64) {
        if v < INT_LEN_1 as u64 {
            self.write_byte(sign | (v as u8));
        } else if v < MAX_1 as u64 {
            self.write_byte(sign | INT_LEN_1);
            self.write_byte(v as u8);
        } else if v < MAX_2 as u64 {
            self.write_byte(sign | INT_LEN_2);
            self.write_byte((v >> 8) as u8);
            self.write_byte(v as u8);
        } else if v < MAX_3 as u64 {
            self.write_byte(sign | INT_LEN_3);
            self.write_byte((v >> 16) as u8);
            self.write_byte((v >> 8) as u8);
            self.write_byte(v as u8);
        } else if v < MAX_4 as u64 {
            self.write_byte(sign | INT_LEN_4);
            self.write_byte((v >> 24) as u8);
            self.write_byte((v >> 16) as u8);
            self.write_byte((v >> 8) as u8);
            self.write_byte(v as u8);
        } else if v < MAX_5 {
            self.write_byte(sign | INT_LEN_5);
            self.write_byte((v >> 32) as u8);
            self.write_byte((v >> 24) as u8);
            self.write_byte((v >> 16) as u8);
            self.write_byte((v >> 8) as u8);
            self.write_byte(v as u8);
        } else if v < MAX_6 {
            self.write_byte(sign | INT_LEN_6);
            self.write_byte((v >> 40) as u8);
            self.write_byte((v >> 32) as u8);
            self.write_byte((v >> 24) as u8);
            self.write_byte((v >> 16) as u8);
            self.write_byte((v >> 8) as u8);
            self.write_byte(v as u8);
        } else if v < MAX_7 {
            self.write_byte(sign | INT_LEN_7);
            self.write_byte((v >> 48) as u8);
            self.write_byte((v >> 40) as u8);
            self.write_byte((v >> 32) as u8);
            self.write_byte((v >> 24) as u8);
            self.write_byte((v >> 16) as u8);
            self.write_byte((v >> 8) as u8);
            self.write_byte(v as u8);
        } else {
            self.write_byte(sign | INT_LEN_8);
            self.write_byte((v >> 56) as u8);
            self.write_byte((v >> 48) as u8);
            self.write_byte((v >> 40) as u8);
            self.write_byte((v >> 32) as u8);
            self.write_byte((v >> 24) as u8);
            self.write_byte((v >> 16) as u8);
            self.write_byte((v >> 8) as u8);
            self.write_byte(v as u8);
        }
    }

    pub fn encode_float(&mut self, f: f64) {
        let mut buffer = ryu::Buffer::new();
        let bytes = buffer.format(f).as_bytes();
        self.write_byte(PREFIX_FLOAT | FLOAT_POSITIVE_NEGATIVE_MASK);
        self.write_bytes(bytes);
    }

    pub fn encode_string(&mut self, s: &str) {
        let utf = s.as_bytes();
        let len = utf.len();
        let mut prefix = PREFIX_STRING;

        if len < 254 {
            self.write_byte(prefix | len as u8);
            self.write_bytes(utf);
        } else if len < 65536 {
            prefix |= STRING_LEN_1;
            self.write_byte(prefix);
            self.write_byte((len & 0xFF) as u8);
            self.write_bytes(utf);
        } else {
            prefix |= STRING_LEN_2;
            self.write_byte(prefix);
            self.write_byte((len >> 8 & 0xFF) as u8);
            self.write_byte((len & 0xFF) as u8);
            self.write_bytes(utf);
        }
    }

    pub fn encode_bytes(&mut self, bytes: &[u8]) {
        let len = bytes.len();
        let mut prefix = PREFIX_BYTES;

        if len < 254 {
            self.write_byte(prefix | len as u8);
            self.write_bytes(bytes);
        } else if len < 65536 {
            prefix |= BYTES_LEN_1;
            self.write_byte(prefix);
            self.write_byte((len & 0xFF) as u8);
            self.write_bytes(bytes);
        } else {
            prefix |= BYTES_LEN_2;
            self.write_byte(prefix);
            self.write_byte((len >> 8 & 0xFF) as u8);
            self.write_byte((len & 0xFF) as u8);
            self.write_bytes(bytes);
        }
    }

    fn encode_tag(&mut self, tag_bytes: &[u8]) {
        let len = tag_bytes.len();
        let mut prefix = PREFIX_TAG;

        if len < 254 {
            self.write_byte(prefix);
            self.write_byte(len as u8);
            self.write_bytes(tag_bytes);
        } else if len < 65536 {
            prefix |= TAG_LEN_1;
            self.write_byte(prefix);
            self.write_byte(len as u8);
            self.write_bytes(tag_bytes);
        } else {
            prefix |= TAG_LEN_2;
            self.write_byte(prefix);
            self.write_byte((len >> 8 & 0xFF) as u8);
            self.write_byte((len & 0xFF) as u8);
            self.write_bytes(tag_bytes);
        }
    }
}

impl Default for Encoder {
    fn default() -> Self {
        Self::new()
    }
}
