use std::{
    collections::HashMap,
    ops::Bound,
    sync::{Arc, Mutex},
};

use chrono::{DateTime, Utc};
use ftypes::{
    AmendmentMessageContent, ErrorKind, Message, MessageId, MessageSubId, ReceivedMessage,
    ReceivedMessageContent, Result, Store, ThreadId, Timer,
    convert::from_presage::MessageConversionAdditionalDataProvider,
};
use futures::{Stream, StreamExt};
use presage::{
    libsignal_service::zkgroup::GroupMasterKeyBytes,
    store::{ContentsStore, StateStore},
};

#[derive(Clone)]
pub struct DirectPresageWrapperStore {
    store: presage_store_sqlite::SqliteStore,
    // Note: presage updates the in-store group before we can see the message.
    // We therefore cannot notice changes to the group only via the store, but will also need cached revision values in-memory.
    // Analogous for expire timers.
    group_revision_cache: Arc<Mutex<HashMap<GroupMasterKeyBytes, u32>>>,
    expire_timer_cache: Arc<Mutex<HashMap<ThreadId, Timer>>>,
}

impl DirectPresageWrapperStore {
    pub async fn new(path: &str, password: Option<&str>) -> Result<Self> {
        let s = DirectPresageWrapperStore {
            store: presage_store_sqlite::SqliteStore::open_with_passphrase(
                path,
                password,
                presage_store_sqlite::OnNewIdentity::Trust,
            )
            .await
            .map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to set up presage store")
                    .add_path_context(path)
            })?,
            group_revision_cache: Arc::new(Mutex::new(HashMap::new())),
            expire_timer_cache: Arc::new(Mutex::new(HashMap::new())),
        };

        s.update_group_revision_cache().await?;
        s.update_expire_timer_cache().await?;

        Ok(s)
    }

    async fn update_group_revision_cache(&self) -> Result<()> {
        let map = self
            .groups()
            .await?
            .into_iter()
            .map(|g| (g.secrets.master_key, g.secrets.revision))
            .collect();
        *self.group_revision_cache.lock().expect("Poisoned lock") = map;

        Ok(())
    }

    async fn update_expire_timer_cache(&self) -> Result<()> {
        let map = self
            .threads()
            .await?
            .into_iter()
            .map(|g| (g.id(), g.disappearing_messages_timer()))
            .collect();
        *self.expire_timer_cache.lock().expect("Poisoned lock") = map;

        Ok(())
    }

    /// Don't have a nice way to convert from ftypes::ThreadId to presage::Thread, as the presage one requires group master keys, but the ftypes one requires group IDs.
    async fn ftypes_thread_to_presage_thread(
        &self,
        thread: ftypes::ThreadId,
    ) -> Result<presage::store::Thread> {
        match thread {
            ftypes::ThreadId::Contact(account_id) => {
                Ok(presage::store::Thread::Contact(account_id.0))
            }
            ftypes::ThreadId::Group(group_id) => {
                let Some(group) = self.group(group_id).await? else {
                    return Err(ErrorKind::Database
                        .error()
                        .with_message("thread for group not found")
                        .add_value_dbg_context(group_id));
                };
                Ok(presage::store::Thread::Group(group.secrets.master_key))
            }
        }
    }
}

impl Store for DirectPresageWrapperStore {
    type PresageStore = presage_store_sqlite::SqliteStore;

    fn presage_store(&self) -> Self::PresageStore {
        self.store.clone()
    }

