//! JSON view over the Automerge document with conflict-aware merging.

use anyhow::anyhow;
use automerge::{
    transaction::Transactable, ActorId, AutoCommit, ObjId, ObjType, ReadDoc, ScalarValue, Value,
};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use serde::Serialize;
use serde_json::Value as JsonValue;
use std::collections::BTreeSet;

#[derive(Debug, thiserror::Error)]
pub enum StoreError {
    #[error("path crosses a non-map value")]
    Conflict,
    #[error("value does not exist")]
    NotFound,
    #[error(transparent)]
    Doc(#[from] automerge::AutomergeError),
    #[error(transparent)]
    Other(#[from] anyhow::Error),
}

/// Automerge's total order on operations: Lamport counter, then actor.
fn op_order(id: &ObjId) -> (u64, Option<&ActorId>) {
    match id {
        ObjId::Root => (0, None),
        ObjId::Id(counter, actor, _) => (*counter, Some(actor)),
    }
}

#[derive(Debug, Clone)]
enum Resolved {
    /// Concurrent maps presented as one. Newest first; `[0]` is canonical for writes.
    Maps(Vec<ObjId>),
    /// Any other value: the newest concurrent candidate.
    Other(Value<'static>, ObjId),
}

/// All concurrent values of `key` across a merged set of parent maps, newest first.
fn candidates(
    doc: &AutoCommit,
    parents: &[ObjId],
    key: &str,
) -> Result<Vec<(Value<'static>, ObjId)>, StoreError> {
    let mut all = Vec::new();
    for parent in parents {
        for (value, id) in doc.get_all(parent, key)? {
            all.push((value.into_owned(), id));
        }
    }
    all.sort_by(|a, b| op_order(&b.1).cmp(&op_order(&a.1)));
    Ok(all)
}

fn resolve(
    doc: &AutoCommit,
    parents: &[ObjId],
    key: &str,
    log: Option<&mut ConflictLog>,
) -> Result<Option<Resolved>, StoreError> {
    let mut all = candidates(doc, parents, key)?;
    if all.is_empty() {
        return Ok(None);
    }

    let merged = matches!(all[0].0, Value::Object(ObjType::Map));

    if all.len() > 1 {
        if let Some(log) = log {
            log.record(doc, key, &all, merged)?;
        }
    }

    if merged {
        let maps = all
            .into_iter()
            .filter(|(value, _)| matches!(value, Value::Object(ObjType::Map)))
            .map(|(_, id)| id)
            .collect();
        Ok(Some(Resolved::Maps(maps)))
    } else {
        let (value, id) = all.swap_remove(0);
        Ok(Some(Resolved::Other(value, id)))
    }
}

fn merged_keys(doc: &AutoCommit, maps: &[ObjId]) -> BTreeSet<String> {
    maps.iter().flat_map(|map| doc.keys(map)).collect()
}

fn resolved_to_json(
    doc: &AutoCommit,
    resolved: &Resolved,
    mut log: Option<&mut ConflictLog>,
) -> Result<JsonValue, StoreError> {
    let maps = match resolved {
        Resolved::Maps(maps) => maps,
        Resolved::Other(Value::Scalar(scalar), _) => {
            return scalar_to_json(scalar.clone().into_owned())
        }
        Resolved::Other(Value::Object(_), id) => return object_to_json(doc, id, log),
    };

    let mut out = serde_json::Map::new();

    if let [map] = maps.as_slice() {
        // Fast path: one pass; get_all only for keys Automerge flags as conflicted.
        for item in doc.map_range(map, ..) {
            let key = item.key.to_string();

            let child = if item.conflict {
                match resolve(doc, maps, &key, log.as_deref_mut())? {
                    Some(child) => child,
                    None => continue,
                }
            } else {
                let id = item.id();
                match Value::from(item.value) {
                    Value::Object(ObjType::Map) => Resolved::Maps(vec![id]),
                    value => Resolved::Other(value.into_owned(), id),
                }
            };

            let json = nested(&mut log, &key, |log| resolved_to_json(doc, &child, log))?;
            out.insert(key, json);
        }
    } else {
        // Several concurrent maps: union of their keys.
        for key in merged_keys(doc, maps) {
            let Some(child) = resolve(doc, maps, &key, log.as_deref_mut())? else {
                continue;
            };
            let json = nested(&mut log, &key, |log| resolved_to_json(doc, &child, log))?;
            out.insert(key, json);
        }
    }

    Ok(JsonValue::Object(out))
}

fn object_to_json(
    doc: &AutoCommit,
    id: &ObjId,
    mut log: Option<&mut ConflictLog>,
) -> Result<JsonValue, StoreError> {
    match doc.object_type(id)? {
        ObjType::Map => resolved_to_json(doc, &Resolved::Maps(vec![id.clone()]), log),
        ObjType::List => {
            let mut out = Vec::new();
            for (index, item) in doc.list_range(id, ..).enumerate() {
                let child_id = item.id();
                out.push(match Value::from(item.value) {
                    Value::Object(_) => nested(&mut log, &index.to_string(), |log| {
                        object_to_json(doc, &child_id, log)
                    })?,
                    Value::Scalar(scalar) => scalar_to_json(scalar.into_owned())?,
                });
            }
            Ok(JsonValue::Array(out))
        }
        ObjType::Text => Ok(JsonValue::String(doc.text(id)?)),
        ObjType::Table => Err(anyhow!("Automerge tables are not supported").into()),
    }
}

fn delete_in_all(doc: &mut AutoCommit, maps: &[ObjId], key: &str) -> Result<(), StoreError> {
    for map in maps {
        if doc.get(map, key)?.is_some() {
            doc.delete(map, key)?;
        }
    }
    Ok(())
}

fn read_path_inner(
    doc: &AutoCommit,
    path: &[String],
    mut log: Option<&mut ConflictLog>,
) -> Result<Option<JsonValue>, StoreError> {
    let mut current = Resolved::Maps(vec![automerge::ROOT]);

    for (index, segment) in path.iter().enumerate() {
        let Resolved::Maps(maps) = &current else {
            return Ok(None);
        };
        if let Some(log) = log.as_deref_mut() {
            // The last segment is the requested value itself, not an ancestor.
            log.in_ancestors = index + 1 < path.len();
        }
        match resolve(doc, maps, segment, log.as_deref_mut())? {
            Some(next) => current = next,
            None => return Ok(None),
        }
        if let Some(log) = log.as_deref_mut() {
            log.path.push(segment.clone());
        }
    }

    if let Some(log) = log.as_deref_mut() {
        log.in_ancestors = false;
    }
    Ok(Some(resolved_to_json(doc, &current, log)?))
}

/// Reads a path; the empty path is the whole store.
pub fn read_path(doc: &AutoCommit, path: &[String]) -> Result<Option<JsonValue>, StoreError> {
    read_path_inner(doc, path, None)
}

pub fn read_path_with_conflicts(
    doc: &AutoCommit,
    path: &[String],
) -> Result<Option<(JsonValue, Vec<ConflictReport>)>, StoreError> {
    let mut log = ConflictLog {
        want_values: true,
        ..Default::default()
    };
    let value = read_path_inner(doc, path, Some(&mut log))?;
    Ok(value.map(|value| (value, log.reports)))
}

/// For streams: the value, plus whether the path or subtree involves any conflict.
/// Reports are collected without values, so this costs little more than read_path.
pub fn read_path_traced(
    doc: &AutoCommit,
    path: &[String],
) -> Result<(Option<JsonValue>, bool), StoreError> {
    let mut log = ConflictLog::default();
    let value = read_path_inner(doc, path, Some(&mut log))?;
    Ok((value, !log.reports.is_empty()))
}

/// Keys of the merged map at `path`; empty if it doesn't exist or isn't a map.
pub fn keys_at(doc: &AutoCommit, path: &[String]) -> Result<Vec<String>, StoreError> {
    let mut maps = vec![automerge::ROOT];
    for segment in path {
        maps = match resolve(doc, &maps, segment, None)? {
            Some(Resolved::Maps(next)) => next,
            _ => return Ok(Vec::new()),
        };
    }
    Ok(merged_keys(doc, &maps).into_iter().collect())
}

/// Writes (replaces) the value at `path`, creating intermediate maps.
pub fn write_path(
    doc: &mut AutoCommit,
    path: &[String],
    value: &JsonValue,
) -> Result<(), StoreError> {
    let (leaf, parents) = path.split_last().ok_or(StoreError::NotFound)?;

    let mut maps = vec![automerge::ROOT];
    for segment in parents {
        maps = match resolve(doc, &maps, segment, None)? {
            Some(Resolved::Maps(next)) => next,
            Some(Resolved::Other(..)) => return Err(StoreError::Conflict),
            // A concurrent creation elsewhere is merged on read.
            None => vec![doc.put_object(&maps[0], segment.as_str(), ObjType::Map)?],
        };
    }

    // Replace semantics: clear the key in every merged parent, or an old map
    // would merge into the new value.
    delete_in_all(doc, &maps, leaf)?;
    json_to_automerge(doc, &maps[0], leaf, value)?;
    Ok(())
}

/// Deletes the value at `path`, optionally pruning ancestors that become empty.
pub fn delete_path(doc: &mut AutoCommit, path: &[String], prune: bool) -> Result<(), StoreError> {
    let (leaf, parents) = path.split_last().ok_or(StoreError::NotFound)?;

    // levels[d] = merged maps at path[..d]
    let mut levels = vec![vec![automerge::ROOT]];
    for segment in parents {
        let next = match resolve(doc, levels.last().expect("non-empty"), segment, None)? {
            Some(Resolved::Maps(maps)) => maps,
            _ => return Err(StoreError::NotFound),
        };
        levels.push(next);
    }

    let maps = levels.last().expect("non-empty");
    if resolve(doc, maps, leaf, None)?.is_none() {
        return Err(StoreError::NotFound);
    }
    delete_in_all(doc, maps, leaf)?;

    if prune {
        // Never prune the root (depth 0).
        for depth in (1..levels.len()).rev() {
            if !merged_keys(doc, &levels[depth]).is_empty() {
                break;
            }
            delete_in_all(doc, &levels[depth - 1], &path[depth - 1])?;
        }
    }

    Ok(())
}

/// Where a value is placed inside its parent object.
#[derive(Clone, Copy)]
enum Slot<'a> {
    /// A key in a map.
    Key(&'a str),
    /// An index in a list. The value is inserted here, not overwritten.
    Index(usize),
}

fn put_object(
    doc: &mut AutoCommit,
    parent: &ObjId,
    slot: Slot,
    obj_type: ObjType,
) -> Result<ObjId, StoreError> {
    Ok(match slot {
        Slot::Key(key) => doc.put_object(parent, key, obj_type)?,
        Slot::Index(index) => doc.insert_object(parent, index, obj_type)?,
    })
}

fn put_scalar(
    doc: &mut AutoCommit,
    parent: &ObjId,
    slot: Slot,
    value: impl Into<ScalarValue>,
) -> Result<(), StoreError> {
    match slot {
        Slot::Key(key) => doc.put(parent, key, value)?,
        Slot::Index(index) => doc.insert(parent, index, value)?,
    }
    Ok(())
}

fn write_value(
    doc: &mut AutoCommit,
    parent: &ObjId,
    slot: Slot,
    value: &JsonValue,
) -> Result<(), StoreError> {
    match value {
        JsonValue::Object(object) => {
            let id = put_object(doc, parent, slot, ObjType::Map)?;
            for (child_key, child_value) in object {
                write_value(doc, &id, Slot::Key(child_key), child_value)?;
            }
        }
        JsonValue::Array(items) => {
            let id = put_object(doc, parent, slot, ObjType::List)?;
            for (index, item) in items.iter().enumerate() {
                write_value(doc, &id, Slot::Index(index), item)?;
            }
        }
        JsonValue::String(value) => put_scalar(doc, parent, slot, value.as_str())?,
        JsonValue::Bool(value) => put_scalar(doc, parent, slot, *value)?,
        JsonValue::Number(value) => {
            if let Some(value) = value.as_i64() {
                put_scalar(doc, parent, slot, value)?;
            } else if let Some(value) = value.as_u64() {
                put_scalar(doc, parent, slot, value)?;
            } else if let Some(value) = value.as_f64() {
                put_scalar(doc, parent, slot, value)?;
            } else {
                return Err(anyhow!("Unsupported JSON number").into());
            }
        }
        JsonValue::Null => put_scalar(doc, parent, slot, ScalarValue::Null)?,
    }

    Ok(())
}

fn json_to_automerge(
    doc: &mut AutoCommit,
    parent: &ObjId,
    key: &str,
    value: &JsonValue,
) -> Result<(), StoreError> {
    write_value(doc, parent, Slot::Key(key), value)
}

fn scalar_to_json(scalar: ScalarValue) -> Result<JsonValue, StoreError> {
    match scalar {
        ScalarValue::Null => Ok(JsonValue::Null),
        ScalarValue::Boolean(b) => Ok(JsonValue::Bool(b)),
        ScalarValue::Int(i) => Ok(JsonValue::Number(i.into())),
        ScalarValue::Uint(u) => Ok(JsonValue::Number(u.into())),
        ScalarValue::F64(f) => serde_json::Number::from_f64(f)
            .map(JsonValue::Number)
            .ok_or(anyhow!("Non-finite float cannot be represented as JSON").into()),
        ScalarValue::Str(s) => Ok(JsonValue::String(s.to_string())),
        ScalarValue::Bytes(b) => Ok(JsonValue::String(BASE64.encode(b))),
        ScalarValue::Counter(c) => Ok(JsonValue::Number(i64::from(c).into())),
        ScalarValue::Timestamp(t) => Ok(JsonValue::Number(t.into())),
        ScalarValue::Unknown { .. } => Err(anyhow!("Unsupported Automerge scalar type").into()),
    }
}

// --- conflict log ------------------------------------------------------------

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub enum Resolution {
    MergedMaps,
    NewestWins,
}

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CandidateReport {
    pub op_id: String,
    pub selected: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub value: Option<JsonValue>,
}

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ConflictReport {
    pub path: Vec<String>,
    /// On the way to the requested path; reported without candidate values.
    #[serde(skip_serializing_if = "std::ops::Not::not")]
    pub ancestor: bool,
    pub resolution: Resolution,
    pub candidates: Vec<CandidateReport>,
}

#[derive(Default)]
struct ConflictLog {
    path: Vec<String>,
    reports: Vec<ConflictReport>,
    /// Render candidate values (only within the requested subtree).
    want_values: bool,
    /// Currently walking the ancestors of the requested path.
    in_ancestors: bool,
}

fn op_id_string(id: &ObjId) -> String {
    match id {
        ObjId::Root => "_root".to_owned(),
        ObjId::Id(counter, actor, _) => format!("{counter}@{actor}"),
    }
}

fn value_to_json(doc: &AutoCommit, value: &Value<'_>, id: &ObjId) -> Result<JsonValue, StoreError> {
    match value {
        Value::Scalar(scalar) => Ok(scalar_to_json(scalar.clone().into_owned())?),
        Value::Object(_) => object_to_json(doc, id, None),
    }
}

impl ConflictLog {
    fn record(
        &mut self,
        doc: &AutoCommit,
        key: &str,
        all: &[(Value<'static>, ObjId)],
        merged: bool,
    ) -> Result<(), StoreError> {
        let with_values = self.want_values && !self.in_ancestors;

        let candidates = all
            .iter()
            .enumerate()
            .map(|(index, (value, id))| {
                let selected = if merged {
                    matches!(value, Value::Object(ObjType::Map))
                } else {
                    index == 0
                };
                Ok(CandidateReport {
                    op_id: op_id_string(id),
                    selected,
                    value: if with_values {
                        Some(value_to_json(doc, value, id)?)
                    } else {
                        None
                    },
                })
            })
            .collect::<Result<Vec<_>, StoreError>>()?;

        let mut path = self.path.clone();
        path.push(key.to_owned());

        self.reports.push(ConflictReport {
            path,
            ancestor: self.in_ancestors,
            resolution: if merged {
                Resolution::MergedMaps
            } else {
                Resolution::NewestWins
            },
            candidates,
        });
        Ok(())
    }
}

/// Runs `f` with `segment` appended to the conflict path.
fn nested<T>(
    log: &mut Option<&mut ConflictLog>,
    segment: &str,
    f: impl FnOnce(Option<&mut ConflictLog>) -> T,
) -> T {
    if let Some(log) = log.as_deref_mut() {
        log.path.push(segment.to_owned());
    }
    let out = f(log.as_deref_mut());
    if let Some(log) = log.as_deref_mut() {
        log.path.pop();
    }
    out
}