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

/// Pretty-prints valid JSON using two-space indentation.
pub fn pretty_print(input: &str) -> Result<String, String> {
    let mut parser = Parser {
        input: input.as_bytes(),
        pos: 0,
        output: String::new(),
    };

    parser.skip_whitespace();
    parser.parse_value(0)?;
    parser.skip_whitespace();

    if parser.pos != parser.input.len() {
        return Err(format!("unexpected character at byte {}", parser.pos));
    }

    Ok(parser.output)
}

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

impl Parser<'_> {
    fn parse_value(&mut self, depth: usize) -> Result<(), String> {
        self.skip_whitespace();

        match self.peek() {
            Some(b'{') => self.parse_object(depth),
            Some(b'[') => self.parse_array(depth),
            Some(b'"') => self.parse_string(),
            Some(b'-' | b'0'..=b'9') => self.parse_number(),
            Some(b't') => self.parse_literal(b"true"),
            Some(b'f') => self.parse_literal(b"false"),
            Some(b'n') => self.parse_literal(b"null"),
            Some(_) => Err(format!("expected JSON value at byte {}", self.pos)),
            None => Err("unexpected end of input".into()),
        }
    }

    fn parse_object(&mut self, depth: usize) -> Result<(), String> {
        self.pos += 1;
        self.output.push('{');
        self.skip_whitespace();

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

        let mut first = true;
        loop {
            if !first {
                self.expect(b',')?;
                self.output.push(',');
            }

            self.output.push('\n');
            self.indent(depth + 1);
            self.skip_whitespace();

            if self.peek() != Some(b'"') {
                return Err(format!("expected object key at byte {}", self.pos));
            }
            self.parse_string()?;

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

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

            first = false;
        }
    }

    fn parse_array(&mut self, depth: usize) -> Result<(), String> {
        self.pos += 1;
        self.output.push('[');
        self.skip_whitespace();

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

        let mut first = true;
        loop {
            if !first {
                self.expect(b',')?;
                self.output.push(',');
            }

            self.output.push('\n');
            self.indent(depth + 1);
            self.parse_value(depth + 1)?;
            self.skip_whitespace();

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

            first = false;
        }
    }

    // Strings are copied verbatim after validating JSON escape syntax.
    fn parse_string(&mut self) -> Result<(), String> {
        let start = self.pos;
        self.expect(b'"')?;

        while let Some(byte) = self.peek() {
            match byte {
                b'"' => {
                    self.pos += 1;
                    let text = std::str::from_utf8(&self.input[start..self.pos])
                        .map_err(|_| format!("invalid UTF-8 at byte {}", start))?;
                    self.output.push_str(text);
                    return Ok(());
                }
                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(format!(
                                            "invalid Unicode escape at byte {}",
                                            self.pos
                                        ));
                                    }
                                }
                            }
                        }
                        _ => return Err(format!("invalid escape at byte {}", self.pos)),
                    }
                }
                0x00..=0x1f => {
                    return Err(format!("control character in string at byte {}", self.pos));
                }
                _ => self.pos += 1,
            }
        }

        Err("unterminated string".into())
    }

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

        self.consume(b'-');

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

        if self.consume(b'.') {
            let fraction = self.pos;
            while matches!(self.peek(), Some(b'0'..=b'9')) {
                self.pos += 1;
            }
            if self.pos == fraction {
                return Err(format!("expected digit at byte {}", self.pos));
            }
        }

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

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

        let number = std::str::from_utf8(&self.input[start..self.pos]).unwrap();
        self.output.push_str(number);
        Ok(())
    }

    fn parse_literal(&mut self, literal: &[u8]) -> Result<(), String> {
        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).unwrap());
            Ok(())
        } else {
            Err(format!("invalid literal at byte {}", self.pos))
        }
    }

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

    fn expect(&mut self, expected: u8) -> Result<(), String> {
        self.skip_whitespace();
        if self.consume(expected) {
            Ok(())
        } else {
            Err(format!(
                "expected '{}' at byte {}",
                expected as char, self.pos
            ))
        }
    }

    fn consume(&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 indent(&mut self, depth: usize) {
        for _ in 0..depth {
            self.output.push_str("  ");
        }
    }
}

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::pretty_print;

    #[test]
    fn formats_nested_json() {
        let input = r#"{"name":"Ada","values":[1,true,null,{"x":2}]}"#;
        let expected = r#"{
  "name": "Ada",
  "values": [
    1,
    true,
    null,
    {
      "x": 2
    }
  ]
}"#;
        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());
    }
}