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

/// A small recursive-descent JSON formatter.
struct Formatter<'a> {
    input: &'a [u8],
    pos: usize,
    output: String,
}

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

    fn format(mut self) -> Result<String, String> {
        self.skip_whitespace();
        self.value(0)?;
        self.skip_whitespace();

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

        self.output.push('\n');
        Ok(self.output)
    }

    fn value(&mut self, depth: usize) -> Result<(), String> {
        self.skip_whitespace();

        match self.peek() {
            Some(b'{') => self.object(depth),
            Some(b'[') => self.array(depth),
            Some(b'"') => {
                let string = self.string()?;
                self.output.push_str(string);
                Ok(())
            }
            Some(b't') => self.literal(b"true"),
            Some(b'f') => self.literal(b"false"),
            Some(b'n') => self.literal(b"null"),
            Some(b'-' | b'0'..=b'9') => {
                let number = self.number()?;
                self.output.push_str(number);
                Ok(())
            }
            Some(_) => self.error("expected a JSON value"),
            None => self.error("unexpected end of input"),
        }
    }

    fn object(&mut self, depth: usize) -> Result<(), String> {
        self.consume(b'{')?;
        self.output.push('{');
        self.skip_whitespace();

        if self.take(b'}') {
            self.output.push('}');
            return Ok(());
        }

        self.output.push('\n');

        loop {
            self.indent(depth + 1);
            let key = self.string()?;
            self.output.push_str(key);
            self.skip_whitespace();
            self.consume(b':')?;
            self.output.push_str(": ");
            self.value(depth + 1)?;
            self.skip_whitespace();

            if self.take(b',') {
                self.output.push_str(",\n");
                self.skip_whitespace();
            } else {
                self.consume(b'}')?;
                self.output.push('\n');
                self.indent(depth);
                self.output.push('}');
                return Ok(());
            }
        }
    }

    fn array(&mut self, depth: usize) -> Result<(), String> {
        self.consume(b'[')?;
        self.output.push('[');
        self.skip_whitespace();

        if self.take(b']') {
            self.output.push(']');
            return Ok(());
        }

        self.output.push('\n');

        loop {
            self.indent(depth + 1);
            self.value(depth + 1)?;
            self.skip_whitespace();

            if self.take(b',') {
                self.output.push_str(",\n");
                self.skip_whitespace();
            } else {
                self.consume(b']')?;
                self.output.push('\n');
                self.indent(depth);
                self.output.push(']');
                return Ok(());
            }
        }
    }

    /// Validates a JSON string while retaining its original escape spelling.
    fn string(&mut self) -> Result<&'a str, String> {
        let start = self.pos;
        self.consume(b'"')?;

        loop {
            match self.next() {
                Some(b'"') => {
                    return std::str::from_utf8(&self.input[start..self.pos])
                        .map_err(|_| format!("invalid UTF-8 at byte {}", start));
                }
                Some(b'\\') => match self.next() {
                    Some(b'"' | b'\\' | b'/' | b'b' | b'f' | b'n' | b'r' | b't') => {}
                    Some(b'u') => {
                        for _ in 0..4 {
                            match self.next() {
                                Some(b'0'..=b'9' | b'a'..=b'f' | b'A'..=b'F') => {}
                                _ => return self.error("invalid Unicode escape"),
                            }
                        }
                    }
                    _ => return self.error("invalid string escape"),
                },
                Some(0x00..=0x1f) => {
                    return self.error("unescaped control character in string")
                }
                Some(_) => {}
                None => return self.error("unterminated string"),
            }
        }
    }

    fn number(&mut self) -> Result<&'a str, String> {
        let start = self.pos;
        self.take(b'-');

        match self.peek() {
            Some(b'0') => {
                self.pos += 1;
                if matches!(self.peek(), Some(b'0'..=b'9')) {
                    return 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 self.error("invalid number"),
        }

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

        if matches!(self.peek(), Some(b'e' | b'E')) {
            self.pos += 1;
            if matches!(self.peek(), Some(b'+' | b'-')) {
                self.pos += 1;
            }
            if !matches!(self.peek(), Some(b'0'..=b'9')) {
                return self.error("expected digit in exponent");
            }
            while matches!(self.peek(), Some(b'0'..=b'9')) {
                self.pos += 1;
            }
        }

        std::str::from_utf8(&self.input[start..self.pos])
            .map_err(|_| format!("invalid number at byte {}", start))
    }

    fn literal(&mut self, expected: &[u8]) -> Result<(), String> {
        if self.input.get(self.pos..self.pos + expected.len()) == Some(expected) {
            self.pos += expected.len();
            self.output
                .push_str(std::str::from_utf8(expected).expect("ASCII literal"));
            Ok(())
        } else {
            self.error("invalid literal")
        }
    }

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

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

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

    fn next(&mut self) -> Option<u8> {
        let byte = self.peek()?;
        self.pos += 1;
        Some(byte)
    }

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

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

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

fn pretty_print(input: &str) -> Result<String, String> {
    Formatter::new(input).format()
}

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) => print!("{output}"),
        Err(error) => {
            eprintln!("invalid JSON: {error}");
            std::process::exit(1);
        }
    }
}

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

    #[test]
    fn formats_nested_json() {
        let input = r#"{"name":"Ada","values":[1,true,null,{"x":2}]}"#;
        let expected = concat!(
            "{\n",
            "  \"name\": \"Ada\",\n",
            "  \"values\": [\n",
            "    1,\n",
            "    true,\n",
            "    null,\n",
            "    {\n",
            "      \"x\": 2\n",
            "    }\n",
            "  ]\n",
            "}\n"
        );
        assert_eq!(pretty_print(input).unwrap(), expected);
    }

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