//! The document file (`--document-file`): this node's document on disk, so a restart
//! resumes with it instead of starting empty.
//!
//! Format: the line `continuum-document/1`, a header line (JSON: the actor IDs this
//! node wrote with, so a later rejoin replays those writes as its own), then records,
//! each a little-endian `u32` length and that many bytes: first a full save of the
//! document, then increments (the changes after the heads written before). A record cut
//! short by a crash is dropped when reading.
//!
//! Written by a worker (`swarm::document_file`), only while the node is ready and not
//! rejoining (a rejoin keeps the old file until it is done):
//! * every `--document-append`, what changed is appended;
//! * the whole file is rewritten, atomically, when the epoch or this node's actors
//!   changed (a switch, a rejoin, a restart), and when the appended increments outgrew
//!   the full save, at most every `--document-rewrite`;
//! * at shutdown, what changed is appended.

use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;

use super::{ActorId, ChangeHash, Document};
use anyhow::{Context, Result, anyhow, bail};
use serde::{Deserialize, Serialize};
use tokio::time::Instant;
use tracing::{debug, error, warn};

use crate::document::replica::Epoch;
use crate::files::{FileKind, Files};
use crate::locks::MutexExt;

const MAGIC: &[u8] = b"continuum-document/1\n";

#[derive(Debug, Default, Serialize, Deserialize)]
struct Header {
    /// Hex actor IDs.
    #[serde(default)]
    actors: Vec<String>,
}

/// A document file as read at start.
pub(crate) struct Loaded {
    /// Under a fresh actor.
    pub(crate) doc: Document,
    /// The actors this node wrote with before.
    pub(crate) actors: Vec<ActorId>,
}

/// Reads the document file; None if there is none. An error if it can't be read or
/// isn't one: the caller decides (`--force badDocumentFile`).
pub(crate) fn load(path: &Path, files: &Files) -> Result<Option<Loaded>> {
    let Some(bytes) = files
        .read(FileKind::Document, path)
        .with_context(|| format!("Reading the document file {}", path.display()))?
    else {
        return Ok(None);
    };
    let (loaded, dropped) =
        parse(&bytes).with_context(|| format!("The document file {}", path.display()))?;
    if dropped > 0 {
        warn!(
            path = %path.display(),
            bytes = dropped,
            "Dropped a record cut short at the end of the document file (a crash while appending)"
        );
    }
    Ok(Some(loaded))
}

/// The file's content, and how many bytes at its end were cut short.
fn parse(bytes: &[u8]) -> Result<(Loaded, usize)> {
    let rest = bytes
        .strip_prefix(MAGIC)
        .ok_or_else(|| anyhow!("not a continuum document file"))?;
    let line_end = rest
        .iter()
        .position(|byte| *byte == b'\n')
        .ok_or_else(|| anyhow!("no header"))?;
    let header: Header = serde_json::from_slice(&rest[..line_end]).context("header")?;
    let actors = header
        .actors
        .iter()
        .map(|hex| {
            ActorId::try_from(hex.as_str())
                .map_err(|_| anyhow!("invalid actor {hex:?} in the header"))
        })
        .collect::<Result<Vec<_>>>()?;

    let mut records = &rest[line_end + 1..];
    let mut doc: Option<Document> = None;
    while !records.is_empty() {
        let Some((length, after)) = records.split_first_chunk::<4>() else {
            break;
        };
        let length = u32::from_le_bytes(*length) as usize;
        if after.len() < length {
            break;
        }
        let (record, after) = after.split_at(length);
        match &mut doc {
            None => doc = Some(super::load(record).context("the full save")?),
            Some(doc) => {
                doc.load_incremental(record).context("an increment")?;
            }
        }
        records = after;
    }
    let Some(doc) = doc else {
        bail!("no full save");
    };
    let loaded = Loaded { doc, actors };
    Ok((loaded, records.len()))
}

