#![no_std]

extern crate alloc;

use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use alloc::format;
use alloc::string::{String, ToString};
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::fmt;
use core::future::Future;
use core::marker::PhantomData;
use core::pin::Pin;
use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use core::task::{Context, Poll};
use futures_util::stream::FuturesUnordered;
use futures_util::{FutureExt, StreamExt, future::BoxFuture, pin_mut, select_biased};
use smoo_host_core::control::{ConfigExportsV0, read_status};
use smoo_host_core::pump::{BulkInHandle, HostPumpHandle, RequestRx, start_host_pump};
use smoo_host_core::{
    BlockSource, BlockSourceError, BlockSourceErrorKind, BlockSourceHandle, ControlTransport,
    HostErrorKind, SmooHost, Transport, TransportError,
};
use smoo_proto::{ConfigExport, OpCode, Request, Response, SmooStatusV0};

pub use smoo_host_core::ExportIdentity;

pub const DEFAULT_HEARTBEAT_MISS_BUDGET: u32 = 5;

#[derive(Clone, Debug)]
pub struct HostSessionConfig {
    pub status_retry_attempts: usize,
}

impl Default for HostSessionConfig {
    fn default() -> Self {
        Self {
            status_retry_attempts: 5,
        }
    }
}

pub struct HostSession {
    sources: BTreeMap<u32, BlockSourceHandle>,
    config: HostSessionConfig,
}

impl HostSession {
    pub fn new(
        sources: BTreeMap<u32, BlockSourceHandle>,
        config: HostSessionConfig,
    ) -> Result<Self, HostSessionError> {
        if sources.is_empty() {
            return Err(HostSessionError::with_message(
                HostSessionErrorKind::InvalidConfiguration,
                "host session requires at least one export",
            ));
        }
        if config.status_retry_attempts == 0 {
            return Err(HostSessionError::with_message(
                HostSessionErrorKind::InvalidConfiguration,
                "status_retry_attempts must be greater than zero",
            ));
        }
        Ok(Self { sources, config })
    }

    pub async fn start<T, C>(
        self,
        transport: T,
        control: &mut C,
    ) -> Result<HostSessionTask<T>, HostSessionError>
    where
        T: Transport + Clone + Send + Sync + 'static,
        C: ControlTransport + Sync,
    {
        // SmooHost is used briefly for setup (IDENT + CONFIG_EXPORTS) over
        // ep0; we then drop it and drive the data plane through the pipelined
        // pump directly. The pump consumes the transport, so we clone before
        // setup so we still have one to hand off.
        let mut host = SmooHost::new(transport.clone(), self.sources.clone());
        let session_id = setup_session(&mut host, control, &self.sources, &self.config).await?;
        drop(host);

        let state = Arc::new(SessionState::new(session_id));
        let driver = Box::pin(run_session(self, transport, state.clone()));

        Ok(HostSessionTask {
            state,
            driver,
            _transport: PhantomData,
        })
    }
}

pub struct HostSessionTask<T>
where
    T: Transport + Clone + Send + Sync + 'static,
{
    state: Arc<SessionState>,
    driver: Pin<Box<dyn Future<Output = HostSessionFinish> + 'static>>,
    _transport: PhantomData<T>,
}

impl<T> Unpin for HostSessionTask<T> where T: Transport + Clone + Send + Sync + 'static {}