    async fn groups(&self) -> Result<Vec<ftypes::Group>> {
        let groups: std::result::Result<
            Vec<(GroupMasterKeyBytes, presage::model::groups::Group)>,
            presage_store_sqlite::SqliteStoreError,
        > = self
            .store
            .groups()
            .await
            .map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to query groups")
                    .add_field_context("groups")
            })?
            .collect();
        Ok(groups
            .map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to query groups")
                    .add_field_context("groups")
            })?
            .into_iter()
            .map(|(b, g)| ftypes::convert::from_presage::group(b, g))
            .collect())
    }

    async fn contacts(&self) -> Result<Vec<ftypes::Contact>> {
        let contacts: std::result::Result<
            Vec<presage::model::contacts::Contact>,
            presage_store_sqlite::SqliteStoreError,
        > = self
            .store
            .contacts()
            .await
            .map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to query contacts")
                    .add_field_context("contacts")
            })?
            .collect();
        let contacts = contacts.map_err(|e| {
            ErrorKind::Database
                .caused_by(e)
                .with_message("failed to query contacts")
                .add_field_context("contacts")
        })?;
        let mut result = Vec::with_capacity(contacts.len());

        for c in contacts {
            let profile_key = if let Ok(p) = c.profile_key.clone().try_into() {
                Some(presage::libsignal_service::prelude::ProfileKey { bytes: p })
            } else {
                self.store
                    .profile_key(&presage::libsignal_service::protocol::ServiceId::Aci(
                        c.uuid.into(),
                    ))
                    .await
                    .map_err(|e| {
                        ErrorKind::Database
                            .caused_by(e)
                            .with_message("failed to query profile key for uuid")
                            .add_field_context("contacts")
                            .add_value_dbg_context(c.uuid)
                    })?
            };
            if let Some(profile_key) = profile_key {
                let p = self.store.profile(c.uuid, profile_key).await.map_err(|e| {
                    ErrorKind::Database
                        .caused_by(e)
                        .with_message("failed to query profile for contact")
                        .add_field_context("contacts")
                })?;
                result.push(ftypes::convert::from_presage::contact(
                    c,
                    Some(profile_key),
                    p,
                    None,
                    None,
                )?);
            } else {
                result.push(ftypes::convert::from_presage::contact(
                    c, None, None, None, None,
                )?);
            }
        }

        Ok(result)
    }

    async fn group(&self, id: ftypes::GroupId) -> Result<Option<ftypes::Group>> {
        Ok(self.groups().await?.into_iter().find(|g| g.id == id))
    }

    async fn upsert_contact(&self, _contact: ftypes::Contact) -> Result<()> {
        // Only need to refresh the expire timer cache; the rest is stored by presage.
        self.update_expire_timer_cache().await?;
        Ok(())
    }

    async fn upsert_group(&self, _group: ftypes::Group) -> Result<()> {
        // Only need to refresh the group revision cache; the rest is stored by presage.
        self.update_group_revision_cache().await?;
        Ok(())
    }

    async fn contact(&self, id: ftypes::ContactId) -> Result<Option<ftypes::Contact>> {
        if let Some(group) = id.in_group {
            let Some(group) = self.group(group).await? else {
                return Err(ErrorKind::Database
                    .error()
                    .with_message("group from queried contact does not exist")
                    .add_group_context(group)
                    .add_contact_context(id));
            };
            let Some(member) = group.members.into_iter().find(|m| m.id == id) else {
                return Err(ErrorKind::Database
                    .error()
                    .with_message("contact not in queried group")
                    .add_group_context(group.id)
                    .add_contact_context(id));
            };

            let Some(profile_key) = member.secrets.profile_key else {
                return Err(ErrorKind::Database
                    .error()
                    .with_message("contact in queried group does not have a profile key")
                    .add_group_context(group.id)
                    .add_contact_context(id));
            };

            let profile = self
                .store
                .profile(id.account_id.0.raw_uuid(), profile_key)
                .await
                .map_err(|e| {
                    ErrorKind::Database
                        .caused_by(e)
                        .with_message("failed to query profile for contact")
                        .add_field_context("contacts")
                })?;

            let Some(contact) = self
                .store
                .contact_by_id(&id.account_id.0)
                .await
                .map_err(|e| {
                    ErrorKind::Database
                        .caused_by(e)
                        .with_message("failed to query contact for group from database")
                        .add_group_context(group.id)
                        .add_contact_context(id)
                })?
            else {
                return Ok(None);
            };

            ftypes::convert::from_presage::contact(
                contact,
                Some(profile_key),
                profile,
                Some(group.id),
                Some(member),
            )
            .map(Some)
        } else {
            Ok(self.contacts().await?.into_iter().find(|c| c.id == id))
        }
    }

    async fn message(&self, id: MessageId) -> Result<Option<ftypes::RichMessage>> {
        let thread = self
            .thread_id_from_message_sender_timestamps(&id.sender, &[id.timestamp])
            .await?;
        let messages = self
            .messages(
                thread[0].ok_or_else(|| {
                    ErrorKind::Database
                        .error()
                        .with_message("cannot find thread for message")
                        .add_message_id_context(id)
                })?,
                id.timestamp..=id.timestamp,
            )
            .await?;

        futures::pin_mut!(messages);

        if let Some(res) = messages.next().await {
            Ok(Some(res?))
        } else {
            Ok(None)
        }
    }

    async fn messages(
        &self,
        thread: ftypes::ThreadId,
        range: impl std::ops::RangeBounds<DateTime<Utc>>,
    ) -> Result<impl Stream<Item = Result<ftypes::RichMessage>>> {
        let presage_thread = self.ftypes_thread_to_presage_thread(thread).await?;

        let presage_range = (
            range.start_bound().map(|b| b.timestamp_millis() as u64),
            Bound::Unbounded, // Note: Cannot use given end bound; reaction messages and other amendment messages may come later.
        );

        let presage_result = self
            .store
            .messages(&presage_thread, presage_range)
            .await
            .map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to query messages from database")
                    .add_thread_context(thread)
            })?;

        let plain_messages = futures::stream::iter(presage_result).then(move |m| async move {
            match m {
                Ok(m) => ftypes::convert::from_presage::message(m, self.clone())
                    .await
                    .map(<Vec<ReceivedMessage>>::from),
                Err(e) => Err(ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to query messages from database (single message failed)")
                    .add_thread_context(thread)),
            }
        });

        Ok(async_stream::try_stream! {
            let mut amendments: HashMap<MessageId, Vec<Message<AmendmentMessageContent>>> = HashMap::new();

            for await plain_message in plain_messages {
                for m in plain_message? {
                    if let Some(mut rich) = m.as_rich_message() {
                        let amends = amendments.remove(&rich.id).unwrap_or_default();
                        for a in amends {
                            rich = rich.amend(a);
                        }
                        if range.contains(&rich.id.timestamp) {
                            yield rich;
                        }
                    } else {
                        let Some(amendment) = m.as_amendment_message() else {
                            log::warn!("Stored non-main and non-amendment message in store");
                            continue;
                        };

                        for amends in amendment.amends() {
                            amendments.entry(amends).and_modify(|v| v.insert(0, amendment.clone())).or_insert_with(|| vec![amendment.clone()]);
                        }
                    }
                }
            }
        })
    }

    async fn registration_data(&self) -> Result<ftypes::RegistrationData> {
        let presage_data = self.store.load_registration_data().await;

        match presage_data {
            Ok(Some(data)) => Ok(ftypes::convert::from_presage::registration_data(data)),
            Ok(None) => Err(ErrorKind::Database
                .error()
                .with_message("store does not have access to registration data")),
            Err(e) => Err(ErrorKind::Database
                .caused_by(e)
                .with_message("failed to load registration data")),
        }
    }

    async fn clear_registration(&self) -> Result<()> {
        self.presage_store()
            .clear_registration()
            .await
            .map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to clear registration")
            })
    }

    async fn clear_contents(&self) -> Result<()> {
        self.presage_store().clear_contents().await.map_err(|e| {
            ErrorKind::Database
                .caused_by(e)
                .with_message("failed to clear contents")
        })
    }

    async fn clear_messages(&self) -> Result<()> {
        self.presage_store().clear_messages().await.map_err(|e| {
            ErrorKind::Database
                .caused_by(e)
                .with_message("failed to clear messages")
        })
    }

    async fn clear_thread_messages(&self, thread: ftypes::ThreadId) -> Result<()> {
        let presage_thread = self.ftypes_thread_to_presage_thread(thread).await?;

        self.presage_store()
            .clear_thread(&presage_thread)
            .await
            .map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to clear messages for thread")
            })
    }

    async fn store_message(&self, message: ReceivedMessage) -> Result<()> {
        // Most things are correctly handled by presage.

        // Receipt messages in groups are stored in PMs instead.
        // Note: One cannot delete the message from the presage thread; one receipt message could contain messages from multiple threads.
        if let ReceivedMessageContent::Amendment(AmendmentMessageContent::Receipt(_)) =
            &message.content
            && message.thread.group().is_some()
        {
            let wrong_thread = presage::store::Thread::Contact(message.id.sender.0);
            let correct_thread = self.ftypes_thread_to_presage_thread(message.thread).await?;

            let presage_store = self.presage_store();
            let content = presage_store
                .message(
                    &wrong_thread,
                    message.id.timestamp.timestamp_millis() as u64,
                )
                .await
                .map_err(|e| {
                    ErrorKind::Database
                        .caused_by(e)
                        .with_message(
                            "failed to query receipt message from store to move to group thread",
                        )
                        .add_message_context(message.clone())
                })?;

            if let Some(content) = content {
                presage_store
                    .save_message(&correct_thread, content)
                    .await
                    .map_err(|e| {
                        ErrorKind::Database
                            .caused_by(e)
                            .with_message("failed to move receipt message to group thread")
                            .add_message_context(message.clone())
                    })?;
            }
        }

        Ok(())
    }
}

