use armour_typ::{ScalarTyp, SchemaFields, SchemaTyp};

use crate::max_cap::{VEC_MAX_CAP, VEC_MAX_SIZE_OF};
use crate::{Rapira, RapiraError, Result, byte_rapira, push};

// Wire tags are protocol commitments and must remain stable for backward compatibility.
const SCALAR_BOOL: u8 = 0;
const SCALAR_U8: u8 = 1;
const SCALAR_U16: u8 = 2;
const SCALAR_U32: u8 = 3;
const SCALAR_U64: u8 = 4;
const SCALAR_I32: u8 = 5;
const SCALAR_I64: u8 = 6;
const SCALAR_F32: u8 = 7;
const SCALAR_F64: u8 = 8;
const SCALAR_STR: u8 = 9;
const SCALAR_DATETIME: u8 = 10;
const SCALAR_TIMESTAMP: u8 = 11;
const SCALAR_DECIMAL: u8 = 12;
const SCALAR_ID32: u8 = 13;
const SCALAR_ID64: u8 = 14;
const SCALAR_FUID: u8 = 15;
const SCALAR_LOW_ID: u8 = 16;
const SCALAR_BYTES: u8 = 17;
const SCALAR_VOID: u8 = 18;
const SCALAR_RUST_JSON: u8 = 19;
const SCALAR_JSON_BYTES: u8 = 20;
const SCALAR_ARRAY_BYTES: u8 = 21;

const SCHEMA_SCALAR: u8 = 0;
const SCHEMA_ARRAY: u8 = 1;
const SCHEMA_VEC: u8 = 2;
const SCHEMA_OPTIONAL: u8 = 3;
const SCHEMA_SIMPLE_ENUM: u8 = 4;
const SCHEMA_STRUCT: u8 = 5;
const SCHEMA_ENUM: u8 = 6;
const SCHEMA_CUSTOM: u8 = 7;

const FIELDS_NAMED: u8 = 0;
const FIELDS_UNNAMED: u8 = 1;

impl Rapira for ScalarTyp {
    const STATIC_SIZE: Option<usize> = None;
    const MIN_SIZE: usize = 1;

    fn size(&self) -> usize {
        match self {
            ScalarTyp::ArrayBytes(_) => 1 + u32::MIN_SIZE,
            _ => 1,
        }
    }

    fn check_bytes(slice: &mut &[u8]) -> Result<()> {
        let tag = byte_rapira::from_slice(slice)?;
        match tag {
            SCALAR_BOOL | SCALAR_U8 | SCALAR_U16 | SCALAR_U32 | SCALAR_U64 | SCALAR_I32
            | SCALAR_I64 | SCALAR_F32 | SCALAR_F64 | SCALAR_STR | SCALAR_DATETIME
            | SCALAR_TIMESTAMP | SCALAR_DECIMAL | SCALAR_ID32 | SCALAR_ID64 | SCALAR_FUID
            | SCALAR_LOW_ID | SCALAR_BYTES | SCALAR_VOID | SCALAR_RUST_JSON | SCALAR_JSON_BYTES => {
                Ok(())
            }
            SCALAR_ARRAY_BYTES => {
                u32::check_bytes(slice)?;
                Ok(())
            }
            _ => Err(RapiraError::EnumVariant),
        }
    }

    fn from_slice(slice: &mut &[u8]) -> Result<Self> {
        let tag = byte_rapira::from_slice(slice)?;
        match tag {
            SCALAR_BOOL => Ok(ScalarTyp::Bool),
            SCALAR_U8 => Ok(ScalarTyp::U8),
            SCALAR_U16 => Ok(ScalarTyp::U16),
            SCALAR_U32 => Ok(ScalarTyp::U32),
            SCALAR_U64 => Ok(ScalarTyp::U64),
            SCALAR_I32 => Ok(ScalarTyp::I32),
            SCALAR_I64 => Ok(ScalarTyp::I64),
            SCALAR_F32 => Ok(ScalarTyp::F32),
            SCALAR_F64 => Ok(ScalarTyp::F64),
            SCALAR_STR => Ok(ScalarTyp::Str),
            SCALAR_DATETIME => Ok(ScalarTyp::Datetime),
            SCALAR_TIMESTAMP => Ok(ScalarTyp::Timestamp),
            SCALAR_DECIMAL => Ok(ScalarTyp::Decimal),
            SCALAR_ID32 => Ok(ScalarTyp::Id32),
            SCALAR_ID64 => Ok(ScalarTyp::Id64),
            SCALAR_FUID => Ok(ScalarTyp::Fuid),
            SCALAR_LOW_ID => Ok(ScalarTyp::LowId),
            SCALAR_BYTES => Ok(ScalarTyp::Bytes),
            SCALAR_VOID => Ok(ScalarTyp::Void),
            SCALAR_RUST_JSON => Ok(ScalarTyp::RustJson),
            SCALAR_JSON_BYTES => Ok(ScalarTyp::JsonBytes),
            SCALAR_ARRAY_BYTES => Ok(ScalarTyp::ArrayBytes(u32::from_slice(slice)?)),
            _ => Err(RapiraError::EnumVariant),
        }
    }

