use std::fmt;
use std::io::Cursor;

use base64::engine::general_purpose::STANDARD;
use base64::Engine;
use rmpv::Value;
use serde_json::{Map, Number};

#[derive(Debug)]
pub enum MessageError {
    Json(serde_json::Error),
    Encode(rmpv::encode::Error),
    Decode(rmpv::decode::Error),
    InvalidBase64(base64::DecodeError),
    InvalidNumber,
    TrailingBytes,
}

impl fmt::Display for MessageError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Json(error) => error.fmt(formatter),
            Self::Encode(error) => error.fmt(formatter),
            Self::Decode(error) => error.fmt(formatter),
            Self::InvalidBase64(error) => error.fmt(formatter),
            Self::InvalidNumber => {
                formatter.write_str("number cannot be represented in MessagePack")
            }
            Self::TrailingBytes => formatter.write_str("MessagePack payload has trailing bytes"),
        }
    }
}

impl std::error::Error for MessageError {}

pub fn encode_tagged_json(json: &str) -> Result<Vec<u8>, MessageError> {
    let value: serde_json::Value = serde_json::from_str(json).map_err(MessageError::Json)?;
    let message = json_to_message(value)?;
    let mut output = Vec::new();
    rmpv::encode::write_value(&mut output, &message).map_err(MessageError::Encode)?;
    Ok(output)
}

pub fn decode_tagged_json(bytes: &[u8]) -> Result<String, MessageError> {
    if bytes.is_empty() {
        return Ok("null".to_string());
    }
    let mut cursor = Cursor::new(bytes);
    let value = rmpv::decode::read_value(&mut cursor).map_err(MessageError::Decode)?;
    if cursor.position() != bytes.len() as u64 {
        return Err(MessageError::TrailingBytes);
    }
    serde_json::to_string(&message_to_json(value)?).map_err(MessageError::Json)
}

fn json_to_message(value: serde_json::Value) -> Result<Value, MessageError> {
    match value {
        serde_json::Value::Null => Ok(Value::Nil),
        serde_json::Value::Bool(value) => Ok(Value::Boolean(value)),
        serde_json::Value::Number(value) => {
            if let Some(value) = value.as_i64() {
                Ok(Value::from(value))
            } else if let Some(value) = value.as_u64() {
                Ok(Value::from(value))
            } else {
                value
                    .as_f64()
                    .map(Value::F64)
                    .ok_or(MessageError::InvalidNumber)
            }
        }
        serde_json::Value::String(value) => Ok(Value::from(value)),
        serde_json::Value::Array(values) => values
            .into_iter()
            .map(json_to_message)
            .collect::<Result<Vec<_>, _>>()
            .map(Value::Array),
        serde_json::Value::Object(mut values) => {
            if values.len() == 1 {
                if let Some(serde_json::Value::String(encoded)) = values.remove("$bin") {
                    return STANDARD
                        .decode(encoded)
                        .map(Value::Binary)
                        .map_err(MessageError::InvalidBase64);
                }
            }
            values
                .into_iter()
                .map(|(key, value)| {
                    let key = key
                        .strip_prefix("$int:")
                        .and_then(|raw| raw.parse::<i64>().ok())
                        .map(Value::from)
                        .unwrap_or_else(|| Value::from(key));
                    Ok((key, json_to_message(value)?))
                })
                .collect::<Result<Vec<_>, MessageError>>()
                .map(Value::Map)
        }
    }
}

fn message_to_json(value: Value) -> Result<serde_json::Value, MessageError> {
    match value {
        Value::Nil => Ok(serde_json::Value::Null),
        Value::Boolean(value) => Ok(serde_json::Value::Bool(value)),
        Value::Integer(value) => value
            .as_i64()
            .map(|value| serde_json::Value::Number(Number::from(value)))
            .or_else(|| {
                value
                    .as_u64()
                    .map(|value| serde_json::Value::Number(Number::from(value)))
            })
            .ok_or(MessageError::InvalidNumber),
        Value::F32(value) => Number::from_f64(f64::from(value))
            .map(serde_json::Value::Number)
            .ok_or(MessageError::InvalidNumber),
        Value::F64(value) => Number::from_f64(value)
            .map(serde_json::Value::Number)
            .ok_or(MessageError::InvalidNumber),
        Value::String(value) => Ok(serde_json::Value::String(
            value.as_str().unwrap_or_default().to_string(),
        )),
        Value::Binary(value) => {
            let mut output = Map::new();
            output.insert(
                "$bin".to_string(),
                serde_json::Value::String(STANDARD.encode(value)),
            );
            Ok(serde_json::Value::Object(output))
        }
        Value::Array(values) => values
            .into_iter()
            .map(message_to_json)
            .collect::<Result<Vec<_>, _>>()
            .map(serde_json::Value::Array),
        Value::Map(values) => {
            let mut output = Map::new();
            for (key, value) in values {
                let key = match key {
                    Value::String(value) => value.as_str().unwrap_or_default().to_string(),
                    Value::Integer(value) => value.to_string(),
                    other => other.to_string(),
                };
                output.insert(key, message_to_json(value)?);
            }
            Ok(serde_json::Value::Object(output))
        }
        Value::Ext(1, bytes) => {
            let mut cursor = Cursor::new(bytes);
            let nested = rmpv::decode::read_value(&mut cursor).map_err(MessageError::Decode)?;
            message_to_json(nested)
        }
        Value::Ext(kind, bytes) => {
            let mut output = Map::new();
            output.insert(
                "$ext".to_string(),
                serde_json::Value::Number(Number::from(kind)),
            );
            output.insert(
                "$bin".to_string(),
                serde_json::Value::String(STANDARD.encode(bytes)),
            );
            Ok(serde_json::Value::Object(output))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{decode_tagged_json, encode_tagged_json};

    #[test]
    fn tagged_json_encodes_integer_keys_and_decodes_compatible_json() {
        let source = r#"{"$int:7":{"$bin":"AQID"},"enabled":true}"#;
        let encoded = encode_tagged_json(source).unwrap();
        let decoded: serde_json::Value =
            serde_json::from_str(&decode_tagged_json(&encoded).unwrap()).unwrap();
        assert_eq!(decoded["7"]["$bin"], "AQID");
        assert_eq!(decoded["enabled"], true);
    }
}
