use serde_json::Value as JsonValue;

use crate::syntax::{self, Value, ValueKind};
use crate::{Error, Options, Span, format};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PathSegment<'a> {
    Key(&'a str),
    Index(usize),
}

impl<'a> From<&'a str> for PathSegment<'a> {
    fn from(value: &'a str) -> Self {
        Self::Key(value)
    }
}

impl From<usize> for PathSegment<'_> {
    fn from(value: usize) -> Self {
        Self::Index(value)
    }
}

#[derive(Debug, Clone)]
pub struct Document {
    source: String,
    root: Value,
    options: Options,
}

impl Document {
    pub fn parse(source: impl Into<String>) -> Result<Self, Error> {
        Self::parse_with(source, &Options::default())
    }

    pub fn parse_with(source: impl Into<String>, options: &Options) -> Result<Self, Error> {
        let source = source.into();
        let root = syntax::parse_with(&source, options)?;
        Ok(Self { source, root, options: *options })
    }

    pub fn parse_object(source: impl Into<String>) -> Result<Self, Error> {
        Self::parse_object_with(source, &Options::default())
    }

    pub fn parse_object_with(source: impl Into<String>, options: &Options) -> Result<Self, Error> {
        let source = source.into();
        let root = syntax::parse_object_with(&source, options)?;
        Ok(Self { source, root, options: *options })
    }

    pub fn options(&self) -> Options {
        self.options
    }

    pub fn source(&self) -> &str {
        &self.source
    }

    pub fn into_source(self) -> String {
        self.source
    }

    pub fn root(&self) -> &Value {
        &self.root
    }

    pub fn to_value(&self) -> Result<JsonValue, Error> {
        self.root.to_value()
    }

