//! Stage `ATrium` blocks, sign the final root, and atomically publish the index, head, and outbox.
use super::{MAX_REPO_RETRIES, RepositoryService, car::map_block_error};
use crate::error::ErrorContext;
use crate::{
    database::{PendingRepoCommit, RepoBlockStore, RepoHead},
    domain::{
        account::Account,
        event::RepositoryEvent,
        record::{Record, RecordMutation},
    },
    error::{Error, Result},
};
use atrium_api::com::atproto::repo::defs::CommitMetaData;
use atrium_api::types::string::{Cid as AtriumCid, Did, Tid};
use atrium_crypto::verify::verify_signature;
use atrium_repo::{
    Repository,
    blockstore::{AsyncBlockStoreWrite as _, DAG_CBOR as ATRIUM_DAG_CBOR, SHA2_256},
    repo::Error as AtriumRepoError,
};
use cid::Cid;
use cpds_protocol::atproto_record;
use serde::Serialize;

#[derive(Serialize)]
struct UnsignedCommit<'a> {
    did: Did,
    version: i64,
    data: Cid,
    rev: &'a Tid,
    prev: Option<Cid>,
}

#[derive(Serialize)]
struct SignedCommit<'a> {
    #[serde(flatten)]
    content: UnsignedCommit<'a>,
    #[serde(rename = "sig", with = "serde_bytes")]
    signature: Vec<u8>,
}

fn record_path(record: &Record) -> String {
    format!("{}/{}", record.collection, record.rkey)
}

async fn stage_record_blocks(
    blocks: &mut RepoBlockStore,
    mutations: &[RecordMutation],
) -> Result<()> {
    for mutation in mutations {
        let record = match mutation {
            RecordMutation::Create(record)
            | RecordMutation::Update {
                replacement: record,
                ..
            } => record,
            RecordMutation::Delete(_) => continue,
        };
        let data = atproto_record(&record.value)?;
        let bytes = serde_ipld_dagcbor::to_vec(&data)
            .context(Error::Internal, "validated record cannot be encoded")?;
        let cid = blocks
            .write_block(ATRIUM_DAG_CBOR, SHA2_256, &bytes)
            .await
            .map_err(map_block_error)?;
        if cid != record.cid {
            return Err(Error::Internal(
                "ATrium record CID differs from validated record CID".to_owned(),
            ));
        }
    }
    Ok(())
}

async fn apply_repo_mutations(
    repo: &mut Repository<&mut RepoBlockStore>,
    mutations: &[RecordMutation],
) -> Result<Cid> {
    let mut tree = repo.tree();
    for mutation in mutations {
        match mutation {
            RecordMutation::Create(record) => tree.add(&record_path(record), record.cid).await,
            RecordMutation::Update { replacement, .. } => {
                tree.update(&record_path(replacement), replacement.cid)
                    .await
            }
            RecordMutation::Delete(record) => tree.delete(&record_path(record)).await,
        }
        .map_err(|error| map_repo_error(error.into()))?;
    }
    Ok(tree.root())
}

async fn event_block_cids(blocks: &mut RepoBlockStore, root: Cid) -> Result<Vec<Cid>> {
    let reachable = {
        let mut repo = Repository::open(&mut *blocks, root)
            .await
            .map_err(map_repo_error)?;
        repo.export()
            .await
            .map_err(map_repo_error)?
            .collect::<Vec<_>>()
    };
    Ok(blocks
        .reachable_writes(reachable)?
        .map(|(cid, _)| *cid)
        .collect())
}

pub(super) fn invalid_swap(message: &str) -> Error {
    Error::xrpc("InvalidSwap", message)
}

pub(super) fn map_repo_error(error: AtriumRepoError) -> Error {
    match error {
        AtriumRepoError::BlockStore(error) => map_block_error(error),
        AtriumRepoError::Mst(atrium_repo::mst::Error::KeyAlreadyExists) => {
            Error::Conflict("record already exists".to_owned())
        }
        AtriumRepoError::Mst(atrium_repo::mst::Error::KeyNotFound) => Error::NotFound,
        error => Error::Internal(format!("ATrium repository operation failed: {error}")),
    }
}

