use std::fmt;
use std::io::{self, Read};

#[derive(Debug, Clone, PartialEq)]
enum Json {
    Null,
    Bool(bool),
    Number(String),
    String(String),
    Array(Vec<Json>),
    Object(Vec<(String, Json)>),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Error {
    byte: usize,
    message: String,
}

impl Error {
    fn new(byte: usize, message: impl Into<String>) -> Self {
        Self {
            byte,
            message: message.into(),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} at byte {}", self.message, self.byte)
    }
}

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

pub fn pretty_print(input: &str) -> Result<String, Error> {
    let mut parser = Parser::new(input);
    let value = parser.parse_value()?;
    parser.skip_ws();

    if !parser.is_done() {
        return Err(Error::new(parser.pos, "trailing characters"));
    }

    let mut out = String::new();
    write_json(&value, 0, &mut out);
    out.push('\n');
    Ok(out)
}

struct Parser<'a> {
    src: &'a str,
    pos: usize,
}

impl<'a> Parser<'a> {
    fn new(src: &'a str) -> Self {
        Self { src, pos: 0 }
    }

    fn is_done(&self) -> bool {
        self.pos >= self.src.len()
    }

    fn peek(&self) -> Option<char> {
        self.src[self.pos..].chars().next()
    }

    fn bump(&mut self) -> Option<char> {
        let ch = self.peek()?;
        self.pos += ch.len_utf8();
        Some(ch)
    }

    fn consume(&mut self, expected: char) -> Result<(), Error> {
        match self.bump() {
            Some(ch) if ch == expected => Ok(()),
            _ => Err(Error::new(self.pos, format!("expected '{}'", expected))),
        }
    }

    fn consume_word(&mut self, word: &str) -> bool {
        if self.src[self.pos..].starts_with(word) {
            self.pos += word.len();
            true
        } else {
            false
        }
    }

    fn skip_ws(&mut self) {
        while matches!(self.peek(), Some(' ' | '\n' | '\r' | '\t')) {
            self.bump();
        }
    }

    fn parse_value(&mut self) -> Result<Json, Error> {
        self.skip_ws();

        match self.peek() {
            Some('n') if self.consume_word("null") => Ok(Json::Null),
            Some('t') if self.consume_word("true") => Ok(Json::Bool(true)),
            Some('f') if self.consume_word("false") => Ok(Json::Bool(false)),
            Some('"') => Ok(Json::String(self.parse_string()?)),
            Some('[') => self.parse_array(),
            Some('{') => self.parse_object(),
            Some('-' | '0'..='9') => self.parse_number(),
            Some(_) => Err(Error::new(self.pos, "unexpected character")),
            None => Err(Error::new(self.pos, "unexpected end of input")),
        }
    }

    fn parse_array(&mut self) -> Result<Json, Error> {
        self.consume('[')?;
        self.skip_ws();

        let mut values = Vec::new();
        if self.peek() == Some(']') {
            self.bump();
            return Ok(Json::Array(values));
        }

        loop {
            values.push(self.parse_value()?);
            self.skip_ws();

            match self.peek() {
                Some(',') => {
                    self.bump();
                    self.skip_ws();
                }
                Some(']') => {
                    self.bump();
                    break;
                }
                _ => return Err(Error::new(self.pos, "expected ',' or ']'")),
            }
        }

        Ok(Json::Array(values))
    }

    fn parse_object(&mut self) -> Result<Json, Error> {
        self.consume('{')?;
        self.skip_ws();

        let mut pairs = Vec::new();
        if self.peek() == Some('}') {
            self.bump();
            return Ok(Json::Object(pairs));
        }

        loop {
            if self.peek() != Some('"') {
                return Err(Error::new(self.pos, "expected object key"));
            }

            let key = self.parse_string()?;
            self.skip_ws();
            self.consume(':')?;
            let value = self.parse_value()?;
            pairs.push((key, value));
            self.skip_ws();

            match self.peek() {
                Some(',') => {
                    self.bump();
                    self.skip_ws();
                }
                Some('}') => {
                    self.bump();
                    break;
                }
                _ => return Err(Error::new(self.pos, "expected ',' or '}'")),
            }
        }

        Ok(Json::Object(pairs))
    }

    fn parse_string(&mut self) -> Result<String, Error> {
        self.consume('"')?;
        let mut out = String::new();

        loop {
            match self.bump() {
                Some('"') => return Ok(out),
                Some('\\') => out.push(self.parse_escape()?),
                Some(ch) if ch <= '\u{1f}' => {
                    return Err(Error::new(self.pos, "control character in string"));
                }
                Some(ch) => out.push(ch),
                None => return Err(Error::new(self.pos, "unterminated string")),
            }
        }
    }

    fn parse_escape(&mut self) -> Result<char, Error> {
        match self.bump() {
            Some('"') => Ok('"'),
            Some('\\') => Ok('\\'),
            Some('/') => Ok('/'),
            Some('b') => Ok('\u{08}'),
            Some('f') => Ok('\u{0c}'),
            Some('n') => Ok('\n'),
            Some('r') => Ok('\r'),
            Some('t') => Ok('\t'),
            Some('u') => {
                let code = self.parse_hex4()?;

                if (0xd800..=0xdbff).contains(&code) {
                    let saved = self.pos;
                    if self.bump() == Some('\\') && self.bump() == Some('u') {
                        let low = self.parse_hex4()?;
                        if (0xdc00..=0xdfff).contains(&low) {
                            let full = 0x10000 + ((code - 0xd800) << 10) + (low - 0xdc00);
                            return char::from_u32(full)
                                .ok_or_else(|| Error::new(self.pos, "invalid unicode escape"));
                        }
                    }
                    self.pos = saved;
                    Err(Error::new(self.pos, "missing low surrogate"))
                } else if (0xdc00..=0xdfff).contains(&code) {
                    Err(Error::new(self.pos, "unexpected low surrogate"))
                } else {
                    char::from_u32(code).ok_or_else(|| Error::new(self.pos, "invalid unicode escape"))
                }
            }
            _ => Err(Error::new(self.pos, "invalid escape sequence")),
        }
    }

