// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;
use anyhow;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Action {
    GetNumNewMatchedTokens(GetNumNewMatchedTokensInput, GetNumNewMatchedTokensOutput),
    UpdateStateAfterAlloc(UpdateStateAfterAllocInput, UpdateStateAfterAllocOutput),
    BuildConnectorMeta(BuildConnectorMetaInput, BuildConnectorMetaOutput),
    RequestFinished(RequestFinishedInput, RequestFinishedOutput),
    HasSlot(HasSlotInput, HasSlotOutput),
    CreateSlot(CreateSlotInput, CreateSlotOutput),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GetNumNewMatchedTokensInput {
    request_id: String,
    request_num_tokens: usize,
    num_computed_tokens: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GetNumNewMatchedTokensOutput {
    num_new_matched_tokens: usize,
    has_matched: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateStateAfterAllocInput {
    request_id: String,
    block_ids: Vec<BlockId>,
    num_external_tokens: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateStateAfterAllocOutput {}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BuildConnectorMetaInput {
    scheduler_output: SchedulerOutput,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BuildConnectorMetaOutput {
    metadata: ConnectorMetadata,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RequestFinishedInput {
    request_id: String,
    block_ids: Vec<BlockId>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RequestFinishedOutput {
    is_finished: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HasSlotInput {
    request_id: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HasSlotOutput {
    result: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateSlotInput {
    request: KvbmRequest,
    tokens: Vec<u32>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateSlotOutput {}

#[derive(Debug)]
pub struct KvConnectorLeaderRecorder {
    _recorder: Recorder<Action>, // Keep recorder alive
    unbounded_tx: mpsc::UnboundedSender<Action>,
    connector_leader: Box<dyn Leader>,
}

impl KvConnectorLeaderRecorder {
    pub fn new(
        worker_id: String,
        drt: PyDistributedRuntime,
        block_manager: PyBlockManager,
        leader: PyKvbmLeader,
    ) -> Self {
        tracing::info!(
            "KvConnectorLeaderRecorder initialized with worker_id: {}",
            worker_id
        );

        // if drt is none, then we must construct a runtime and distributed runtime
        let block_manager = block_manager.get_block_manager().clone();
        let block_size = block_manager.block_size();

        let leader = leader.get_inner();

        // if we need a drt, get it from here
        let drt = drt.inner().clone();

        let token = CancellationToken::new();
        let output_path = "/tmp/records.jsonl";
        tracing::info!("recording events to {}", output_path);

        let ns = drt.namespace("kvbm_connector_leader").unwrap();

        let kvbm_metrics = KvbmMetrics::new(&ns);

        let recorder = drt
            .runtime()
            .primary()
            .block_on(async { Recorder::new(token, &output_path, None, None, None).await })
            .unwrap();

        let connector_leader = KvConnectorLeader {
            slot_manager: ConnectorSlotManager::new(
                block_manager.clone(),
                leader,
                drt.clone(),
                kvbm_metrics.clone(),
            ),
            block_size,
            inflight_requests: HashSet::new(),
            onboarding_slots: HashSet::new(),
            iteration_counter: 0,
            kvbm_metrics,
        };

        let (unbounded_tx, unbounded_rx) = mpsc::unbounded_channel();
        let recorder_tx = recorder.event_sender();

        // todo(kvbm): make this a critical task
        drt.runtime()
            .primary()
            .spawn(Self::forward_unbounded_to_sender(unbounded_rx, recorder_tx));

        Self {
            _recorder: recorder,
            unbounded_tx,
            connector_leader: Box::new(connector_leader),
        }
    }

    async fn forward_unbounded_to_sender<T: Send + 'static>(
        mut unbounded_rx: mpsc::UnboundedReceiver<T>,
        bounded_tx: mpsc::Sender<T>,
    ) {
        while let Some(msg) = unbounded_rx.recv().await {
            if bounded_tx.send(msg).await.is_err() {
                tracing::error!("Failed to send message to bounded channel");
            }
        }
    }
}

impl Leader for KvConnectorLeaderRecorder {
    /// Match the tokens in the request with the available block pools.
    /// Note: the necessary details of the request are captured prior to this call. For vllm,
    /// we make a create slot call prior to this call, so a slot is guaranteed to exist.
    ///
    /// To align with the connector interface, we must ensure that if no blocks are matched, we return (0, false).
    /// In our implementation, if we match any block, we return (num_matched_tokens, true).
    fn get_num_new_matched_tokens(
        &self,
        request_id: String,
        request_num_tokens: usize,
        num_computed_tokens: usize,
    ) -> anyhow::Result<(usize, bool)> {
        let input_copy = GetNumNewMatchedTokensInput {
            request_id: request_id.clone(),
            request_num_tokens,
            num_computed_tokens,
        };
        let output = self.connector_leader.get_num_new_matched_tokens(
            request_id,
            request_num_tokens,
            num_computed_tokens,
        )?;
        let _ = self.unbounded_tx.send(Action::GetNumNewMatchedTokens(
            input_copy,
            GetNumNewMatchedTokensOutput {
                num_new_matched_tokens: output.0,
                has_matched: output.1,
            },
        ));
        Ok(output)
    }

    /// We drop the need to pass in the KvCacheBlocks and the num_external_tokens as they are captured
    /// statefully in the [`VllmLeaderKvCacheManagerAndConnector::get_num_new_matched_tokens`] function.
    ///
    /// Note: vLLM will not provide any scheduler output data for requests that are onboarding. it is entirely
    /// on the connector's implementation to handle this case.
    fn update_state_after_alloc(
        &mut self,
        request_id: String,
        block_ids: Vec<BlockId>,
        num_external_tokens: usize,
    ) -> anyhow::Result<()> {
        let input_copy = UpdateStateAfterAllocInput {
            request_id: request_id.clone(),
            block_ids: block_ids.clone(),
            num_external_tokens,
        };
        self.connector_leader.update_state_after_alloc(
            request_id,
            block_ids,
            num_external_tokens,
        )?;
        let _ = self.unbounded_tx.send(Action::UpdateStateAfterAlloc(
            input_copy,
            UpdateStateAfterAllocOutput {},
        ));
        Ok(())
    }

    fn build_connector_metadata(
        &mut self,
        scheduler_output: SchedulerOutput,
    ) -> anyhow::Result<Vec<u8>> {
        let input_copy = BuildConnectorMetaInput {
            scheduler_output: scheduler_output.clone(),
        };
        let output = self
            .connector_leader
            .build_connector_metadata(scheduler_output)?;
        let _ = self.unbounded_tx.send(Action::BuildConnectorMeta(
            input_copy,
            BuildConnectorMetaOutput {
                metadata: serde_json::from_slice(&output)?,
            },
        ));
        Ok(output)
    }

    fn request_finished(
        &mut self,
        request_id: String,
        block_ids: Vec<BlockId>,
    ) -> anyhow::Result<bool> {
        let input_copy = RequestFinishedInput {
            request_id: request_id.clone(),
            block_ids: block_ids.clone(),
        };
        let output = self
            .connector_leader
            .request_finished(request_id, block_ids)?;
        let _ = self.unbounded_tx.send(Action::RequestFinished(
            input_copy,
            RequestFinishedOutput {
                is_finished: output,
            },
        ));
        Ok(output)
    }

    fn has_slot(&self, request_id: String) -> bool {
        let input_copy = HasSlotInput {
            request_id: request_id.clone(),
        };
        let output = self.connector_leader.has_slot(request_id);
        let _ = self.unbounded_tx.send(Action::HasSlot(
            input_copy,
            HasSlotOutput { result: output },
        ));
        output
    }

    /// Create a new slot for the given request ID.
    /// This is used to create a new slot for the request.
    fn create_slot(&mut self, request: KvbmRequest, tokens: Vec<u32>) -> anyhow::Result<()> {
        let input_copy = CreateSlotInput {
            request: request.clone(),
            tokens: tokens.clone(),
        };
        let _ = self.connector_leader.create_slot(request, tokens);
        let _ = self
            .unbounded_tx
            .send(Action::CreateSlot(input_copy, CreateSlotOutput {}));
        Ok(())
    }
}
