/* Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

use super::COMPONENT;
use crate::shard::IggyShard;
use crate::shard::transmission::frame::ShardResponse;
use crate::shard::transmission::message::{
    ResolvedPartition, ResolvedTopic, ShardRequest, ShardRequestPayload,
};
use crate::streaming::partitions::journal::Journal;
use crate::streaming::polling_consumer::PollingConsumer;
use crate::streaming::segments::{IggyIndexesMut, IggyMessagesBatchMut, IggyMessagesBatchSet};
use err_trail::ErrContext;
use iggy_common::IggyPollMetadata;
use iggy_common::PooledBuffer;
use iggy_common::sharding::IggyNamespace;
use iggy_common::{
    Consumer, EncryptorKind, IGGY_MESSAGE_HEADER_SIZE, Identifier, IggyError, PollingStrategy,
};
use std::sync::atomic::Ordering;
use tracing::error;

impl IggyShard {
    /// Appends messages to partition. Permission must be checked by caller via
    /// `resolve_topic_for_append()` before calling this method.
    pub async fn append_messages(
        &self,
        partition: ResolvedPartition,
        batch: IggyMessagesBatchMut,
    ) -> Result<(), IggyError> {
        if batch.count() == 0 {
            return Ok(());
        }

        let namespace = IggyNamespace::new(
            partition.stream_id,
            partition.topic_id,
            partition.partition_id,
        );

        let payload = ShardRequestPayload::SendMessages { batch };
        let request = ShardRequest::data_plane(namespace, payload);

        match self.send_to_data_plane(request).await? {
            ShardResponse::SendMessages => Ok(()),
            ShardResponse::ErrorResponse(err) => Err(err),
            _ => unreachable!("Expected SendMessages response"),
        }
    }

    /// Polls messages from partition. Permission must be checked by caller via
    /// `resolve_topic_for_poll()` before calling this method.
    pub async fn poll_messages(
        &self,
        client_id: u32,
        topic: ResolvedTopic,
        consumer: Consumer,
        maybe_partition_id: Option<u32>,
        args: PollingArgs,
    ) -> Result<(IggyPollMetadata, IggyMessagesBatchSet), IggyError> {
        let Some((consumer, partition_id)) = self.resolve_consumer_with_partition_id(
            topic,
            &consumer,
            client_id,
            maybe_partition_id,
            true,
        )?
        else {
            return Ok((IggyPollMetadata::new(0, 0), IggyMessagesBatchSet::empty()));
        };

        let namespace = IggyNamespace::new(topic.stream_id, topic.topic_id, partition_id);

        let payload = ShardRequestPayload::PollMessages { consumer, args };
        let request = ShardRequest::data_plane(namespace, payload);

        let (metadata, batch) = match self.send_to_data_plane(request).await? {
            ShardResponse::PollMessages(result) => result,
            ShardResponse::ErrorResponse(err) => return Err(err),
            _ => unreachable!("Expected PollMessages response"),
        };

        let batch = if let Some(encryptor) = &self.encryptor {
            self.decrypt_messages(batch, encryptor).await?
        } else {
            batch
        };

        // Track last offset sent to CG member for cooperative rebalance.
        if let PollingConsumer::ConsumerGroup(group_id, _) = &consumer
            && let Some(last_offset) = batch.last_offset()
        {
            self.metadata.record_polled_offset(
                topic.stream_id,
                topic.topic_id,
                group_id.0,
                partition_id,
                last_offset,
            );
        }

        Ok((metadata, batch))
    }

    pub async fn flush_unsaved_buffer(
        &self,
        user_id: u32,
        partition: ResolvedPartition,
        fsync: bool,
    ) -> Result<(), IggyError> {
        self.metadata
            .perm_append_messages(user_id, partition.stream_id, partition.topic_id)
            .error(|e: &IggyError| {
                format!("{COMPONENT} (error: {e}) - permission denied to flush unsaved buffer for user {} on stream ID: {}, topic ID: {}", user_id, partition.stream_id as u32, partition.topic_id as u32)
            })?;

        let namespace = IggyNamespace::new(
            partition.stream_id,
            partition.topic_id,
            partition.partition_id,
        );
        let payload = ShardRequestPayload::FlushUnsavedBuffer { fsync };
        let request = ShardRequest::data_plane(namespace, payload);

        match self.send_to_data_plane(request).await? {
            ShardResponse::FlushUnsavedBuffer { .. } => Ok(()),
            ShardResponse::ErrorResponse(err) => Err(err),
            _ => unreachable!("Expected FlushUnsavedBuffer response"),
        }
    }

    /// Flushes unsaved messages from the partition store to disk.
    /// Returns the number of messages saved.
    pub(crate) async fn flush_unsaved_buffer_from_local_partitions(
        &self,
        namespace: &IggyNamespace,
        fsync: bool,
    ) -> Result<u32, IggyError> {
        let frozen_batches = {
            let mut partitions = self.local_partitions.borrow_mut();
            let Some(partition) = partitions.get_mut(namespace) else {
                return Ok(0);
            };
            if !partition.log.has_segments() || partition.log.journal().is_empty() {
                return Ok(0);
            }
            let batches = partition.log.journal_mut().commit();
            partition.log.ensure_indexes();
            batches.append_indexes_to(partition.log.active_indexes_mut().unwrap());

            let frozen: Vec<_> = batches
                .into_inner()
                .into_iter()
                .map(|mut b| b.freeze())
                .collect();
            partition.log.set_in_flight(frozen.clone());
            frozen
        };

        let saved_count = self
            .persist_frozen_batches_to_disk(namespace, frozen_batches)
            .await?;

        if fsync {
            self.fsync_all_messages_from_local_partitions(namespace)
                .await?;
        }

        Ok(saved_count)
    }

    pub(crate) async fn fsync_all_messages_from_local_partitions(
        &self,
        namespace: &IggyNamespace,
    ) -> Result<(), IggyError> {
        let storage = {
            let partitions = self.local_partitions.borrow();
            let Some(partition) = partitions.get(namespace) else {
                return Ok(());
            };
            if !partition.log.has_segments() {
                return Ok(());
            }
            partition.log.active_storage().clone()
        };

        if storage.messages_writer.is_none() || storage.index_writer.is_none() {
            return Ok(());
        }

        if let Some(ref messages_writer) = storage.messages_writer
            && let Err(e) = messages_writer.fsync().await
        {
            tracing::error!(
                "Failed to fsync messages writer for partition {:?}: {}",
                namespace,
                e
            );
            return Err(e);
        }

        if let Some(ref index_writer) = storage.index_writer
            && let Err(e) = index_writer.fsync().await
        {
            tracing::error!(
                "Failed to fsync index writer for partition {:?}: {}",
                namespace,
                e
            );
            return Err(e);
        }

        Ok(())
    }

    pub(crate) async fn auto_commit_consumer_offset_from_local_partition(
        &self,
        namespace: &IggyNamespace,
        consumer: PollingConsumer,
        offset: u64,
    ) -> Result<(), IggyError> {
        let (offset_value, path) = {
            let partitions = self.local_partitions.borrow();
            let partition = partitions.get(namespace).ok_or_else(|| {
                IggyError::PartitionNotFound(
                    namespace.partition_id(),
                    Identifier::numeric(namespace.topic_id() as u32).unwrap(),
                    Identifier::numeric(namespace.stream_id() as u32).unwrap(),
                )
            })?;

            match consumer {
                PollingConsumer::Consumer(consumer_id, _) => {
                    tracing::trace!(
                        "Auto-committing offset {} for consumer {} on partition {:?}",
                        offset,
                        consumer_id,
                        namespace
                    );
                    let hdl = partition.consumer_offsets.pin();
                    let item = hdl.get_or_insert(
                        consumer_id,
                        crate::streaming::partitions::consumer_offset::ConsumerOffset::default_for_consumer(
                            consumer_id as u32,
                            &self.config.system.get_consumer_offsets_path(
                                namespace.stream_id(),
                                namespace.topic_id(),
                                namespace.partition_id(),
                            ),
                        ),
                    );
                    item.offset.store(offset, Ordering::Release);
                    (item.offset.load(Ordering::Relaxed), item.path.clone())
                }
                PollingConsumer::ConsumerGroup(consumer_group_id, _) => {
                    tracing::trace!(
                        "Auto-committing offset {} for consumer group {} on partition {:?}",
                        offset,
                        consumer_group_id.0,
                        namespace
                    );
                    let hdl = partition.consumer_group_offsets.pin();
                    let item = hdl.get_or_insert(
                        consumer_group_id,
                        crate::streaming::partitions::consumer_offset::ConsumerOffset::default_for_consumer_group(
                            consumer_group_id,
                            &self.config.system.get_consumer_group_offsets_path(
                                namespace.stream_id(),
                                namespace.topic_id(),
                                namespace.partition_id(),
                            ),
                        ),
                    );
                    item.offset.store(offset, Ordering::Release);
                    (item.offset.load(Ordering::Relaxed), item.path.clone())
                }
            }
        };

        crate::streaming::partitions::storage::persist_offset(&path, offset_value).await?;

        self.maybe_complete_pending_revocation(
            &consumer,
            namespace.stream_id(),
            namespace.topic_id(),
            namespace.partition_id(),
        )
        .await;

        Ok(())
    }

    /// Appends a batch to the active segment, flushing to disk and rotating if needed.
    ///
    /// Safety: called exclusively from the message pump — segment indices captured before
    /// internal `.await` points (prepare_for_persistence, persist, rotate) remain valid
    /// because no other handler can modify the segment vec while this frame is in progress.
    pub(crate) async fn append_messages_to_local_partition(
        &self,
        namespace: &IggyNamespace,
        mut batch: IggyMessagesBatchMut,
        config: &crate::configs::system::SystemConfig,
    ) -> Result<(), IggyError> {
        let (
            current_offset,
            current_position,
            segment_start_offset,
            segment_index,
            message_deduplicator,
        ) = {
            let partitions = self.local_partitions.borrow();
            let partition = partitions
                .get(namespace)
                .expect("local_partitions: partition must exist");

            let current_offset = if partition.should_increment_offset {
                partition.offset.load(Ordering::Relaxed) + 1
            } else {
                0
            };

            let segment = partition.log.active_segment();
            let segment_index = partition.log.segments().len() - 1;

            (
                current_offset,
                segment.current_position,
                segment.start_offset,
                segment_index,
                partition.message_deduplicator.clone(),
            )
        };

        batch
            .prepare_for_persistence(
                segment_start_offset,
                current_offset,
                current_position,
                message_deduplicator.as_ref(),
            )
            .await;

        let (journal_messages_count, journal_size, is_full) = {
            let mut partitions = self.local_partitions.borrow_mut();
            let partition = partitions
                .get_mut(namespace)
                .expect("local_partitions: partition must exist");

            let segment = &mut partition.log.segments_mut()[segment_index];

            if segment.start_timestamp == 0 {
                segment.start_timestamp = batch.first_timestamp().unwrap();
            }

            let batch_messages_size = batch.size();
            let batch_messages_count = batch.count();

            partition
                .stats
                .increment_size_bytes(batch_messages_size as u64);
            partition
                .stats
                .increment_messages_count(batch_messages_count as u64);

            segment.end_timestamp = batch.last_timestamp().unwrap();
            segment.end_offset = batch.last_offset().unwrap();

            let (journal_messages_count, journal_size) =
                partition.log.journal_mut().append(batch)?;

            let last_offset = if batch_messages_count == 0 {
                current_offset
            } else {
                current_offset + batch_messages_count as u64 - 1
            };

            if partition.should_increment_offset {
                partition.offset.store(last_offset, Ordering::Relaxed);
            } else {
                partition.should_increment_offset = true;
                partition.offset.store(last_offset, Ordering::Relaxed);
            }
            partition.stats.set_current_offset(last_offset);
            partition.log.segments_mut()[segment_index].current_position += batch_messages_size;

            let is_full = partition.log.segments()[segment_index].is_full();

            (journal_messages_count, journal_size, is_full)
        };

        let unsaved_messages_count_exceeded =
            journal_messages_count >= config.partition.messages_required_to_save;
        let unsaved_messages_size_exceeded = journal_size
            >= config
                .partition
                .size_of_messages_required_to_save
                .as_bytes_u64() as u32;

        if is_full || unsaved_messages_count_exceeded || unsaved_messages_size_exceeded {
            let frozen_batches = {
                let mut partitions = self.local_partitions.borrow_mut();
                let partition = partitions
                    .get_mut(namespace)
                    .expect("local_partitions: partition must exist");
                let batches = partition.log.journal_mut().commit();
                partition.log.ensure_indexes();
                batches.append_indexes_to(partition.log.active_indexes_mut().unwrap());

                let frozen: Vec<_> = batches
                    .into_inner()
                    .into_iter()
                    .map(|mut b| b.freeze())
                    .collect();
                partition.log.set_in_flight(frozen.clone());
                frozen
            };

            self.persist_frozen_batches_to_disk(namespace, frozen_batches)
                .await?;

            if is_full {
                self.rotate_segment_in_local_partitions(namespace).await?;
            }
        }

        Ok(())
    }

    /// Persists already-frozen batches to disk. Caller must have set in_flight buffer.
    async fn persist_frozen_batches_to_disk(
        &self,
        namespace: &IggyNamespace,
        frozen_batches: Vec<iggy_common::IggyMessagesBatch>,
    ) -> Result<u32, IggyError> {
        let batch_count: u32 = frozen_batches.iter().map(|b| b.count()).sum();

        if batch_count == 0 {
            return Ok(0);
        }

        let (messages_writer, index_writer) = {
            let partitions = self.local_partitions.borrow();
            let partition = partitions
                .get(namespace)
                .expect("local_partitions: partition must exist");

            if !partition.log.has_segments() {
                return Ok(0);
            }

            let messages_writer = partition
                .log
                .active_storage()
                .messages_writer
                .as_ref()
                .expect("Messages writer not initialized")
                .clone();
            let index_writer = partition
                .log
                .active_storage()
                .index_writer
                .as_ref()
                .expect("Index writer not initialized")
                .clone();
            (messages_writer, index_writer)
        };

        let saved = messages_writer
            .as_ref()
            .save_frozen_batches(&frozen_batches)
            .await?;

        let unsaved_indexes_slice = {
            let partitions = self.local_partitions.borrow();
            let partition = partitions
                .get(namespace)
                .expect("local_partitions: partition must exist");
            let segment_index = partition.log.segments().len() - 1;
            partition.log.indexes()[segment_index]
                .as_ref()
                .expect("indexes must exist for segment being persisted")
                .unsaved_slice()
        };

        index_writer
            .as_ref()
            .save_indexes(unsaved_indexes_slice)
            .await?;

        tracing::trace!(
            "Persisted {} messages on disk for partition: {:?}, total bytes written: {}.",
            batch_count,
            namespace,
            saved
        );

        {
            let mut partitions = self.local_partitions.borrow_mut();
            let partition = partitions
                .get_mut(namespace)
                .expect("local_partitions: partition must exist");

            let segment_index = partition.log.segments().len() - 1;
            let indexes = partition.log.indexes_mut()[segment_index]
                .as_mut()
                .expect("indexes must exist for segment being persisted");
            indexes.mark_saved();

            let segment = &mut partition.log.segments_mut()[segment_index];
            segment.size =
                iggy_common::IggyByteSize::from(segment.size.as_bytes_u64() + saved.as_bytes_u64());

            partition.log.clear_in_flight();
        }

        Ok(batch_count)
    }

    pub(crate) async fn poll_messages_from_local_partition(
        &self,
        namespace: &IggyNamespace,
        consumer: crate::streaming::polling_consumer::PollingConsumer,
        args: PollingArgs,
    ) -> Result<(IggyPollMetadata, IggyMessagesBatchSet), IggyError> {
        crate::streaming::partitions::ops::poll_messages(
            &self.local_partitions,
            namespace,
            consumer,
            args,
        )
        .await
    }

    async fn decrypt_messages(
        &self,
        batches: IggyMessagesBatchSet,
        encryptor: &EncryptorKind,
    ) -> Result<IggyMessagesBatchSet, IggyError> {
        let mut decrypted_batches = Vec::with_capacity(batches.containers_count());
        for batch in batches.iter() {
            let mut indexes = IggyIndexesMut::with_capacity(batch.count() as usize, 0);
            let mut decrypted_messages = PooledBuffer::with_capacity(batch.size() as usize);
            let mut position = 0;

            for message in batch.iter() {
                let mut header = message.header().to_header();
                let offset = header.offset;
                let payload = encryptor.decrypt(message.payload());
                match payload {
                    Ok(payload) => {
                        // Update the header with the decrypted payload length
                        header.payload_length = payload.len() as u32;

                        // Decrypt user headers if present
                        let decrypted_user_headers = if let Some(user_headers) =
                            message.user_headers()
                        {
                            match encryptor.decrypt(user_headers) {
                                Ok(decrypted) => {
                                    header.user_headers_length = decrypted.len() as u32;
                                    Some(decrypted)
                                }
                                Err(error) => {
                                    error!(
                                        "Cannot decrypt the message user headers at offset: {offset}. Error: {error}"
                                    );
                                    continue;
                                }
                            }
                        } else {
                            None
                        };

                        decrypted_messages.extend_from_slice(&header.to_bytes());
                        decrypted_messages.extend_from_slice(&payload);
                        if let Some(ref user_headers) = decrypted_user_headers {
                            decrypted_messages.extend_from_slice(user_headers);
                        }
                        position += IGGY_MESSAGE_HEADER_SIZE
                            + payload.len()
                            + header.user_headers_length as usize;
                        indexes.insert(0, position as u32, 0);
                    }
                    Err(error) => {
                        error!("Cannot decrypt the message at offset: {offset}. Error: {error}",);
                        continue;
                    }
                }
            }
            let decrypted_batch =
                IggyMessagesBatchMut::from_indexes_and_messages(indexes, decrypted_messages);
            decrypted_batches.push(decrypted_batch);
        }

        Ok(IggyMessagesBatchSet::from_vec(decrypted_batches))
    }

    pub fn maybe_encrypt_messages(
        &self,
        batch: IggyMessagesBatchMut,
    ) -> Result<IggyMessagesBatchMut, IggyError> {
        let encryptor = match self.encryptor.as_ref() {
            Some(encryptor) => encryptor,
            None => return Ok(batch),
        };
        let mut encrypted_messages = PooledBuffer::with_capacity(batch.size() as usize * 2);
        let mut indexes = IggyIndexesMut::with_capacity(batch.count() as usize, 0);
        let mut position = 0;

        for message in batch.iter() {
            let header = message.header().to_header();
            let offset = header.offset;
            let payload_bytes = message.payload();
            let user_headers_bytes = message.user_headers();

            let encrypted_payload = encryptor.encrypt(payload_bytes);
            match encrypted_payload {
                Ok(encrypted_payload) => {
                    let mut updated_header = header;
                    updated_header.payload_length = encrypted_payload.len() as u32;

                    // Encrypt user headers if present
                    let encrypted_user_headers = if let Some(user_headers_bytes) =
                        user_headers_bytes
                    {
                        match encryptor.encrypt(user_headers_bytes) {
                            Ok(encrypted) => {
                                updated_header.user_headers_length = encrypted.len() as u32;
                                Some(encrypted)
                            }
                            Err(error) => {
                                error!(
                                    "Cannot encrypt the message user headers at offset: {offset}. Error: {error}"
                                );
                                continue;
                            }
                        }
                    } else {
                        None
                    };

                    encrypted_messages.extend_from_slice(&updated_header.to_bytes());
                    encrypted_messages.extend_from_slice(&encrypted_payload);
                    if let Some(ref encrypted_user_headers) = encrypted_user_headers {
                        encrypted_messages.extend_from_slice(encrypted_user_headers);
                    }
                    position += IGGY_MESSAGE_HEADER_SIZE
                        + encrypted_payload.len()
                        + updated_header.user_headers_length as usize;
                    indexes.insert(0, position as u32, 0);
                }
                Err(error) => {
                    error!("Cannot encrypt the message at offset: {offset}. Error: {error}",);
                    continue;
                }
            }
        }

        Ok(IggyMessagesBatchMut::from_indexes_and_messages(
            indexes,
            encrypted_messages,
        ))
    }
}

#[derive(Debug)]
pub struct PollingArgs {
    pub strategy: PollingStrategy,
    pub count: u32,
    pub auto_commit: bool,
}

impl PollingArgs {
    pub fn new(strategy: PollingStrategy, count: u32, auto_commit: bool) -> Self {
        Self {
            strategy,
            count,
            auto_commit,
        }
    }
}