    fn parse_hex4(&mut self) -> Result<u32, Error> {
        let mut value = 0;

        for _ in 0..4 {
            let ch = self
                .bump()
                .ok_or_else(|| Error::new(self.pos, "short unicode escape"))?;
            value = value * 16
                + ch.to_digit(16)
                    .ok_or_else(|| Error::new(self.pos, "invalid unicode escape"))?;
        }

        Ok(value)
    }

    fn parse_number(&mut self) -> Result<Json, Error> {
        let start = self.pos;

        if self.peek() == Some('-') {
            self.bump();
        }

        match self.peek() {
            Some('0') => {
                self.bump();
            }
            Some('1'..='9') => {
                self.bump();
                while matches!(self.peek(), Some('0'..='9')) {
                    self.bump();
                }
            }
            _ => return Err(Error::new(self.pos, "invalid number")),
        }

        if self.peek() == Some('.') {
            self.bump();
            if !matches!(self.peek(), Some('0'..='9')) {
                return Err(Error::new(self.pos, "expected digit after decimal point"));
            }
            while matches!(self.peek(), Some('0'..='9')) {
                self.bump();
            }
        }

        if matches!(self.peek(), Some('e' | 'E')) {
            self.bump();
            if matches!(self.peek(), Some('+' | '-')) {
                self.bump();
            }
            if !matches!(self.peek(), Some('0'..='9')) {
                return Err(Error::new(self.pos, "expected exponent digit"));
            }
            while matches!(self.peek(), Some('0'..='9')) {
                self.bump();
            }
        }

        Ok(Json::Number(self.src[start..self.pos].to_string()))
    }
}

fn write_json(value: &Json, depth: usize, out: &mut String) {
    match value {
        Json::Null => out.push_str("null"),
        Json::Bool(value) => out.push_str(if *value { "true" } else { "false" }),
        Json::Number(value) => out.push_str(value),
        Json::String(value) => write_string(value, out),
        Json::Array(values) => write_array(values, depth, out),
        Json::Object(pairs) => write_object(pairs, depth, out),
    }
}

fn write_array(values: &[Json], depth: usize, out: &mut String) {
    if values.is_empty() {
        out.push_str("[]");
        return;
    }

    out.push('[');
    for (index, value) in values.iter().enumerate() {
        out.push('\n');
        write_indent(depth + 1, out);
        write_json(value, depth + 1, out);
        if index + 1 != values.len() {
            out.push(',');
        }
    }
    out.push('\n');
    write_indent(depth, out);
    out.push(']');
}

fn write_object(pairs: &[(String, Json)], depth: usize, out: &mut String) {
    if pairs.is_empty() {
        out.push_str("{}");
        return;
    }

    out.push('{');
    for (index, (key, value)) in pairs.iter().enumerate() {
        out.push('\n');
        write_indent(depth + 1, out);
        write_string(key, out);
        out.push_str(": ");
        write_json(value, depth + 1, out);
        if index + 1 != pairs.len() {
            out.push(',');
        }
    }
    out.push('\n');
    write_indent(depth, out);
    out.push('}');
}

fn write_string(value: &str, out: &mut String) {
    out.push('"');

    for ch in value.chars() {
        match ch {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            '\u{08}' => out.push_str("\\b"),
            '\u{0c}' => out.push_str("\\f"),
            ch if ch <= '\u{1f}' => out.push_str(&format!("\\u{:04x}", ch as u32)),
            ch => out.push(ch),
        }
    }

    out.push('"');
}

fn write_indent(depth: usize, out: &mut String) {
    for _ in 0..depth {
        out.push_str("  ");
    }
}

fn main() {
    let mut input = String::new();

    if let Err(error) = io::stdin().read_to_string(&mut input) {
        eprintln!("read error: {error}");
        std::process::exit(1);
    }

    match pretty_print(&input) {
        Ok(output) => print!("{output}"),
        Err(error) => {
            eprintln!("json error: {error}");
            std::process::exit(1);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pretty_prints_nested_values() {
        let input = r#"{"name":"Ada","items":[1,true,null,{"x":"y"}]}"#;
        let expected = concat!(
            "{\n",
            "  \"name\": \"Ada\",\n",
            "  \"items\": [\n",
            "    1,\n",
            "    true,\n",
            "    null,\n",
            "    {\n",
            "      \"x\": \"y\"\n",
            "    }\n",
            "  ]\n",
            "}\n",
        );

        assert_eq!(pretty_print(input).unwrap(), expected);
    }

    #[test]
    fn escapes_strings_on_output() {
        let input = r#"{"text":"line\nquote: \" and snowman: \u2603"}"#;
        let output = pretty_print(input).unwrap();

        assert!(output.contains(r#""text": "line\nquote: \" and snowman: ☃""#));
    }

    #[test]
    fn rejects_trailing_commas() {
        assert!(pretty_print(r#"{"a":1,}"#).is_err());
        assert!(pretty_print(r#"[1,2,]"#).is_err());
    }
}