use std::collections::{BTreeSet, HashSet};

use futures::{future, stream, FutureExt, StreamExt, TryStreamExt};
use litl::{impl_debug_as_litl, impl_nested_tagged_data_serde, NestedTaggedData};
use ridl::{
    hashing::HashOf,
    signing::{SignatureError, Signed, SignerID, SignerSecret},
};
use serde::Serialize;
use serde_derive::{Deserialize, Serialize};
use thiserror::Error;
use tracing::{error, debug, warn};

use crate::{
    causal_set::{CausalSet, CausalSetItem, OptimisticCausalSetFrontier},
    telepathic::{
        ApplyDiffErrorFor, ApplyDiffResult, ApplyDiffSuccess, Telepathic, TelepathicDiff,
    },
    StorageBackend, ObjectID,
};

#[derive(Clone, Serialize, Deserialize)]
pub struct SetHeader {
    pub inserter: SignerID,
    pub meta: Option<litl::Val>,
}

impl_debug_as_litl!(SetHeader);

#[derive(Debug)]
pub struct SetState {
    pub id: SetID,
    pub header: Option<SetHeader>,
    pub items: CausalSet<Signed<SetItem>>,
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct SetItem {
    pub data: litl::Val,
    pub prev: BTreeSet<SetItemID>,
}

#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct SetItemID(HashOf<SetItem>);

impl SetItemID {
    pub fn test_random() -> Self {
        SetItemID(HashOf::hash(&SetItem {
            data: litl::to_val(&rand::random::<u32>()).unwrap(),
            prev: BTreeSet::new(),
        }))
    }
}

impl NestedTaggedData for SetItemID {
    const TAG: &'static str = "setItemID";

    type Inner = HashOf<SetItem>;

    fn as_inner(&self) -> &Self::Inner {
        &self.0
    }

    fn from_inner(inner: Self::Inner) -> Self
    where
        Self: Sized,
    {
        SetItemID(inner)
    }
}

impl_nested_tagged_data_serde!(SetItemID);
impl_debug_as_litl!(SetItemID);

impl SetItem {
    pub fn new<D: Serialize, I: IntoIterator<Item = SetItemID>>(data: D, prev: I) -> Self {
        SetItem {
            data: litl::to_val(&data).unwrap(),
            prev: prev.into_iter().collect(),
        }
    }

    pub fn id(&self) -> SetItemID {
        SetItemID(HashOf::hash(self))
    }
}

impl CausalSetItem for Signed<SetItem> {
    type ID = SetItemID;

    fn id(&self) -> Self::ID {
        self.attested.id()
    }

    fn prev(&self) -> HashSet<&Self::ID> {
        self.prev.iter().collect()
    }
}

impl SetState {
    pub fn new_empty(id: SetID) -> Self {
        Self {
            id,
            header: None,
            items: CausalSet::new(),
        }
    }

    pub fn new<M: Serialize>(meta: Option<M>) -> (Self, SetWriteAccess) {
        let signer_secret = SignerSecret::new_random();

        let header = SetHeader {
            inserter: signer_secret.pub_id(),
            meta: meta.map(|meta| litl::to_val(&meta).unwrap()),
        };

        let id = SetID(HashOf::hash(&header));

        (
            Self {
                id,
                header: Some(header),
                items: CausalSet::new(),
            },
            SetWriteAccess(signer_secret),
        )
    }
}

#[derive(Copy, Clone, PartialEq, Eq, Hash)]
pub struct SetID(HashOf<SetHeader>);

impl NestedTaggedData for SetID {
    const TAG: &'static str = "set";

    type Inner = HashOf<SetHeader>;

    fn as_inner(&self) -> &Self::Inner {
        &self.0
    }

    fn from_inner(inner: Self::Inner) -> Self
    where
        Self: Sized,
    {
        SetID(inner)
    }
}

impl_nested_tagged_data_serde!(SetID);

impl_debug_as_litl!(SetID);

pub struct SetWriteAccess(pub(crate) SignerSecret);

impl NestedTaggedData for SetWriteAccess {
    const TAG: &'static str = "setWriteAccess";

    type Inner = SignerSecret;

    fn as_inner(&self) -> &Self::Inner {
        &self.0
    }

    fn from_inner(inner: Self::Inner) -> Self
    where
        Self: Sized,
    {
        SetWriteAccess(inner)
    }
}

impl_nested_tagged_data_serde!(SetWriteAccess);
impl_debug_as_litl!(SetWriteAccess);

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct SetDiff {
    pub id: SetID,
    pub header: Option<SetHeader>,
    pub new_items: Vec<Signed<SetItem>>,
}

impl TelepathicDiff for SetDiff {
    type ID = SetID;

    fn id(&self) -> SetID {
        self.id
    }
}

#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SetStateInfo {
    frontier: OptimisticCausalSetFrontier<Signed<SetItem>>,
    has_header: bool,
}

impl_debug_as_litl!(SetStateInfo);

impl Telepathic for SetState {
    type ID = SetID;
    type WriteAccess = SetWriteAccess;
    type StateInfo = SetStateInfo;
    type Diff = SetDiff;
    type Error = SetError;