fn record(bytes: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(bytes.len() + 4);
    out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
    out.extend_from_slice(bytes);
    out
}

/// A whole file: header and one full save.
fn whole(actors: &[ActorId], save: &[u8]) -> Result<Vec<u8>> {
    let header = Header {
        actors: actors.iter().map(ActorId::to_hex_string).collect(),
    };
    let mut out = MAGIC.to_vec();
    out.extend(serde_json::to_vec(&header)?);
    out.push(b'\n');
    out.extend(record(save));
    Ok(out)
}

/// The document file and what was last written to it.
pub(crate) struct DocumentFile {
    path: PathBuf,
    files: Arc<Files>,
    append_every: Duration,
    rewrite_every: Duration,
    // A std mutex: never held across an await.
    written: Mutex<Option<Written>>,
}

#[derive(Debug, Clone)]
struct Written {
    epoch: Epoch,
    /// This node's actors when the file was last rewritten (`RemakeState::own_actors`).
    actors: usize,
    heads: Vec<ChangeHash>,
    /// The full save.
    base_bytes: u64,
    /// Appended since.
    tail_bytes: u64,
    rewritten_at: Instant,
}

impl DocumentFile {
    pub(crate) fn new(
        path: PathBuf,
        files: Arc<Files>,
        append_every: Duration,
        rewrite_every: Duration,
    ) -> Self {
        Self {
            path,
            files,
            append_every,
            rewrite_every,
            written: Mutex::new(None),
        }
    }

    fn lock(&self) -> MutexGuard<'_, Option<Written>> {
        self.written.lock_or_recover()
    }
}

/// What the document file needs now (`DocumentFile::due`).
pub(crate) enum Due {
    /// Written whole: for this reason.
    Rewrite(&'static str),
    /// What changed since `since` appended, as long as the document is still of `epoch`
    /// (an increment of another epoch's document must never be appended).
    Append {
        epoch: Epoch,
        since: Vec<ChangeHash>,
    },
    Nothing,
}

impl DocumentFile {
    /// How often what changed is appended (`--document-append`).
    pub(crate) fn append_every(&self) -> Duration {
        self.append_every
    }

    /// What is due for a document of `epoch` at `heads`, written by this node's `actors`
    /// (see the module docs).
    /// How many of this node's actors the file holds; None before it was written.
    pub(crate) fn actors(&self) -> Option<usize> {
        self.lock().as_ref().map(|written| written.actors)
    }

    pub(crate) fn due(&self, epoch: Epoch, actors: usize, heads: &[ChangeHash]) -> Due {
        let written = self.lock().clone();
        let reason = match &written {
            None => Some("first write"),
            Some(w) if w.epoch != epoch => Some("new epoch"),
            Some(w) if w.actors != actors => Some("new actor"),
            Some(w)
                if w.tail_bytes > w.base_bytes
                    && w.rewritten_at.elapsed() >= self.rewrite_every =>
            {
                Some("increments outgrew the full save")
            }
            Some(_) => None,
        };
        match (reason, written) {
            (Some(reason), _) => Due::Rewrite(reason),
            (None, Some(written)) if written.heads != heads => Due::Append {
                epoch: written.epoch,
                since: written.heads,
            },
            (None, _) => Due::Nothing,
        }
    }