impl RepositoryService {
    fn record_event(mutation: &RecordMutation) -> RepositoryEvent {
        match mutation {
            RecordMutation::Create(record) => RepositoryEvent::RecordCreated {
                did: record.did.clone(),
                collection: record.collection.clone(),
                rkey: record.rkey.clone(),
                cid: record.cid,
            },
            RecordMutation::Update {
                current,
                replacement,
            } => RepositoryEvent::RecordUpdated {
                did: replacement.did.clone(),
                collection: replacement.collection.clone(),
                rkey: replacement.rkey.clone(),
                cid: replacement.cid,
                previous_cid: current.cid,
            },
            RecordMutation::Delete(record) => RepositoryEvent::RecordDeleted {
                did: record.did.clone(),
                collection: record.collection.clone(),
                rkey: record.rkey.clone(),
                previous_cid: record.cid,
            },
        }
    }

    pub(super) async fn commit_record_mutations(
        &self,
        actor: &Account,
        expected_commit: Option<&Cid>,
        mutations: &[RecordMutation],
    ) -> Result<CommitMetaData> {
        for _ in 0..MAX_REPO_RETRIES {
            let previous = self.database.repo_head(&actor.did)?;
            if expected_commit
                .is_some_and(|expected| previous.as_ref().is_none_or(|head| expected != &head.root))
            {
                return Err(invalid_swap("repository commit does not match"));
            }
            let rev = previous.as_ref().map_or_else(
                || Ok(self.tids.next()),
                |head| {
                    self.tids.next_after(&head.rev).ok_or_else(|| {
                        Error::Internal("stored repository revision is invalid".to_owned())
                    })
                },
            )?;
            let tid = Tid::new(rev.clone())
                .context(Error::Internal, "generated repository revision is invalid")?;
            let mut blocks = RepoBlockStore::new(self.database.clone());
            stage_record_blocks(&mut blocks, mutations).await?;
            let mut repo = self
                .open_or_create_repo(&mut blocks, actor, previous.as_ref())
                .await?;
            let previous_data = previous.as_ref().map(|_| repo.commit().data());
            let data = apply_repo_mutations(&mut repo, mutations).await?;
            drop(repo);
            let root = self
                .write_signed_commit(&mut blocks, actor, data, &tid)
                .await?;
            let event_blocks = event_block_cids(&mut blocks, root).await?;
            let staged = blocks.writes();
            let commit = PendingRepoCommit {
                did: actor.did.clone(),
                expected_head: previous.as_ref().map(|head| head.root),
                previous_rev: previous.as_ref().map(|head| head.rev.clone()),
                previous_data,
                head: RepoHead {
                    root,
                    rev: rev.clone(),
                },
                blocks: staged,
                event_blocks,
                mutations: mutations.to_vec(),
            };
            match self
                .database
                .commit_repo(&commit, self.repo_event_retention)
            {
                Ok(()) => {
                    self.publish_events(mutations.iter().map(Self::record_event))
                        .await;
                    return Ok(CommitMetaData {
                        cid: AtriumCid::new(root),
                        rev: tid,
                    });
                }
                Err(Error::Conflict(message)) if message == "repository head changed" => {}
                Err(error) => return Err(error),
            }
        }
        Err(Error::Conflict(
            "repository changed too frequently; retry the request".to_owned(),
        ))
    }

    async fn write_signed_commit(
        &self,
        blocks: &mut RepoBlockStore,
        actor: &Account,
        data: Cid,
        rev: &Tid,
    ) -> Result<Cid> {
        let did = Did::new(actor.did.clone())
            .context(Error::Internal, "hosted account DID is invalid")?;
        let commit = UnsignedCommit {
            did,
            version: 3,
            data,
            rev,
            prev: None,
        };
        let signature = self.sign(
            &serde_ipld_dagcbor::to_vec(&commit)
                .context(Error::Internal, "repository commit cannot be encoded")?,
        )?;
        let signed = SignedCommit {
            content: commit,
            signature,
        };
        let bytes = serde_ipld_dagcbor::to_vec(&signed).context(
            Error::Internal,
            "signed repository commit cannot be encoded",
        )?;
        blocks
            .write_block(ATRIUM_DAG_CBOR, SHA2_256, &bytes)
            .await
            .map_err(map_block_error)
    }