    fn convert_to_bytes(&self, slice: &mut [u8], cursor: &mut usize) {
        match self {
            ScalarTyp::Bool => push(slice, cursor, SCALAR_BOOL),
            ScalarTyp::U8 => push(slice, cursor, SCALAR_U8),
            ScalarTyp::U16 => push(slice, cursor, SCALAR_U16),
            ScalarTyp::U32 => push(slice, cursor, SCALAR_U32),
            ScalarTyp::U64 => push(slice, cursor, SCALAR_U64),
            ScalarTyp::I32 => push(slice, cursor, SCALAR_I32),
            ScalarTyp::I64 => push(slice, cursor, SCALAR_I64),
            ScalarTyp::F32 => push(slice, cursor, SCALAR_F32),
            ScalarTyp::F64 => push(slice, cursor, SCALAR_F64),
            ScalarTyp::Str => push(slice, cursor, SCALAR_STR),
            ScalarTyp::Datetime => push(slice, cursor, SCALAR_DATETIME),
            ScalarTyp::Timestamp => push(slice, cursor, SCALAR_TIMESTAMP),
            ScalarTyp::Decimal => push(slice, cursor, SCALAR_DECIMAL),
            ScalarTyp::Id32 => push(slice, cursor, SCALAR_ID32),
            ScalarTyp::Id64 => push(slice, cursor, SCALAR_ID64),
            ScalarTyp::Fuid => push(slice, cursor, SCALAR_FUID),
            ScalarTyp::LowId => push(slice, cursor, SCALAR_LOW_ID),
            ScalarTyp::Bytes => push(slice, cursor, SCALAR_BYTES),
            ScalarTyp::Void => push(slice, cursor, SCALAR_VOID),
            ScalarTyp::RustJson => push(slice, cursor, SCALAR_RUST_JSON),
            ScalarTyp::JsonBytes => push(slice, cursor, SCALAR_JSON_BYTES),
            ScalarTyp::ArrayBytes(n) => {
                push(slice, cursor, SCALAR_ARRAY_BYTES);
                n.convert_to_bytes(slice, cursor);
            }
        }
    }
}

impl Rapira for SchemaFields {
    const STATIC_SIZE: Option<usize> = None;
    const MIN_SIZE: usize = 1 + u32::MIN_SIZE;

    fn size(&self) -> usize {
        1 + match self {
            SchemaFields::Named(fields) => fields.size(),
            SchemaFields::Unnamed(fields) => fields.size(),
        }
    }

    fn check_bytes(slice: &mut &[u8]) -> Result<()> {
        let tag = byte_rapira::from_slice(slice)?;
        match tag {
            FIELDS_NAMED => Vec::<(String, SchemaTyp)>::check_bytes(slice),
            FIELDS_UNNAMED => Vec::<SchemaTyp>::check_bytes(slice),
            _ => Err(RapiraError::EnumVariant),
        }
    }

    fn from_slice(slice: &mut &[u8]) -> Result<Self> {
        let tag = byte_rapira::from_slice(slice)?;
        match tag {
            FIELDS_NAMED => Ok(SchemaFields::Named(Vec::<(String, SchemaTyp)>::from_slice(
                slice,
            )?)),
            FIELDS_UNNAMED => Ok(SchemaFields::Unnamed(Vec::<SchemaTyp>::from_slice(slice)?)),
            _ => Err(RapiraError::EnumVariant),
        }
    }

    fn convert_to_bytes(&self, slice: &mut [u8], cursor: &mut usize) {
        match self {
            SchemaFields::Named(fields) => {
                push(slice, cursor, FIELDS_NAMED);
                fields.convert_to_bytes(slice, cursor);
            }
            SchemaFields::Unnamed(fields) => {
                push(slice, cursor, FIELDS_UNNAMED);
                fields.convert_to_bytes(slice, cursor);
            }
        }
    }
}