    /// Writes the whole file, atomically: `copy` (a copy of the document, of `epoch`)
    /// and this node's actors `own` that wrote in it.
    pub(crate) async fn rewrite(
        self: &Arc<Self>,
        epoch: Epoch,
        copy: Document,
        own: Vec<ActorId>,
        reason: &str,
    ) {
        let start = Instant::now();
        let own_count = own.len();
        let writer = Arc::clone(self);
        let written = tokio::task::spawn_blocking(move || -> Result<(Vec<ChangeHash>, u64)> {
            let heads = copy.get_heads();
            // Only the actors with changes in this document matter for a later rejoin;
            // the current one in any case.
            let present: HashSet<ActorId> = copy
                .get_changes(&[])
                .iter()
                .map(|change| change.actor_id().clone())
                .collect();
            let current = copy.get_actor().clone();
            let actors: Vec<ActorId> = own
                .into_iter()
                .filter(|actor| *actor == current || present.contains(actor))
                .collect();
            let save = super::save(&copy);
            let bytes = whole(&actors, &save)?;
            writer
                .files
                .write_atomic(FileKind::Document, &writer.path, &bytes)?;
            Ok((heads, save.len() as u64))
        })
        .await;
        match written {
            Ok(Ok((heads, base_bytes))) => {
                debug!(
                    reason,
                    bytes = base_bytes,
                    took_ms = start.elapsed().as_millis() as u64,
                    "Rewrote the document file"
                );
                *self.lock() = Some(Written {
                    epoch,
                    actors: own_count,
                    heads,
                    base_bytes,
                    tail_bytes: 0,
                    rewritten_at: Instant::now(),
                });
            }
            Ok(Err(error)) => {
                error!(path = %self.path.display(), error = %format!("{error:#}"), "Couldn't rewrite the document file");
            }
            Err(error) => error!(error = %error, "Document file rewrite task failed"),
        }
    }

    /// Appends `increment`: what changed from `since` (the heads written last) to
    /// `heads`.
    pub(crate) async fn append(
        self: &Arc<Self>,
        since: Vec<ChangeHash>,
        heads: Vec<ChangeHash>,
        increment: Vec<u8>,
    ) {
        let bytes = record(&increment);
        let writer = Arc::clone(self);
        let appended = tokio::task::spawn_blocking(move || {
            writer
                .files
                .append(FileKind::Document, &writer.path, &bytes)
        })
        .await;
        match appended {
            Ok(Ok(())) => {
                let mut guard = self.lock();
                if let Some(current) = guard.as_mut().filter(|current| current.heads == since) {
                    current.heads = heads;
                    current.tail_bytes += increment.len() as u64 + 4;
                    debug!(
                        bytes = increment.len(),
                        tail_bytes = current.tail_bytes,
                        base_bytes = current.base_bytes,
                        "Appended to the document file"
                    );
                }
            }
            Ok(Err(error)) => {
                // The file may end in a partial record now; the next rewrite replaces it.
                error!(path = %self.path.display(), error = %error, "Couldn't append to the document file; rewriting it next");
                *self.lock() = None;
            }
            Err(error) => error!(error = %error, "Document file append task failed"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::document::testing::change;
    use automerge::{ROOT, ReadDoc, transaction::Transactable};

    #[test]
    fn whole_then_increments() {
        let mut doc = Document::new();
        change(&mut doc, |doc| doc.put(ROOT, "a", 1_i64).unwrap());
        let actor = doc.get_actor().clone();
        let heads = doc.get_heads();
        let mut bytes = whole(std::slice::from_ref(&actor), &doc.save()).unwrap();

        change(&mut doc, |doc| doc.put(ROOT, "b", 2_i64).unwrap());
        bytes.extend(record(&doc.save_after(&heads)));

        let (loaded, dropped) = parse(&bytes).unwrap();
        assert_eq!(dropped, 0);
        assert_eq!(loaded.actors, [actor.clone()]);
        assert_ne!(loaded.doc.get_actor(), &actor);
        assert!(loaded.doc.get(ROOT, "b").unwrap().is_some());

        // A record cut short by a crash is dropped.
        let mut cut = bytes.clone();
        cut.extend(record(b"whatever")[..6].iter());
        let (loaded, dropped) = parse(&cut).unwrap();
        assert_eq!(dropped, 6);
        assert!(loaded.doc.get(ROOT, "b").unwrap().is_some());

        assert!(parse(b"something else").is_err());
        assert!(parse(&bytes[..MAGIC.len() + 3]).is_err());
    }
}