    async fn open_or_create_repo<'a>(
        &self,
        blocks: &'a mut RepoBlockStore,
        actor: &Account,
        head: Option<&RepoHead>,
    ) -> Result<Repository<&'a mut RepoBlockStore>> {
        if let Some(head) = head {
            return self.open_repo(blocks, head.root, &head.rev).await;
        }

        let did = Did::new(actor.did.clone())
            .context(Error::Internal, "hosted account DID is invalid")?;
        let builder = Repository::create(&mut *blocks, did)
            .await
            .map_err(map_repo_error)?;
        let signature = self.sign(&builder.bytes())?;
        let mut repo = builder.finalize(signature).await.map_err(map_repo_error)?;

        for record in self.database.all_records(&actor.did)? {
            let path = format!("{}/{}", record.collection, record.rkey);
            let data = atproto_record(&record.value).context(
                Error::Internal,
                format_args!("stored record {record} is invalid"),
            )?;
            let (builder, cid) = repo.add_raw(&path, &data).await.map_err(map_repo_error)?;
            if cid != record.cid {
                return Err(Error::Internal(format!(
                    "stored record {record} has noncanonical CID"
                )));
            }
            let signature = self.sign(&builder.bytes())?;
            builder.finalize(signature).await.map_err(map_repo_error)?;
        }
        Ok(repo)
    }

    pub(super) async fn open_repo<'a>(
        &self,
        blocks: &'a mut RepoBlockStore,
        root: Cid,
        expected_rev: &str,
    ) -> Result<Repository<&'a mut RepoBlockStore>> {
        let repo = Repository::open(&mut *blocks, root)
            .await
            .map_err(map_repo_error)?;
        let commit = repo.commit();
        if commit.rev().to_string() != expected_rev {
            return Err(Error::Internal(
                "stored repository head revision does not match its commit".to_owned(),
            ));
        }
        verify_signature(&self.signing_did, &commit.bytes(), commit.sig())
            .context(Error::Internal, "stored repository signature is invalid")?;
        Ok(repo)
    }

    fn sign(&self, bytes: &[u8]) -> Result<Vec<u8>> {
        self.signer
            .sign(bytes)
            .context(Error::Internal, "repository commit signing failed")
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(Serialize)]
    struct PreviousUnsignedCommit {
        did: Did,
        version: i64,
        data: Cid,
        rev: Tid,
        prev: Option<Cid>,
    }

    #[derive(Serialize)]
    struct PreviousSignedCommit {
        #[serde(flatten)]
        content: PreviousUnsignedCommit,
        #[serde(rename = "sig", with = "serde_bytes")]
        signature: Vec<u8>,
    }

    #[test]
    fn borrowed_revision_keeps_unsigned_and_signed_commit_bytes() {
        let data = cpds_protocol::dag_cbor_cid(b"data");
        let previous = cpds_protocol::dag_cbor_cid(b"previous");
        for rev in ["2222222222222", "3jqfcqzm3fo2j", "jzzzzzzzzzzzz"] {
            let rev = Tid::new(rev.to_owned()).unwrap();
            for prev in [None, Some(previous)] {
                for signature in [Vec::new(), vec![0, 255], (0..=255).collect()] {
                    let owned = PreviousUnsignedCommit {
                        did: Did::new("did:web:alice.test".to_owned()).unwrap(),
                        version: 3,
                        data,
                        rev: rev.clone(),
                        prev,
                    };
                    let borrowed = UnsignedCommit {
                        did: owned.did.clone(),
                        version: owned.version,
                        data,
                        rev: &rev,
                        prev,
                    };
                    assert_eq!(
                        serde_ipld_dagcbor::to_vec(&borrowed).unwrap(),
                        serde_ipld_dagcbor::to_vec(&owned).unwrap()
                    );
                    assert_eq!(
                        serde_json::to_vec(&borrowed).unwrap(),
                        serde_json::to_vec(&owned).unwrap()
                    );
                    let owned = PreviousSignedCommit {
                        content: owned,
                        signature: signature.clone(),
                    };
                    let borrowed = SignedCommit {
                        content: borrowed,
                        signature,
                    };
                    let expected = serde_ipld_dagcbor::to_vec(&owned).unwrap();
                    let actual = serde_ipld_dagcbor::to_vec(&borrowed).unwrap();
                    assert_eq!(actual, expected);
                    assert_eq!(
                        cpds_protocol::dag_cbor_cid(&actual),
                        cpds_protocol::dag_cbor_cid(&expected)
                    );
                    assert_eq!(
                        serde_json::to_vec(&borrowed).unwrap(),
                        serde_json::to_vec(&owned).unwrap()
                    );
                }
            }
        }
    }
}