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

#[derive(Debug)]
struct JsonError {
    message: String,
    offset: usize,
}

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

struct Printer<'a> {
    input: &'a [u8],
    pos: usize,
    output: String,
    indent: usize,
}

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

    fn error(&self, message: impl Into<String>) -> JsonError {
        JsonError {
            message: message.into(),
            offset: self.pos,
        }
    }

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

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

    fn write_indent(&mut self, level: usize) {
        self.output
            .extend(std::iter::repeat(' ').take(level * self.indent));
    }

    fn expect(&mut self, byte: u8) -> Result<(), JsonError> {
        if self.peek() == Some(byte) {
            self.pos += 1;
            Ok(())
        } else {
            Err(self.error(format!("expected '{}'", byte as char)))
        }
    }

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

        match self.peek() {
            Some(b'{') => self.object(level),
            Some(b'[') => self.array(level),
            Some(b'"') => self.string(),
            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') => self.number(),
            Some(_) => Err(self.error("expected a JSON value")),
            None => Err(self.error("unexpected end of input")),
        }
    }

    fn object(&mut self, level: usize) -> Result<(), JsonError> {
        self.expect(b'{')?;
        self.output.push('{');
        self.skip_whitespace();

        if self.peek() == Some(b'}') {
            self.pos += 1;
            self.output.push('}');
            return Ok(());
        }

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

        loop {
            self.write_indent(level + 1);
            self.skip_whitespace();

            if self.peek() != Some(b'"') {
                return Err(self.error("object keys must be strings"));
            }
            self.string()?;

            self.skip_whitespace();
            self.expect(b':')?;
            self.output.push_str(": ");
            self.value(level + 1)?;

            self.skip_whitespace();
            match self.peek() {
                Some(b',') => {
                    self.pos += 1;
                    self.output.push_str(",\n");
                }
                Some(b'}') => {
                    self.pos += 1;
                    self.output.push('\n');
                    self.write_indent(level);
                    self.output.push('}');
                    return Ok(());
                }
                _ => return Err(self.error("expected ',' or '}'")),
            }
        }
    }

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

        if self.peek() == Some(b']') {
            self.pos += 1;
            self.output.push(']');
            return Ok(());
        }

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

        loop {
            self.write_indent(level + 1);
            self.value(level + 1)?;
            self.skip_whitespace();

            match self.peek() {
                Some(b',') => {
                    self.pos += 1;
                    self.output.push_str(",\n");
                }
                Some(b']') => {
                    self.pos += 1;
                    self.output.push('\n');
                    self.write_indent(level);
                    self.output.push(']');
                    return Ok(());
                }
                _ => return Err(self.error("expected ',' or ']'")),
            }
        }
    }

    // Copy strings exactly while checking escape syntax and control characters.
    fn string(&mut self) -> Result<(), JsonError> {
        let start = self.pos;
        self.expect(b'"')?;

        loop {
            match self.peek() {
                Some(b'"') => {
                    self.pos += 1;
                    self.output
                        .push_str(std::str::from_utf8(&self.input[start..self.pos])
                            .map_err(|_| self.error("string is not valid UTF-8"))?);
                    return Ok(());
                }
                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(b'0'..=b'9' | b'a'..=b'f' | b'A'..=b'F') => {
                                        self.pos += 1;
                                    }
                                    _ => return Err(self.error("invalid Unicode escape")),
                                }
                            }
                        }
                        _ => return Err(self.error("invalid string escape")),
                    }
                }
                Some(0x00..=0x1f) => {
                    return Err(self.error("unescaped control character in string"));
                }
                Some(_) => self.pos += 1,
                None => return Err(self.error("unterminated string")),
            }
        }
    }

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

    fn number(&mut self) -> Result<(), JsonError> {
        let start = self.pos;

        if self.peek() == Some(b'-') {
            self.pos += 1;
        }

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

        self.output.push_str(
            std::str::from_utf8(&self.input[start..self.pos])
                .expect("JSON numbers contain only ASCII"),
        );
        Ok(())
    }
}

fn pretty_print(input: &str) -> Result<String, JsonError> {
    let mut printer = Printer::new(input, 2);
    printer.value(0)?;
    printer.skip_whitespace();

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

    printer.output.push('\n');
    Ok(printer.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) => 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]}"#;
        let expected = concat!(
            "{\n",
            "  \"name\": \"Ada\",\n",
            "  \"values\": [\n",
            "    1,\n",
            "    true,\n",
            "    null\n",
            "  ]\n",
            "}\n"
        );

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

    #[test]
    fn preserves_escaped_string_contents() {
        let input = r#"{"text":"a\\b\n\u263a","empty":{}}"#;
        let output = pretty_print(input).unwrap();

        assert!(output.contains(r#""text": "a\\b\n\u263a""#));
        assert!(output.contains(r#""empty": {}"#));
    }

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