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

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

struct Parser<'a> {
    input: &'a [u8],
    pos: usize,
}

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

    fn parse(mut self) -> Result<Json, String> {
        self.skip_whitespace();
        let value = self.parse_value()?;
        self.skip_whitespace();

        if self.pos != self.input.len() {
            return Err(self.error("unexpected trailing input"));
        }

        Ok(value)
    }

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

        match self.peek() {
            Some(b'n') => {
                self.expect_bytes(b"null")?;
                Ok(Json::Null)
            }
            Some(b't') => {
                self.expect_bytes(b"true")?;
                Ok(Json::Bool(true))
            }
            Some(b'f') => {
                self.expect_bytes(b"false")?;
                Ok(Json::Bool(false))
            }
            Some(b'"') => Ok(Json::String(self.parse_string()?)),
            Some(b'[') => self.parse_array(),
            Some(b'{') => self.parse_object(),
            Some(b'-' | b'0'..=b'9') => self.parse_number(),
            Some(_) => Err(self.error("expected a JSON value")),
            None => Err(self.error("unexpected end of input")),
        }
    }

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

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

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

            if self.take_if(b']') {
                break;
            }

            self.consume(b',')?;
            self.skip_whitespace();
        }

        Ok(Json::Array(values))
    }

    fn parse_object(&mut self) -> Result<Json, String> {
        self.consume(b'{')?;
        self.skip_whitespace();

        let mut entries = Vec::new();
        if self.take_if(b'}') {
            return Ok(Json::Object(entries));
        }

        loop {
            if self.peek() != Some(b'"') {
                return Err(self.error("expected an object key"));
            }

            let key = self.parse_string()?;
            self.skip_whitespace();
            self.consume(b':')?;
            let value = self.parse_value()?;
            entries.push((key, value));
            self.skip_whitespace();

            if self.take_if(b'}') {
                break;
            }

            self.consume(b',')?;
            self.skip_whitespace();
        }

        Ok(Json::Object(entries))
    }

    // Strings are kept as raw JSON lexemes so valid escapes remain unchanged.
    fn parse_string(&mut self) -> Result<String, String> {
        let start = self.pos;
        self.consume(b'"')?;

        loop {
            match self.peek() {
                Some(b'"') => {
                    self.pos += 1;
                    return String::from_utf8(self.input[start..self.pos].to_vec())
                        .map_err(|_| self.error("invalid UTF-8 in string"));
                }
                Some(b'\\') => {
                    self.pos += 1;
                    match self.peek() {
                        Some(b'"' | b'\\' | b'/' | b'b' | b'f' | b'n' | b'r' | b't') => {
                            self.pos += 1;
                        }
                        Some(b'u') => {
                            self.pos += 1;
                            for _ in 0..4 {
                                match self.peek() {
                                    Some(c) if c.is_ascii_hexdigit() => self.pos += 1,
                                    _ => return Err(self.error("invalid Unicode escape")),
                                }
                            }
                        }
                        _ => return Err(self.error("invalid escape sequence")),
                    }
                }
                Some(0x00..=0x1f) => {
                    return Err(self.error("unescaped control character in string"));
                }
                Some(_) => self.pos += 1,
                None => return Err(self.error("unterminated string")),
            }
        }
    }

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

        self.take_if(b'-');

        match self.peek() {
            Some(b'0') => {
                self.pos += 1;
                if matches!(self.peek(), Some(b'0'..=b'9')) {
                    return Err(self.error("leading zero in number"));
                }
            }
            Some(b'1'..=b'9') => {
                self.pos += 1;
                while matches!(self.peek(), Some(b'0'..=b'9')) {
                    self.pos += 1;
                }
            }
            _ => return Err(self.error("invalid number")),
        }

        if self.take_if(b'.') {
            let digits_start = self.pos;
            while matches!(self.peek(), Some(b'0'..=b'9')) {
                self.pos += 1;
            }
            if self.pos == digits_start {
                return Err(self.error("expected digits after decimal point"));
            }
        }

        if matches!(self.peek(), Some(b'e' | b'E')) {
            self.pos += 1;
            if matches!(self.peek(), Some(b'+' | b'-')) {
                self.pos += 1;
            }

            let digits_start = self.pos;
            while matches!(self.peek(), Some(b'0'..=b'9')) {
                self.pos += 1;
            }
            if self.pos == digits_start {
                return Err(self.error("expected exponent digits"));
            }
        }

        let number = std::str::from_utf8(&self.input[start..self.pos])
            .expect("number tokens contain only ASCII")
            .to_owned();

        Ok(Json::Number(number))
    }

    fn skip_whitespace(&mut self) {
        while matches!(self.peek(), Some(b' ' | b'\n' | b'\r' | b'\t')) {
            self.pos += 1;
        }
    }

    fn expect_bytes(&mut self, expected: &[u8]) -> Result<(), String> {
        if self.input.get(self.pos..self.pos + expected.len()) == Some(expected) {
            self.pos += expected.len();
            Ok(())
        } else {
            Err(self.error("invalid literal"))
        }
    }

    fn consume(&mut self, expected: u8) -> Result<(), String> {
        if self.take_if(expected) {
            Ok(())
        } else {
            Err(self.error(&format!("expected '{}'", expected as char)))
        }
    }

    fn take_if(&mut self, expected: u8) -> bool {
        if self.peek() == Some(expected) {
            self.pos += 1;
            true
        } else {
            false
        }
    }

    fn peek(&self) -> Option<u8> {
        self.input.get(self.pos).copied()
    }

    fn error(&self, message: &str) -> String {
        format!("{message} at byte {}", self.pos)
    }
}