impl Rapira for SchemaTyp {
    const STATIC_SIZE: Option<usize> = None;
    const MIN_SIZE: usize = 1;

    fn size(&self) -> usize {
        1 + match self {
            SchemaTyp::Scalar(s) => s.size(),
            SchemaTyp::Array(_, inner) => u32::MIN_SIZE + inner.size(),
            SchemaTyp::Vec(inner) | SchemaTyp::Optional(inner) => inner.size(),
            SchemaTyp::SimpleEnum { name, variants } => name.size() + simple_enum_size(variants),
            SchemaTyp::Struct { name, fields } => name.size() + fields.size(),
            SchemaTyp::Enum { name, variants } => name.size() + enum_size(variants),
            SchemaTyp::Custom(name, args) => name.size() + args.size(),
        }
    }

    fn check_bytes(slice: &mut &[u8]) -> Result<()> {
        let tag = byte_rapira::from_slice(slice)?;
        match tag {
            SCHEMA_SCALAR => ScalarTyp::check_bytes(slice),
            SCHEMA_ARRAY => {
                u32::check_bytes(slice)?;
                SchemaTyp::check_bytes(slice)
            }
            SCHEMA_VEC | SCHEMA_OPTIONAL => SchemaTyp::check_bytes(slice),
            SCHEMA_SIMPLE_ENUM => {
                String::check_bytes(slice)?;
                check_simple_enum_variants(slice)?;
                Ok(())
            }
            SCHEMA_STRUCT => {
                String::check_bytes(slice)?;
                SchemaFields::check_bytes(slice)
            }
            SCHEMA_ENUM => {
                String::check_bytes(slice)?;
                check_enum_variants(slice)?;
                Ok(())
            }
            SCHEMA_CUSTOM => {
                String::check_bytes(slice)?;
                Vec::<SchemaTyp>::check_bytes(slice)
            }
            _ => Err(RapiraError::EnumVariant),
        }
    }

    fn from_slice(slice: &mut &[u8]) -> Result<Self> {
        let tag = byte_rapira::from_slice(slice)?;
        match tag {
            SCHEMA_SCALAR => Ok(SchemaTyp::Scalar(ScalarTyp::from_slice(slice)?)),
            SCHEMA_ARRAY => Ok(SchemaTyp::Array(
                u32::from_slice(slice)?,
                Box::new(SchemaTyp::from_slice(slice)?),
            )),
            SCHEMA_VEC => Ok(SchemaTyp::Vec(Box::new(SchemaTyp::from_slice(slice)?))),
            SCHEMA_OPTIONAL => Ok(SchemaTyp::Optional(Box::new(SchemaTyp::from_slice(slice)?))),
            SCHEMA_SIMPLE_ENUM => Ok(SchemaTyp::SimpleEnum {
                name: String::from_slice(slice)?,
                variants: read_simple_enum_variants(slice)?,
            }),
            SCHEMA_STRUCT => Ok(SchemaTyp::Struct {
                name: String::from_slice(slice)?,
                fields: SchemaFields::from_slice(slice)?,
            }),
            SCHEMA_ENUM => Ok(SchemaTyp::Enum {
                name: String::from_slice(slice)?,
                variants: read_enum_variants(slice)?,
            }),
            SCHEMA_CUSTOM => Ok(SchemaTyp::Custom(
                String::from_slice(slice)?,
                Vec::<SchemaTyp>::from_slice(slice)?,
            )),
            _ => Err(RapiraError::EnumVariant),
        }
    }