impl<T> HostSessionTask<T>
where
    T: Transport + Clone + Send + Sync + 'static,
{
    pub fn stop(&mut self) {
        self.state.stop_requested.store(true, Ordering::Relaxed);
    }

    pub async fn heartbeat<C>(&self, control: &mut C) -> Result<SmooStatusV0, HostSessionError>
    where
        C: ControlTransport + Sync,
    {
        let status = read_status(control).await.map_err(map_transport_error)?;
        let expected = self.state.expected_session_id.load(Ordering::Relaxed);
        if status.session_id != expected {
            self.state
                .observed_session_id
                .store(status.session_id, Ordering::Relaxed);
            self.state.session_changed.store(true, Ordering::Relaxed);
        }
        Ok(status)
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct HostSessionDriveConfig {
    pub heartbeat_miss_budget: u32,
}

impl Default for HostSessionDriveConfig {
    fn default() -> Self {
        Self {
            heartbeat_miss_budget: DEFAULT_HEARTBEAT_MISS_BUDGET,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum HostSessionDriveEvent {
    HeartbeatStatus {
        status: SmooStatusV0,
    },
    HeartbeatRecovered {
        missed_heartbeats: u32,
    },
    HeartbeatMiss {
        error: HostSessionError,
        missed_heartbeats: u32,
        budget: u32,
    },
    HeartbeatMissBudgetExhausted {
        missed_heartbeats: u32,
        budget: u32,
    },
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum HostSessionDriveOutcome {
    Shutdown,
    TransportLost,
    SessionChanged { previous: u64, current: u64 },
    Failed(HostSessionError),
}

pub async fn drive_host_session<T, C, Shutdown, Tick, TickFuture, OnEvent>(
    mut task: HostSessionTask<T>,
    mut control: C,
    shutdown: Shutdown,
    mut heartbeat_tick: Tick,
    config: HostSessionDriveConfig,
    mut on_event: OnEvent,
) -> HostSessionDriveOutcome
where
    T: Transport + Clone + Send + Sync + 'static,
    C: ControlTransport + Sync,
    Shutdown: Future<Output = ()>,
    Tick: FnMut() -> TickFuture,
    TickFuture: Future<Output = ()>,
    OnEvent: FnMut(HostSessionDriveEvent),
{
    let mut heartbeat_budget = match HeartbeatBudget::new(config.heartbeat_miss_budget) {
        Ok(budget) => budget,
        Err(err) => return HostSessionDriveOutcome::Failed(err),
    };

    let shutdown = shutdown.fuse();
    futures_util::pin_mut!(shutdown);

    loop {
        let heartbeat_tick = heartbeat_tick().fuse();
        futures_util::pin_mut!(heartbeat_tick);

        futures_util::select_biased! {
            _ = shutdown => {
                task.stop();
                return HostSessionDriveOutcome::Shutdown;
            }
            finish = (&mut task).fuse() => {
                return map_drive_finish(finish.outcome);
            }
            _ = heartbeat_tick => {
                match task.heartbeat(&mut control).await {
                    Ok(status) => {
                        if let Some(missed_heartbeats) = heartbeat_budget.record_success() {
                            on_event(HostSessionDriveEvent::HeartbeatRecovered {
                                missed_heartbeats,
                            });
                        }
                        on_event(HostSessionDriveEvent::HeartbeatStatus { status });
                    }
                    Err(error) => {
                        let update = heartbeat_budget.record_miss();
                        on_event(HostSessionDriveEvent::HeartbeatMiss {
                            error,
                            missed_heartbeats: update.missed_heartbeats,
                            budget: update.budget,
                        });
                        if update.exhausted {
                            on_event(HostSessionDriveEvent::HeartbeatMissBudgetExhausted {
                                missed_heartbeats: update.missed_heartbeats,
                                budget: update.budget,
                            });
                            return HostSessionDriveOutcome::TransportLost;
                        }
                    }
                }
            }
        }
    }
}

impl<T> Future for HostSessionTask<T>
where
    T: Transport + Clone + Send + Sync + 'static,
{
    type Output = HostSessionFinish;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // SAFETY: `HostSessionTask` does not project pinned fields by address.
        let this = unsafe { self.get_unchecked_mut() };
        this.driver.as_mut().poll(cx)
    }
}

pub struct HostSessionFinish {
    pub session: HostSession,
    pub outcome: Result<HostSessionOutcome, HostSessionError>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HostSessionOutcome {
    Stopped,
    TransportLost,
    SessionChanged { previous: u64, current: u64 },
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HostSessionErrorKind {
    InvalidConfiguration,
    SessionClosed,
    Transport,
    BlockSource,
    InvalidRequest,
    Unsupported,
    NotReady,
    Internal,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HostSessionError {
    kind: HostSessionErrorKind,
    message: Option<String>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct HeartbeatBudget {
    miss_budget: u32,
    missed_heartbeats: u32,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct HeartbeatBudgetUpdate {
    missed_heartbeats: u32,
    budget: u32,
    exhausted: bool,
}

impl HeartbeatBudget {
    fn new(miss_budget: u32) -> Result<Self, HostSessionError> {
        if miss_budget == 0 {
            return Err(HostSessionError::with_message(
                HostSessionErrorKind::InvalidConfiguration,
                "heartbeat_miss_budget must be greater than zero",
            ));
        }
        Ok(Self {
            miss_budget,
            missed_heartbeats: 0,
        })
    }

    fn record_success(&mut self) -> Option<u32> {
        if self.missed_heartbeats == 0 {
            return None;
        }
        let recovered = self.missed_heartbeats;
        self.missed_heartbeats = 0;
        Some(recovered)
    }

    fn record_miss(&mut self) -> HeartbeatBudgetUpdate {
        self.missed_heartbeats = self.missed_heartbeats.saturating_add(1);
        HeartbeatBudgetUpdate {
            missed_heartbeats: self.missed_heartbeats,
            budget: self.miss_budget,
            exhausted: self.missed_heartbeats >= self.miss_budget,
        }
    }
}

impl HostSessionError {
    pub const fn new(kind: HostSessionErrorKind) -> Self {
        Self {
            kind,
            message: None,
        }
    }

    pub fn with_message(kind: HostSessionErrorKind, message: impl Into<String>) -> Self {
        Self {
            kind,
            message: Some(message.into()),
        }
    }

    pub const fn kind(&self) -> HostSessionErrorKind {
        self.kind
    }

    pub fn message(&self) -> Option<&str> {
        self.message.as_deref()
    }
}

impl fmt::Display for HostSessionError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.message() {
            Some(msg) => write!(f, "{:?}: {msg}", self.kind),
            None => write!(f, "{:?}", self.kind),
        }
    }
}

struct SessionState {
    stop_requested: AtomicBool,
    session_changed: AtomicBool,
    expected_session_id: AtomicU64,
    observed_session_id: AtomicU64,
}

impl SessionState {
    fn new(session_id: u64) -> Self {
        Self {
            stop_requested: AtomicBool::new(false),
            session_changed: AtomicBool::new(false),
            expected_session_id: AtomicU64::new(session_id),
            observed_session_id: AtomicU64::new(session_id),
        }
    }
}

async fn run_session<T>(
    session: HostSession,
    transport: T,
    state: Arc<SessionState>,
) -> HostSessionFinish
where
    T: Transport + Clone + Send + Sync + 'static,
{
    let sources = Arc::new(session.sources.clone());
    let (handle, req_rx, pump_task) = start_host_pump(transport);

    let dispatcher = run_dispatcher(handle, req_rx, sources, state.clone()).fuse();
    let pump_task = pump_task.fuse();
    pin_mut!(dispatcher, pump_task);

    let outcome = loop {
        select_biased! {
            outcome = dispatcher => break outcome,
            _ = pump_task => {
                // Pump exited; the dispatcher's req_rx will see None on the
                // next iteration and resolve naturally, so just keep polling
                // it. This branch will fall through here once dispatcher
                // resolves.
            }
        }
    };

    HostSessionFinish { session, outcome }
}

async fn run_dispatcher(
    handle: HostPumpHandle,
    mut req_rx: RequestRx,
    sources: Arc<BTreeMap<u32, BlockSourceHandle>>,
    state: Arc<SessionState>,
) -> Result<HostSessionOutcome, HostSessionError> {
    let mut handlers: FuturesUnordered<BoxFuture<'static, ()>> = FuturesUnordered::new();

    loop {
        // Stop/session-change exits intentionally cancel active handlers. The
        // data-plane instance is being torn down, so dropping pump handles
        // closes channels and discards stale wire state.
        if state.stop_requested.load(Ordering::Relaxed) {
            return Ok(HostSessionOutcome::Stopped);
        }
        if state.session_changed.load(Ordering::Relaxed) {
            return Ok(observed_session_change(&state));
        }

        let next_req = req_rx.next().fuse();
        let next_done = handlers.select_next_some();
        pin_mut!(next_req, next_done);

        select_biased! {
            _ = next_done => {
                // Handler completed; nothing else to do here. Per-request
                // errors are logged inside the handler so the loop continues.
            }
            req_opt = next_req => {
                let Some(request) = req_opt else {
                    // Pump closed RequestRx — transport gone or peer shut down.
                    drain_handlers(&mut handlers).await;
                    return Ok(HostSessionOutcome::TransportLost);
                };
                let handler = build_handler(&handle, &sources, request).await?;
                handlers.push(handler);
            }
        }
    }
}

async fn drain_handlers(handlers: &mut FuturesUnordered<BoxFuture<'static, ()>>) {
    while handlers.next().await.is_some() {}
}

fn observed_session_change(state: &SessionState) -> HostSessionOutcome {
    let previous = state.expected_session_id.load(Ordering::Relaxed);
    let current = state.observed_session_id.load(Ordering::Relaxed);
    HostSessionOutcome::SessionChanged { previous, current }
}

/// Build the per-Request handler future, queueing any required bulk-IN
/// reservation in dispatch order before the handler is spawned.
async fn build_handler(
    handle: &HostPumpHandle,
    sources: &Arc<BTreeMap<u32, BlockSourceHandle>>,
    request: Request,
) -> Result<BoxFuture<'static, ()>, HostSessionError> {
    tracing::trace!(
        export_id = request.export_id,
        request_id = request.request_id,
        op = ?request.op,
        lba = request.lba,
        blocks = request.num_blocks,
        "session: dispatch Request"
    );

    let source = sources.get(&request.export_id).cloned();

    let bulk_in = if matches!(request.op, OpCode::Write) && request.num_blocks > 0 {
        let block_size = match &source {
            Some(s) => s.block_size(),
            None => {
                // Unknown export with bulk-IN payload coming. We have to
                // surface this rather than discard silently — without the
                // block_size we can't drain the bulk endpoint, and the wire
                // would desync. Promote to a session-level transport error
                // so the outer loop tears down and re-syncs via CONFIG_EXPORTS.
                return Err(HostSessionError::with_message(
                    HostSessionErrorKind::InvalidRequest,
                    format!(
                        "Write request for unknown export {} — cannot drain bulk-IN",
                        request.export_id
                    ),
                ));
            }
        };
        let byte_len = blocks_to_bytes(request.num_blocks, block_size).map_err(|err| {
            HostSessionError::with_message(HostSessionErrorKind::InvalidRequest, err.to_string())
        })?;
        if byte_len > 0 {
            let bulk = handle
                .queue_bulk_in_read(byte_len)
                .await
                .map_err(map_transport_error)?;
            Some(bulk)
        } else {
            None
        }
    } else {
        None
    };

    let pump = handle.clone();
    Ok(async move {
        handle_request(pump, source, request, bulk_in).await;
    }
    .boxed())
}

async fn handle_request(
    pump: HostPumpHandle,
    source: Option<BlockSourceHandle>,
    request: Request,
    bulk_in: Option<BulkInHandle>,
) {
    let result = match (source, request.op) {
        (None, _) => respond_invalid(&pump, request, ERRNO_EINVAL).await,
        (Some(source), OpCode::Read) => handle_read(&pump, &source, request).await,
        (Some(source), OpCode::Write) => handle_write(&pump, &source, request, bulk_in).await,
        (Some(source), OpCode::Flush) => handle_flush(&pump, &source, request).await,
        (Some(source), OpCode::Discard) => handle_discard(&pump, &source, request).await,
    };
    if let Err(err) = result {
        tracing::warn!(
            export_id = request.export_id,
            request_id = request.request_id,
            ?err,
            "session: handler error",
        );
    }
}

async fn handle_read(
    pump: &HostPumpHandle,
    source: &BlockSourceHandle,
    request: Request,
) -> Result<(), HandlerError> {
    let block_size = source.block_size();
    let byte_len =
        blocks_to_bytes(request.num_blocks, block_size).map_err(HandlerError::Invalid)?;

    if byte_len == 0 {
        let response = ok_response(&request, OpCode::Read);
        return pump
            .send_response_with_bulk(response, None)
            .await
            .map_err(HandlerError::Transport);
    }

    let mut buf = alloc::vec![0u8; byte_len];
    let read = match source.read_blocks(request.lba, &mut buf).await {
        Ok(n) => n,
        Err(err) => {
            let response = response_from_block_error(&request, &err);
            return pump
                .send_response_with_bulk(response, None)
                .await
                .map_err(HandlerError::Transport);
        }
    };
    if read != byte_len {
        let response = short_io_response(&request);
        return pump
            .send_response_with_bulk(response, None)
            .await
            .map_err(HandlerError::Transport);
    }

    let response = ok_response(&request, OpCode::Read);
    pump.send_response_with_bulk(response, Some(buf))
        .await
        .map_err(HandlerError::Transport)
}

async fn handle_write(
    pump: &HostPumpHandle,
    source: &BlockSourceHandle,
    request: Request,
    bulk_in: Option<BulkInHandle>,
) -> Result<(), HandlerError> {
    let block_size = source.block_size();
    // Matches the dispatch-time byte count used to reserve this bulk-IN read.
    let byte_len =
        blocks_to_bytes(request.num_blocks, block_size).map_err(HandlerError::Invalid)?;

    let buf = if byte_len > 0 {
        let handle = bulk_in.ok_or_else(|| {
            HandlerError::Invalid(format!(
                "Write request {} expected bulk_in handle but none was queued",
                request.request_id
            ))
        })?;
        handle
            .await_payload()
            .await
            .map_err(HandlerError::Transport)?
    } else {
        Vec::new()
    };

    if byte_len > 0 {
        match source.write_blocks(request.lba, &buf).await {
            Ok(written) if written == byte_len => {}
            Ok(_) => {
                let response = short_io_response(&request);
                return pump
                    .send_response_with_bulk(response, None)
                    .await
                    .map_err(HandlerError::Transport);
            }
            Err(err) => {
                let response = response_from_block_error(&request, &err);
                return pump
                    .send_response_with_bulk(response, None)
                    .await
                    .map_err(HandlerError::Transport);
            }
        }
    }

    let response = ok_response(&request, OpCode::Write);
    pump.send_response(response)
        .await
        .map_err(HandlerError::Transport)
}

async fn handle_flush(
    pump: &HostPumpHandle,
    source: &BlockSourceHandle,
    request: Request,
) -> Result<(), HandlerError> {
    let response = match source.flush().await {
        Ok(()) => ok_response(&request, OpCode::Flush),
        Err(err) => response_from_block_error(&request, &err),
    };
    pump.send_response(response)
        .await
        .map_err(HandlerError::Transport)
}

async fn handle_discard(
    pump: &HostPumpHandle,
    source: &BlockSourceHandle,
    request: Request,
) -> Result<(), HandlerError> {
    let response = match source.discard(request.lba, request.num_blocks).await {
        Ok(()) => ok_response(&request, OpCode::Discard),
        Err(err) => response_from_block_error(&request, &err),
    };
    pump.send_response(response)
        .await
        .map_err(HandlerError::Transport)
}

async fn respond_invalid(
    pump: &HostPumpHandle,
    request: Request,
    errno: u32,
) -> Result<(), HandlerError> {
    let response = Response::new(
        request.export_id,
        request.request_id,
        request.op,
        errno as u8,
        request.lba,
        request.num_blocks,
        0,
    );
    pump.send_response(response)
        .await
        .map_err(HandlerError::Transport)
}

#[derive(Debug)]
enum HandlerError {
    Transport(TransportError),
    Invalid(String),
}

impl fmt::Display for HandlerError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Transport(err) => write!(f, "transport: {err}"),
            Self::Invalid(msg) => write!(f, "invalid: {msg}"),
        }
    }
}

const ERRNO_EINVAL: u32 = 22;
const ERRNO_EIO: u32 = 5;
const ERRNO_EOPNOTSUPP: u32 = 95;

fn ok_response(request: &Request, op: OpCode) -> Response {
    Response::new(
        request.export_id,
        request.request_id,
        op,
        0,
        request.lba,
        request.num_blocks,
        0,
    )
}

fn short_io_response(request: &Request) -> Response {
    Response::new(
        request.export_id,
        request.request_id,
        request.op,
        ERRNO_EIO as u8,
        request.lba,
        request.num_blocks,
        0,
    )
}

fn response_from_block_error(request: &Request, err: &BlockSourceError) -> Response {
    let errno = match err.kind() {
        BlockSourceErrorKind::InvalidInput | BlockSourceErrorKind::OutOfRange => ERRNO_EINVAL,
        BlockSourceErrorKind::Unsupported => ERRNO_EOPNOTSUPP,
        BlockSourceErrorKind::Io | BlockSourceErrorKind::Other => ERRNO_EIO,
    };
    Response::new(
        request.export_id,
        request.request_id,
        request.op,
        errno as u8,
        request.lba,
        request.num_blocks,
        0,
    )
}

fn blocks_to_bytes(num_blocks: u32, block_size: u32) -> Result<usize, String> {
    if block_size == 0 {
        return Err("block size is zero".into());
    }
    num_blocks
        .checked_mul(block_size)
        .map(|v| v as usize)
        .ok_or_else(|| "request size overflow".into())
}

async fn setup_session<T, C>(
    host: &mut SmooHost<T, BlockSourceHandle>,
    control: &mut C,
    sources: &BTreeMap<u32, BlockSourceHandle>,
    config: &HostSessionConfig,
) -> Result<u64, HostSessionError>
where
    T: Transport + Send + Sync,
    C: ControlTransport + Sync,
{
    let _ident = host.setup(control).await.map_err(map_host_error)?;

    let payload = build_config_payload(sources).await?;
    host.configure_exports_v0(control, &payload)
        .await
        .map_err(map_host_error)?;

    fetch_status_with_retry(control, config.status_retry_attempts).await
}

async fn build_config_payload(
    sources: &BTreeMap<u32, BlockSourceHandle>,
) -> Result<ConfigExportsV0, HostSessionError> {
    let mut entries = Vec::with_capacity(sources.len());
    for (export_id, source) in sources {
        let block_size = source.block_size();
        if block_size == 0 {
            return Err(HostSessionError::with_message(
                HostSessionErrorKind::InvalidConfiguration,
                "block size must be non-zero",
            ));
        }
        let total_blocks = source.total_blocks().await.map_err(map_block_error)?;
        let size_bytes = total_blocks.checked_mul(block_size as u64).ok_or_else(|| {
            HostSessionError::with_message(
                HostSessionErrorKind::InvalidConfiguration,
                "export size overflow",
            )
        })?;
        entries.push(ConfigExport {
            export_id: *export_id,
            block_size,
            size_bytes,
        });
    }
    ConfigExportsV0::from_slice(&entries).map_err(|err| {
        HostSessionError::with_message(HostSessionErrorKind::InvalidConfiguration, err.to_string())
    })
}

async fn fetch_status_with_retry<C>(
    control: &mut C,
    attempts: usize,
) -> Result<u64, HostSessionError>
where
    C: ControlTransport + Sync,
{
    let mut remaining = attempts;
    loop {
        match read_status(control).await {
            Ok(status) => return Ok(status.session_id),
            Err(err) => {
                remaining = remaining.saturating_sub(1);
                if remaining == 0 {
                    return Err(map_transport_error(err));
                }
            }
        }
    }
}

fn map_drive_finish(
    outcome: Result<HostSessionOutcome, HostSessionError>,
) -> HostSessionDriveOutcome {
    match outcome {
        Ok(HostSessionOutcome::Stopped) => HostSessionDriveOutcome::Shutdown,
        Ok(HostSessionOutcome::TransportLost) => HostSessionDriveOutcome::TransportLost,
        Ok(HostSessionOutcome::SessionChanged { previous, current }) => {
            HostSessionDriveOutcome::SessionChanged { previous, current }
        }
        Err(err) => HostSessionDriveOutcome::Failed(err),
    }
}

fn map_transport_error(err: smoo_host_core::TransportError) -> HostSessionError {
    HostSessionError::with_message(HostSessionErrorKind::Transport, err.to_string())
}

fn map_block_error(err: smoo_host_core::BlockSourceError) -> HostSessionError {
    HostSessionError::with_message(HostSessionErrorKind::BlockSource, err.to_string())
}

fn map_host_error(err: smoo_host_core::HostError) -> HostSessionError {
    let kind = match err.kind() {
        HostErrorKind::Transport => HostSessionErrorKind::Transport,
        HostErrorKind::BlockSource => HostSessionErrorKind::BlockSource,
        HostErrorKind::InvalidRequest => HostSessionErrorKind::InvalidRequest,
        HostErrorKind::Unsupported => HostSessionErrorKind::Unsupported,
        HostErrorKind::NotReady => HostSessionErrorKind::NotReady,
    };
    HostSessionError::with_message(kind, err.to_string())
}

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

    #[test]
    fn heartbeat_budget_rejects_zero_budget() {
        let err = HeartbeatBudget::new(0).expect_err("zero budget must fail");
        assert_eq!(err.kind(), HostSessionErrorKind::InvalidConfiguration);
    }

    #[test]
    fn heartbeat_budget_counts_and_resets() {
        let mut budget = HeartbeatBudget::new(3).expect("budget");

        let miss1 = budget.record_miss();
        assert_eq!(
            miss1,
            HeartbeatBudgetUpdate {
                missed_heartbeats: 1,
                budget: 3,
                exhausted: false,
            }
        );

        let miss2 = budget.record_miss();
        assert_eq!(
            miss2,
            HeartbeatBudgetUpdate {
                missed_heartbeats: 2,
                budget: 3,
                exhausted: false,
            }
        );

        assert_eq!(budget.record_success(), Some(2));
        assert_eq!(budget.record_success(), None);

        let miss3 = budget.record_miss();
        assert_eq!(
            miss3,
            HeartbeatBudgetUpdate {
                missed_heartbeats: 1,
                budget: 3,
                exhausted: false,
            }
        );
    }

    #[test]
    fn heartbeat_budget_detects_exhaustion() {
        let mut budget = HeartbeatBudget::new(2).expect("budget");
        assert!(!budget.record_miss().exhausted);
        assert!(budget.record_miss().exhausted);
    }
}