impl MessageConversionAdditionalDataProvider for DirectPresageWrapperStore {
    async fn thread_id_from_message_sender_timestamps(
        &self,
        sender: &ftypes::AccountId,
        timestamp: &[DateTime<Utc>],
    ) -> Result<Vec<Option<ftypes::ThreadId>>> {
        let mut result = vec![None; timestamp.len()];

        for (i, ts) in timestamp.iter().enumerate() {
            let ts: u64 = ts.timestamp_millis().try_into().map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("could not convert timestamp to u64")
                    .add_value_dbg_context(ts)
            })?;

            let thread = self
                .store
                .thread_for_sender_and_timestamp(&sender.0, ts)
                .await
                .map_err(|e| {
                    ErrorKind::Database
                        .caused_by(e)
                        .with_message("could not query message from database")
                })?;

            result[i] = thread.map(ftypes::convert::from_presage::thread);
        }

        Ok(result)
    }

    async fn profile_updated(
        &self,
        sender: &ftypes::ContactId,
        profile_key: &presage::libsignal_service::prelude::ProfileKey,
    ) -> Result<bool> {
        Ok(self
            .store
            .profile(sender.account_id.0.raw_uuid(), *profile_key)
            .await
            .map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to query profile for contact")
                    .add_field_context("profile")
                    .add_value_dbg_context(sender)
                    .add_value_dbg_context(profile_key)
            })?
            .is_none())
    }

    async fn group_updated(
        &self,
        _group_id: &ftypes::GroupId,
        secrets: &ftypes::GroupSecrets,
    ) -> Result<bool> {
        Ok(self
            .group_revision_cache
            .lock()
            .expect("Poisoned mutex")
            .get(&secrets.master_key)
            .is_none_or(|r| *r != secrets.revision))
    }

    async fn original_message_id(&self, id: &MessageId) -> Result<MessageId> {
        log::trace!("Querying original message id for {id:?}");
        let Some(thread) = self
            .store
            .thread_for_sender_and_timestamp(&id.sender.0, id.timestamp.timestamp_millis() as u64)
            .await
            .map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to query thread for sender and timestamp")
                    .add_message_id_context(*id)
            })?
        else {
            log::warn!(
                "Cannot find the thread of hte original message id referenced by message {id:?}; assuming the message is the original message."
            );
            return Ok(*id);
        };

        let mut timestamp = id.timestamp.timestamp_millis() as u64;

        loop {
            log::trace!("Current original message id timestamp: {timestamp}");

            let Some(message) = self.store.message(&thread, timestamp).await.map_err(|e| {
                ErrorKind::Database
                    .caused_by(e)
                    .with_message("failed to query original message for message id")
                    .add_message_id_context(*id)
            })?
            else {
                log::warn!(
                    "Cannot find the original message id referenced by message {id:?}; assuming the message is the original message."
                );
                return Ok(*id);
            };

            match message.body {
                presage::libsignal_service::content::ContentBody::EditMessage(edit_message)
                | presage::libsignal_service::content::ContentBody::SynchronizeMessage(
                    presage::libsignal_service::content::SyncMessage {
                        sent:
                            Some(presage::libsignal_service::content::sync_message::Sent {
                                edit_message: Some(edit_message),
                                ..
                            }),
                        ..
                    },
                ) => {
                    log::trace!("Original message id query got an edit message; go back further");
                    assert!(
                        edit_message.target_sent_timestamp() < timestamp,
                        "original message ids should only go backwards in time"
                    );
                    timestamp = edit_message.target_sent_timestamp();
                }
                _ => {
                    log::trace!(
                        "Original message id query got some other message; must be original"
                    );
                    break;
                }
            }
        }

        Ok(MessageId {
            sender: id.sender,
            timestamp: ftypes::convert::from_presage::timestamp(timestamp)?,
            sub_id: MessageSubId::Primary,
        })
    }

    async fn expire_timer_updated(
        &self,
        thread: &ftypes::ThreadId,
        timer: ftypes::Timer,
    ) -> Result<bool> {
        Ok(self
            .expire_timer_cache
            .lock()
            .expect("Poisoned mutex")
            .get(thread)
            != Some(&timer))
    }
}