    fn convert_to_bytes(&self, slice: &mut [u8], cursor: &mut usize) {
        match self {
            SchemaTyp::Scalar(s) => {
                push(slice, cursor, SCHEMA_SCALAR);
                s.convert_to_bytes(slice, cursor);
            }
            SchemaTyp::Array(n, inner) => {
                push(slice, cursor, SCHEMA_ARRAY);
                n.convert_to_bytes(slice, cursor);
                inner.convert_to_bytes(slice, cursor);
            }
            SchemaTyp::Vec(inner) => {
                push(slice, cursor, SCHEMA_VEC);
                inner.convert_to_bytes(slice, cursor);
            }
            SchemaTyp::Optional(inner) => {
                push(slice, cursor, SCHEMA_OPTIONAL);
                inner.convert_to_bytes(slice, cursor);
            }
            SchemaTyp::SimpleEnum { name, variants } => {
                push(slice, cursor, SCHEMA_SIMPLE_ENUM);
                name.convert_to_bytes(slice, cursor);
                write_simple_enum_variants(variants, slice, cursor);
            }
            SchemaTyp::Struct { name, fields } => {
                push(slice, cursor, SCHEMA_STRUCT);
                name.convert_to_bytes(slice, cursor);
                fields.convert_to_bytes(slice, cursor);
            }
            SchemaTyp::Enum { name, variants } => {
                push(slice, cursor, SCHEMA_ENUM);
                name.convert_to_bytes(slice, cursor);
                write_enum_variants(variants, slice, cursor);
            }
            SchemaTyp::Custom(name, args) => {
                push(slice, cursor, SCHEMA_CUSTOM);
                name.convert_to_bytes(slice, cursor);
                args.convert_to_bytes(slice, cursor);
            }
        }
    }
}

#[inline]
fn simple_enum_size(variants: &[(u8, String)]) -> usize {
    variants
        .iter()
        .fold(4, |size, (_, name)| size + 1 + name.size())
}

#[inline]
fn enum_size(variants: &[(u8, String, SchemaTyp)]) -> usize {
    variants
        .iter()
        .fold(4, |size, (_, name, ty)| size + 1 + name.size() + ty.size())
}

fn check_simple_enum_variants(slice: &mut &[u8]) -> Result<()> {
    let len = u32::from_slice(slice)? as usize;
    check_vec_len::<(u8, String)>(len)?;
    for _ in 0..len {
        byte_rapira::from_slice(slice)?;
        String::check_bytes(slice)?;
    }
    Ok(())
}

fn check_enum_variants(slice: &mut &[u8]) -> Result<()> {
    let len = u32::from_slice(slice)? as usize;
    check_vec_len::<(u8, String, SchemaTyp)>(len)?;
    for _ in 0..len {
        byte_rapira::from_slice(slice)?;
        String::check_bytes(slice)?;
        SchemaTyp::check_bytes(slice)?;
    }
    Ok(())
}

fn read_simple_enum_variants(slice: &mut &[u8]) -> Result<Vec<(u8, String)>> {
    let len = u32::from_slice(slice)? as usize;
    check_vec_len::<(u8, String)>(len)?;
    let mut variants = Vec::with_capacity(len);
    for _ in 0..len {
        let tag = byte_rapira::from_slice(slice)?;
        let name = String::from_slice(slice)?;
        variants.push((tag, name));
    }
    Ok(variants)
}

fn read_enum_variants(slice: &mut &[u8]) -> Result<Vec<(u8, String, SchemaTyp)>> {
    let len = u32::from_slice(slice)? as usize;
    check_vec_len::<(u8, String, SchemaTyp)>(len)?;
    let mut variants = Vec::with_capacity(len);
    for _ in 0..len {
        let tag = byte_rapira::from_slice(slice)?;
        let name = String::from_slice(slice)?;
        let payload = SchemaTyp::from_slice(slice)?;
        variants.push((tag, name, payload));
    }
    Ok(variants)
}

#[inline]
fn check_vec_len<T>(len: usize) -> Result<()> {
    if len > VEC_MAX_CAP {
        return Err(RapiraError::MaxCapacity);
    }
    let size = core::mem::size_of::<T>()
        .checked_mul(len)
        .ok_or(RapiraError::MaxSize)?;
    if size > VEC_MAX_SIZE_OF {
        return Err(RapiraError::MaxSize);
    }
    Ok(())
}

#[inline]
fn write_simple_enum_variants(variants: &[(u8, String)], slice: &mut [u8], cursor: &mut usize) {
    (variants.len() as u32).convert_to_bytes(slice, cursor);
    for (tag, name) in variants {
        push(slice, cursor, *tag);
        name.convert_to_bytes(slice, cursor);
    }
}

#[inline]
fn write_enum_variants(variants: &[(u8, String, SchemaTyp)], slice: &mut [u8], cursor: &mut usize) {
    (variants.len() as u32).convert_to_bytes(slice, cursor);
    for (tag, name, payload) in variants {
        push(slice, cursor, *tag);
        name.convert_to_bytes(slice, cursor);
        payload.convert_to_bytes(slice, cursor);
    }
}