fn write_pretty(value: &Json, output: &mut String, depth: usize) {
    match value {
        Json::Null => output.push_str("null"),
        Json::Bool(value) => output.push_str(if *value { "true" } else { "false" }),
        Json::Number(value) | Json::String(value) => output.push_str(value),
        Json::Array(values) => {
            output.push('[');
            if !values.is_empty() {
                output.push('\n');
                for (index, value) in values.iter().enumerate() {
                    indent(output, depth + 1);
                    write_pretty(value, output, depth + 1);
                    if index + 1 != values.len() {
                        output.push(',');
                    }
                    output.push('\n');
                }
                indent(output, depth);
            }
            output.push(']');
        }
        Json::Object(entries) => {
            output.push('{');
            if !entries.is_empty() {
                output.push('\n');
                for (index, (key, value)) in entries.iter().enumerate() {
                    indent(output, depth + 1);
                    output.push_str(key);
                    output.push_str(": ");
                    write_pretty(value, output, depth + 1);
                    if index + 1 != entries.len() {
                        output.push(',');
                    }
                    output.push('\n');
                }
                indent(output, depth);
            }
            output.push('}');
        }
    }
}

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

fn pretty_print(input: &str) -> Result<String, String> {
    let value = Parser::new(input).parse()?;
    let mut output = String::new();
    write_pretty(&value, &mut output, 0);
    Ok(output)
}

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

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

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

    #[test]
    fn formats_nested_json() {
        let input = r#"{"name":"Ada","values":[1,true,null],"meta":{}}"#;
        let expected = r#"{
  "name": "Ada",
  "values": [
    1,
    true,
    null
  ],
  "meta": {}
}"#;

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

    #[test]
    fn preserves_string_escapes_and_numbers() {
        let input = r#"["line\nbreak","\u263a",-1.25e+3]"#;
        let expected = "[\n  \"line\\nbreak\",\n  \"\\u263a\",\n  -1.25e+3\n]";

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

    #[test]
    fn rejects_trailing_input() {
        assert!(pretty_print(r#"{"ok":true} garbage"#).is_err());
    }
}