    fn id(&self) -> Self::ID {
        self.id
    }

    fn try_apply_diff(
        &mut self,
        diff: Self::Diff,
    ) -> ApplyDiffResult<Self::StateInfo, Self::ID, Self::Diff, Self::Error> {
        let (header, got_header_first_time) = match (&mut self.header, &diff.header) {
            (None, None) => {
                return Err(ApplyDiffErrorFor::InvalidKnownStateAssumption(ObjectID::Set(self.id), "Didn't get set header even though we also don't have it".to_owned()));
            }
            (own @ None, Some(diff_header)) => {
                if HashOf::hash(diff_header) != self.id.0 {
                    return Err(SetError::InvalidHeaderHash.into());
                }
                *own = Some(diff_header.clone());
                (own.as_ref().unwrap(), true)
            }
            (Some(own_header), _) => (&*own_header, false),
        };

        let mut redundant_items = 0;

        let mut actually_new_items = Vec::new();

        for new_item in &diff.new_items {
            // performance shortcut: ignore items that are already in the set, potentially ignoring malformed incoming ones
            if self.items.contains_key(&new_item.id()) {
                redundant_items += 1;
                continue;
            }
            new_item
                .ensure_signed_by(&header.inserter)
                .map_err(SetError::InvalidSignature)?;

            self.items.insert(new_item.clone());
            actually_new_items.push(new_item.clone());
        }

        if redundant_items > 0 {
            warn!("Got {} redundant items in applied set diff", redundant_items)
        }

        let effective_diff = SetDiff {
            id: self.id,
            header: if got_header_first_time { Some(header.clone()) } else { None },
            new_items: actually_new_items,
        };
        Ok(ApplyDiffSuccess {
            new_state_info: SetStateInfo {
                frontier: self.items.as_optimistic_frontier(),
                has_header: self.header.is_some(),
            },
            effective_diff,
        })
    }

    fn state_info(&self) -> Option<Self::StateInfo> {
        let frontier = self.items.as_optimistic_frontier();
        if self.header.is_none() {
            None
        } else {
            Some(SetStateInfo {
                frontier,
                has_header: true,
            })
        }
    }

    fn diff_since(&self, state_info: Option<&Self::StateInfo>) -> Option<Self::Diff> {
        if let Some(mut state_info) = state_info.cloned() {
            let frontier = state_info.frontier.resolve(&self.items);
            let (new_items, _) = self.items.items_after(frontier);

            if new_items.is_empty() && state_info.has_header {
                None
            } else {
                Some(SetDiff {
                    id: self.id,
                    header: self.header.clone(),
                    new_items,
                })
            }
        } else {
            Some(SetDiff {
                id: self.id,
                header: self.header.clone(),
                new_items: self.items.values_ordered(),
            })
        }
    }

    fn load(
        id: SetID,
        storage: Box<dyn StorageBackend>,
    ) -> std::pin::Pin<Box<dyn futures::Stream<Item = Self::Diff>>> {
        let key = id.to_string();

        let header_stream = stream::once(storage.get_key(&format!("header_{}", key))).filter_map(
            move |maybe_header_bytes| {
                future::ready(maybe_header_bytes.and_then(|header_bytes| {
                    litl::from_slice(&header_bytes)
                        .map_err(|err| {
                            error!(err = ?err, id = ?id, "Failed to read loaded set header");
                            err
                        })
                        .ok()
                        .map(|header| SetDiff {
                            id,
                            header: Some(header),
                            new_items: vec![],
                        })
                }))
            },
        );

        let items_stream =
            litl::read_newln_sep_stream(storage.get_stream(&key).map(Ok).into_async_read())
                .filter_map(move |read_result| {
                    future::ready(match read_result {
                        Ok(item) => Some(item),
                        Err(err) => {
                            error!(err = ?err, id = ?id, "Failed to read loaded set item");
                            None
                        }
                    })
                })
                .ready_chunks(1000)
                // .map(|item| vec![item])
                .map(move |items| SetDiff {
                    id,
                    header: None,
                    new_items: items,
                });

        header_stream.chain(items_stream).boxed_local()
    }

    fn store(
        effective_diff: Self::Diff,
        storage: Box<dyn StorageBackend>,
    ) -> std::pin::Pin<Box<dyn futures::Future<Output = ()>>> {
        let key = effective_diff.id.to_string();
        async move {
            if let Some(header) = &effective_diff.header {
                storage
                    .set_key(&format!("header_{}", key), litl::to_vec(&header).unwrap())
                    .await;
            }

            if !effective_diff.new_items.is_empty() {
                storage
                    .append_to_stream(
                        &key,
                        litl::to_newln_sep_vec(effective_diff.new_items.iter()).unwrap(),
                        None,
                    )
                    .await;
            }
        }
        .boxed_local()
    }
}

#[derive(Error, Debug)]
pub enum SetError {
    #[error("Invalid header hash")]
    InvalidHeaderHash,
    #[error(transparent)]
    InvalidSignature(#[from] SignatureError),
}