    /// Missing object keys are inserted; an array index equal to its length appends.
    /// Parent containers must exist. An empty path replaces the root.
    pub fn set(&mut self, path: &[PathSegment<'_>], value: &JsonValue) -> Result<(), Error> {
        let edit = if let Some((last, parent)) = path.split_last() {
            let container = self.locate(parent)?;
            if let Some(existing) = child(container, *last)? {
                self.replacement(existing, value)?
            } else {
                let prefix = match (&container.kind, last) {
                    (ValueKind::Object(_), PathSegment::Key(key)) => {
                        format!("{}: ", format::key(key))
                    }
                    (ValueKind::Array(items), PathSegment::Index(index))
                        if *index == items.len() =>
                    {
                        String::new()
                    }
                    _ => {
                        return Err(syntax::invalid(
                            container.span.clone(),
                            "Array index out of bounds",
                        ));
                    }
                };
                Some(vec![self.insertion(container, &prefix, value)?])
            }
        } else {
            self.replacement(&self.root, value)?
        };
        if let Some(edits) = edit {
            self.apply(edits)?;
        }
        Ok(())
    }

    /// Comments outside the value survive, including those between a key and its value.
    /// Removing the root is not supported; replace it with an empty object or null instead.
    pub fn remove(&mut self, path: &[PathSegment<'_>]) -> Result<bool, Error> {
        let Some((last, parent)) = path.split_last() else {
            return Err(syntax::invalid(
                self.root.span.clone(),
                "Cannot remove the root",
            ));
        };
        let container = self.locate(parent)?;
        let Some(value) = child(container, *last)? else {
            return Ok(false);
        };
        let mut edits = vec![(value.span.clone(), String::new())];
        let mut member_start = value.span.start;
        let mut colon_at = None;
        let mut key_end = None;
        let next = match (&container.kind, last) {
            (ValueKind::Object(members), PathSegment::Key(key)) => {
                let index = members
                    .iter()
                    .position(|member| member.key == *key)
                    .unwrap();
                let member = &members[index];
                member_start = member.key_span.start;
                key_end = Some(member.key_span.end);
                edits.push((member.key_span.clone(), String::new()));
                let colon =
                    punctuation(&self.source, member.key_span.end..value.span.start, b':').unwrap();
                colon_at = Some(colon);
                edits.push((colon..colon + 1, String::new()));
                members.get(index + 1).map(|member| member.key_span.start)
            }
            (ValueKind::Array(items), PathSegment::Index(index)) => {
                items.get(index + 1).map(|value| value.span.start)
            }
            _ => unreachable!(),
        };
        let comma = punctuation(
            &self.source,
            value.span.end..next.unwrap_or_else(|| self.container_end(container)),
            b',',
        );
        if let Some(comma) = comma {
            edits.push((comma..comma + 1, String::new()));
        }
        // A member alone on its line takes its indentation and line break
        // with it, rather than leaving a whitespace-only line behind.
        // Comments or other content on the line block both extensions.
        let start_of_line = line_start(&self.source, member_start);
        if start_of_line < member_start && is_blank(&self.source, start_of_line..member_start) {
            edits.push((start_of_line..member_start, String::new()));
        }
        let member_end = comma.map_or(value.span.end, |comma| comma + 1);
        let bytes = self.source.as_bytes();
        let mut end = member_end;
        while end < bytes.len() && matches!(bytes[end], b' ' | b'\t') {
            end += 1;
        }
        if self.source[end..].starts_with("\r\n") {
            end += 2;
        } else if self.source[end..].starts_with(['\r', '\n']) {
            end += 1;
        } else {
            end = member_end;
        }
        if end > member_end {
            edits.push((member_end..end, String::new()));
        }
        // Blank gaps around the colon and comma (e.g. the space in `a: 1`)
        // are not part of any span above and would otherwise survive on the
        // collapsed line. Gaps containing comments are not blank and stay.
        if let (Some(colon), Some(key_end)) = (colon_at, key_end) {
            let mut gaps = vec![key_end..colon, colon + 1..value.span.start];
            if let Some(comma) = comma {
                gaps.push(value.span.end..comma);
            }
            for gap in gaps {
                if gap.start < gap.end && is_blank(&self.source, gap.clone()) {
                    edits.push((gap, String::new()));
                }
            }
        }
        self.apply(edits)?;
        Ok(true)
    }

    fn locate(&self, path: &[PathSegment<'_>]) -> Result<&Value, Error> {
        let mut value = &self.root;
        for segment in path {
            value = child(value, *segment)?
                .ok_or_else(|| syntax::invalid(value.span.clone(), "Path does not exist"))?;
        }
        Ok(value)
    }

    fn replacement(
        &self,
        existing: &Value,
        value: &JsonValue,
    ) -> Result<Option<Vec<(Span, String)>>, Error> {
        if existing.to_value().ok().as_ref() == Some(value) {
            return Ok(None);
        }
        let implicit = matches!(existing.kind, ValueKind::Object(_))
            && !self.source[existing.span.start..].starts_with('{');
        let indent = line_indent(&self.source, existing.span.start);
        let mut text = if let JsonValue::String(text) = value
            && self.source[existing.span.start..].starts_with('\'')
        {
            let json = serde_json::to_string(text).unwrap();
            format!("'{}'", json[1..json.len() - 1].replace('\'', "\\'"))
        } else {
            formatted(
                value,
                implicit,
                indent,
                newline(&self.source),
                &self.options,
            )?
        };
        if existing.span.is_empty()
            && existing.span.start > 0
            && !self.source[..existing.span.start].ends_with(['\r', '\n'])
        {
            text = format!("{}{indent}{text}", newline(&self.source));
        }
        let mut edits = vec![(existing.span.clone(), text)];
        if implicit
            && let Some(comma) =
                punctuation(&self.source, existing.span.end..self.source.len(), b',')
        {
            edits.push((comma..comma + 1, String::new()));
        }
        Ok(Some(edits))
    }

    fn container_end(&self, container: &Value) -> usize {
        if matches!(container.kind, ValueKind::Object(_))
            && !self.source[container.span.start..].starts_with('{')
        {
            self.source.len()
        } else {
            container.span.end
        }
    }

    fn insertion(
        &self,
        container: &Value,
        prefix: &str,
        value: &JsonValue,
    ) -> Result<(Span, String), Error> {
        let braced = self.source[container.span.start..].starts_with(['{', '[']);
        let end = self.container_end(container) - usize::from(braced);
        let base = line_indent(&self.source, container.span.start);
        let first = match &container.kind {
            ValueKind::Object(members) => members.first().map(|member| member.key_span.start),
            ValueKind::Array(items) => items.first().map(|value| value.span.start),
            _ => unreachable!(),
        };
        let indent = first
            .filter(|start| {
                self.source[line_start(&self.source, *start)..*start]
                    .chars()
                    .all(|ch| ch == ' ' || ch == '\t')
            })
            .map(|start| line_indent(&self.source, start).to_owned())
            .unwrap_or_else(|| {
                if braced {
                    format!("{base}  ")
                } else {
                    base.to_owned()
                }
            });
        let line = line_start(&self.source, end);
        let at_line_start = self.source[line..end]
            .chars()
            .all(|ch| ch == ' ' || ch == '\t');
        let at = if braced && at_line_start { line } else { end };
        let newline = newline(&self.source);
        let mut text = String::new();
        if at > 0 && !self.source[..at].ends_with(['\r', '\n']) {
            text.push_str(newline);
        }
        text.push_str(&indent);
        text.push_str(prefix);
        text.push_str(&formatted(value, false, &indent, newline, &self.options)?);
        if braced {
            text.push_str(newline);
            if !at_line_start {
                text.push_str(base);
            }
        } else if self.source.ends_with(['\r', '\n']) {
            text.push_str(newline);
        }
        Ok((at..at, text))
    }

    fn apply(&mut self, mut edits: Vec<(Span, String)>) -> Result<(), Error> {
        let mut source = self.source.clone();
        edits.sort_by_key(|(span, _)| std::cmp::Reverse(span.start));
        for (span, text) in edits {
            source.replace_range(span, &text);
        }
        let root = syntax::parse_with(&source, &self.options).map_err(|error| Error {
            kind: error.kind,
            span: self.root.span.clone(),
            message: format!("Edit would produce invalid fJSON: {}", error.message),
            previous: None,
        })?;
        self.source = source;
        self.root = root;
        Ok(())
    }
}

fn child<'a>(value: &'a Value, segment: PathSegment<'_>) -> Result<Option<&'a Value>, Error> {
    match (&value.kind, segment) {
        (ValueKind::Object(members), PathSegment::Key(key)) => Ok(members
            .iter()
            .find(|member| member.key == key)
            .map(|member| &member.value)),
        (ValueKind::Array(items), PathSegment::Index(index)) => Ok(items.get(index)),
        _ => Err(syntax::invalid(
            value.span.clone(),
            "Path segment does not match its container",
        )),
    }
}

fn formatted(
    value: &JsonValue,
    implicit: bool,
    indent: &str,
    newline: &str,
    options: &Options,
) -> Result<String, Error> {
    Ok(format::print(value, implicit, options)?
        .split('\n')
        .collect::<Vec<_>>()
        .join(&format!("{newline}{indent}")))
}

fn is_blank(source: &str, span: Span) -> bool {
    source[span]
        .bytes()
        .all(|byte| matches!(byte, b' ' | b'\t'))
}

fn line_start(source: &str, at: usize) -> usize {
    source[..at]
        .rfind(['\r', '\n'])
        .map_or(0, |index| index + 1)
}

fn line_indent(source: &str, at: usize) -> &str {
    let start = line_start(source, at);
    let end = source[start..at]
        .find(|ch| ch != ' ' && ch != '\t')
        .unwrap_or(at - start);
    &source[start..start + end]
}

fn newline(source: &str) -> &'static str {
    if source.contains("\r\n") {
        "\r\n"
    } else if source.contains('\r') {
        "\r"
    } else {
        "\n"
    }
}

fn punctuation(source: &str, span: Span, target: u8) -> Option<usize> {
    let bytes = source.as_bytes();
    let mut at = span.start;
    while at < span.end {
        if source[at..].starts_with("//") {
            at += source[at..span.end]
                .find(['\r', '\n'])
                .unwrap_or(span.end - at);
        } else if source[at..].starts_with("/*") {
            at += source[at..span.end].find("*/").unwrap() + 2;
        } else if bytes[at] == target {
            return Some(at);
        } else {
            at += 1;
        }
    }
    None
}