use crate::{CompiledModel, Error, ModelInput, ModelPhase, ModelStatus, Result};
use calculet_calrt::{
    Calbin, ConfiguredRuntime, DeviceIo, Model, PrimitiveType, TaskSlot, Tensor, TensorBuffer,
};
use candle_core::{Device, Tensor as CandleTensor};
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::collections::BTreeMap;

const TOKEN_INPUT: &str = "inputs[0]";
const POSITION_INPUT: &str = "inputs[1]";
const LOGITS_OUTPUT: &str = "outputs[0]";
const CUR_SEQUENCE_CSR: &str = "cur_seq_len[0]";
const PAST_AND_CURRENT_CSR: &str = "past_kv_cur_seq_len[0]";

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct CalrtModelPlan {
    pub prefill_model: String,
    pub decode_model: String,
    pub vocabulary_size: usize,
    pub maximum_sequence_length: usize,
    pub maximum_batch_size: usize,
}

impl CalrtModelPlan {
    pub fn from_calbin(calbin: &Calbin) -> Result<Self> {
        let prefill = select_prefill(&calbin.models)
            .ok_or_else(|| Error::InvalidConfiguration("Calbin has no prefill submodel".into()))?;
        let decode = calbin
            .models
            .iter()
            .find(|model| {
                let name = model.name.to_ascii_lowercase();
                name.contains("decode") && !name.contains("prefill")
            })
            .ok_or_else(|| Error::InvalidConfiguration("Calbin has no decode submodel".into()))?;
        validate_model_contract(prefill)?;
        validate_model_contract(decode)?;
        let vocabulary_size = output_vocabulary_size(decode)?;
        let prefill_vocabulary = output_vocabulary_size(prefill)?;
        if prefill_vocabulary != vocabulary_size {
            return Err(Error::VocabularyMismatch {
                expected: vocabulary_size,
                actual: prefill_vocabulary,
            });
        }
        let maximum_sequence_length = usize::try_from(calbin.llm.max_sequence_length)
            .map_err(|_| Error::InvalidConfiguration("max_seq_len exceeds usize".into()))?;
        if maximum_sequence_length == 0 {
            return Err(Error::InvalidConfiguration(
                "Calbin does not declare max_seq_len".into(),
            ));
        }
        let maximum_batch_size = usize::try_from(calbin.llm.max_batch_size)
            .map_err(|_| Error::InvalidConfiguration("n_batch exceeds usize".into()))?;
        if maximum_batch_size == 0 {
            return Err(Error::InvalidConfiguration(
                "Calbin does not declare n_batch".into(),
            ));
        }
        Ok(Self {
            prefill_model: prefill.name.clone(),
            decode_model: decode.name.clone(),
            vocabulary_size,
            maximum_sequence_length,
            maximum_batch_size,
        })
    }
}

pub struct CalrtCompiledModel<D> {
    runtime: ConfiguredRuntime<D>,
    plan: CalrtModelPlan,
    next_slot: TaskSlot,
    sequence_lengths: BTreeMap<u32, usize>,
}

impl<D: std::fmt::Debug> std::fmt::Debug for CalrtCompiledModel<D> {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("CalrtCompiledModel")
            .field("runtime", &self.runtime)
            .field("plan", &self.plan)
            .field("next_slot", &self.next_slot)
            .field("sequence_lengths", &self.sequence_lengths)
            .finish()
    }
}

impl<D: DeviceIo> CalrtCompiledModel<D> {
    pub fn new(runtime: ConfiguredRuntime<D>) -> Result<Self> {
        let plan = CalrtModelPlan::from_calbin(runtime.calbin())?;
        Ok(Self {
            runtime,
            plan,
            next_slot: TaskSlot::Ping,
            sequence_lengths: BTreeMap::new(),
        })
    }

    pub fn deploy_parameters(&mut self) -> Result<()> {
        self.runtime.deploy_parameters()?;
        Ok(())
    }

    pub fn plan(&self) -> &CalrtModelPlan {
        &self.plan
    }

    pub fn runtime(&self) -> &ConfiguredRuntime<D> {
        &self.runtime
    }

    pub fn runtime_mut(&mut self) -> &mut ConfiguredRuntime<D> {
        &mut self.runtime
    }

    fn run(&mut self, input: &ModelInput) -> Result<CandleTensor> {
        validate_sequence_input(input, self.plan.maximum_sequence_length)?;
        let known_length = self
            .sequence_lengths
            .get(&input.sequence_id)
            .copied()
            .unwrap_or_default();
        if input.past_tokens != known_length {
            return Err(Error::TensorContract(format!(
                "sequence {} expected past length {known_length}, received {}",
                input.sequence_id, input.past_tokens
            )));
        }

        let model_name = match input.phase {
            ModelPhase::Prefill => &self.plan.prefill_model,
            ModelPhase::Decode => &self.plan.decode_model,
        };
        let mut inputs = self.runtime.input_buffer(model_name)?;
        validate_csr_contract(&inputs)?;
        fill_i32(inputs.tensor_mut(TOKEN_INPUT)?, &input.token_ids)?;
        fill_i32(inputs.tensor_mut(POSITION_INPUT)?, &input.positions)?;

        let mut outputs = self.runtime.output_buffer(model_name)?;
        self.runtime.write_inputs(&inputs, self.next_slot)?;

        // The pure-Rust CALRT intentionally returns HardwareExecutionUnavailable here until
        // CCU relocation, launch, job queue completion, and device KV reset are implemented.
        // Keeping this call in the real path makes the remaining boundary explicit.
        self.runtime.submit(model_name, self.next_slot)?;
        self.runtime.read_outputs(&mut outputs, self.next_slot)?;
        let logits = decode_logits(
            outputs.tensor(LOGITS_OUTPUT)?,
            input,
            self.plan.vocabulary_size,
        )?;

        self.sequence_lengths.insert(
            input.sequence_id,
            input
                .past_tokens
                .checked_add(input.token_ids.len())
                .ok_or_else(|| Error::TensorContract("sequence length overflow".into()))?,
        );
        self.next_slot = match self.next_slot {
            TaskSlot::Ping => TaskSlot::Pong,
            TaskSlot::Pong => TaskSlot::Ping,
        };
        Ok(logits)
    }
}

impl<D: DeviceIo> CompiledModel for CalrtCompiledModel<D> {
    fn status(&self) -> ModelStatus {
        ModelStatus {
            ready: false,
            reason: "Rust CALRT still lacks CCU launch, job completion, and device KV reset".into(),
            details: json!({
                "prefill_model": self.plan.prefill_model,
                "decode_model": self.plan.decode_model,
                "maximum_sequence_length": self.plan.maximum_sequence_length,
                "maximum_batch_size": self.plan.maximum_batch_size,
                "vocabulary_size": self.plan.vocabulary_size,
                "host_pipeline_ready": true,
                "hardware_submission_ready": false,
            }),
        }
    }

    fn vocabulary_size(&self) -> usize {
        self.plan.vocabulary_size
    }

    fn maximum_sequence_length(&self) -> usize {
        self.plan.maximum_sequence_length
    }

    fn reset_sequence(&mut self, sequence_id: u32) -> Result<()> {
        self.sequence_lengths.remove(&sequence_id);
        Ok(())
    }

    fn forward(&mut self, input: &ModelInput) -> Result<CandleTensor> {
        self.run(input)
    }
}

fn select_prefill(models: &[Model]) -> Option<&Model> {
    models
        .iter()
        .filter(|model| model.name.to_ascii_lowercase().contains("prefill"))
        .min_by_key(|model| {
            let name = model.name.to_ascii_lowercase();
            !name.contains("by_ids")
        })
}

fn validate_model_contract(model: &Model) -> Result<()> {
    let input_names = model
        .input_tensors()
        .map(|tensor| tensor.name.as_str())
        .collect::<Vec<_>>();
    if !input_names.contains(&TOKEN_INPUT) || !input_names.contains(&POSITION_INPUT) {
        return Err(Error::TensorContract(format!(
            "model {:?} must expose {TOKEN_INPUT} and {POSITION_INPUT}",
            model.name
        )));
    }
    if !model
        .output_tensors()
        .any(|tensor| tensor.name == LOGITS_OUTPUT)
    {
        return Err(Error::TensorContract(format!(
            "model {:?} must expose {LOGITS_OUTPUT}",
            model.name
        )));
    }
    let csr_names = model
        .csr_offsets()
        .map(|(name, _)| name)
        .collect::<Vec<_>>();
    if !csr_names.contains(&CUR_SEQUENCE_CSR) || !csr_names.contains(&PAST_AND_CURRENT_CSR) {
        return Err(Error::TensorContract(format!(
            "model {:?} lacks required LLM CSR fields",
            model.name
        )));
    }
    Ok(())
}

fn validate_csr_contract(inputs: &TensorBuffer) -> Result<()> {
    if inputs.csr(CUR_SEQUENCE_CSR).is_none() || inputs.csr(PAST_AND_CURRENT_CSR).is_none() {
        return Err(Error::TensorContract(
            "CALRT input buffer lacks required LLM CSR offsets".into(),
        ));
    }
    Ok(())
}

fn output_vocabulary_size(model: &Model) -> Result<usize> {
    let output = model
        .output_tensors()
        .find(|tensor| tensor.name == LOGITS_OUTPUT)
        .ok_or_else(|| Error::TensorContract(format!("model {:?} has no logits", model.name)))?;
    match output.shape.as_slice() {
        [_, _, vocabulary] => usize::try_from(*vocabulary)
            .map_err(|_| Error::TensorContract("vocabulary exceeds usize".into())),
        [_, vocabulary_blocks, _, _] => {
            let bytes = usize::from(output.data_type.bit_size()).div_ceil(8);
            if bytes == 0 || 512 % (16 * bytes) != 0 {
                return Err(Error::TensorContract(
                    "unsupported tiled logits primitive type".into(),
                ));
            }
            usize::try_from(*vocabulary_blocks)
                .ok()
                .and_then(|blocks| blocks.checked_mul(512 / 16 / bytes))
                .ok_or_else(|| Error::TensorContract("vocabulary size overflow".into()))
        }
        shape => Err(Error::TensorContract(format!(
            "unsupported logits shape {shape:?}"
        ))),
    }
}

fn validate_sequence_input(input: &ModelInput, maximum: usize) -> Result<()> {
    if input.token_ids.is_empty() {
        return Err(Error::EmptyPrompt);
    }
    if input.token_ids.len() != input.positions.len() {
        return Err(Error::TensorContract(
            "token and position counts differ".into(),
        ));
    }
    let requested =
        input
            .past_tokens
            .checked_add(input.token_ids.len())
            .ok_or(Error::ContextOverflow {
                requested: usize::MAX,
                maximum,
            })?;
    if requested > maximum {
        return Err(Error::ContextOverflow { requested, maximum });
    }
    if input.phase == ModelPhase::Decode && input.token_ids.len() != 1 {
        return Err(Error::TensorContract(
            "decode accepts exactly one token".into(),
        ));
    }
    Ok(())
}

fn fill_i32(tensor: &mut Tensor, values: &[u32]) -> Result<()> {
    if !matches!(
        tensor.info().data_type,
        PrimitiveType::S32 | PrimitiveType::U32 | PrimitiveType::Token
    ) {
        return Err(Error::TensorContract(format!(
            "tensor {:?} must use a 32-bit token type, found {}",
            tensor.info().name,
            tensor.info().data_type
        )));
    }
    let byte_size = values
        .len()
        .checked_mul(size_of::<u32>())
        .ok_or_else(|| Error::TensorContract("input size overflow".into()))?;
    if byte_size > tensor.data().len() {
        return Err(Error::TensorContract(format!(
            "tensor {:?} has {} bytes but needs {byte_size}",
            tensor.info().name,
            tensor.data().len()
        )));
    }
    tensor.data_mut().fill(0);
    for (target, value) in tensor.data_mut()[..byte_size]
        .chunks_exact_mut(size_of::<u32>())
        .zip(values)
    {
        target.copy_from_slice(&value.to_le_bytes());
    }
    tensor.slice(0, byte_size)?;
    Ok(())
}

fn decode_logits(tensor: &Tensor, input: &ModelInput, vocabulary: usize) -> Result<CandleTensor> {
    let info = tensor.info();
    let element_bytes = usize::from(info.data_type.bit_size()).div_ceil(8);
    if !matches!(info.data_type, PrimitiveType::Bf16 | PrimitiveType::F32) {
        return Err(Error::TensorContract(format!(
            "logits must be bf16 or f32, found {}",
            info.data_type
        )));
    }
    let indices = match info.shape.as_slice() {
        [_, _, declared_vocabulary] => {
            let declared = usize::try_from(*declared_vocabulary)
                .map_err(|_| Error::TensorContract("vocabulary exceeds usize".into()))?;
            if declared < vocabulary {
                return Err(Error::VocabularyMismatch {
                    expected: vocabulary,
                    actual: declared,
                });
            }
            (0..vocabulary).collect::<Vec<_>>()
        }
        [sequence_blocks, vocabulary_blocks, _, _] => {
            let d0 = 16_usize;
            let d2 = 512_usize
                .checked_div(d0 * element_bytes)
                .ok_or_else(|| Error::TensorContract("invalid tiled logits type".into()))?;
            let sequence_length = usize::try_from(*sequence_blocks)
                .ok()
                .and_then(|blocks| blocks.checked_mul(d0))
                .ok_or_else(|| Error::TensorContract("sequence shape overflow".into()))?;
            let row = input.token_ids.len() - 1;
            if row >= sequence_length {
                return Err(Error::TensorContract(
                    "prefill row exceeds tiled logits shape".into(),
                ));
            }
            let declared = usize::try_from(*vocabulary_blocks)
                .ok()
                .and_then(|blocks| blocks.checked_mul(d2))
                .ok_or_else(|| Error::TensorContract("vocabulary shape overflow".into()))?;
            if declared < vocabulary {
                return Err(Error::VocabularyMismatch {
                    expected: vocabulary,
                    actual: declared,
                });
            }
            let row_block = row / d0;
            let row_inside = row % d0;
            let row_stride = vocabulary
                .checked_mul(d0)
                .ok_or_else(|| Error::TensorContract("tiled row stride overflow".into()))?;
            let block_stride = d0 * d2;
            let base = row_block
                .checked_mul(row_stride)
                .and_then(|offset| offset.checked_add(row_inside * d2))
                .ok_or_else(|| Error::TensorContract("tiled logits offset overflow".into()))?;
            (0..vocabulary)
                .map(|column| base + (column / d2) * block_stride + column % d2)
                .collect()
        }
        shape => {
            return Err(Error::TensorContract(format!(
                "unsupported logits shape {shape:?}"
            )));
        }
    };
    let mut logits = Vec::with_capacity(vocabulary);
    for index in indices {
        let offset = index
            .checked_mul(element_bytes)
            .ok_or_else(|| Error::TensorContract("logits byte offset overflow".into()))?;
        let bytes = tensor
            .data()
            .get(offset..offset + element_bytes)
            .ok_or_else(|| Error::TensorContract("logits buffer is truncated".into()))?;
        logits.push(match info.data_type {
            PrimitiveType::Bf16 => {
                let bits = u16::from_le_bytes([bytes[0], bytes[1]]);
                f32::from_bits(u32::from(bits) << 16)
            }
            PrimitiveType::F32 => f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]),
            _ => unreachable!(),
        });
    }
    Ok(CandleTensor::from_vec(logits, vocabulary, &Device::Cpu)?)
}

#[cfg(test)]
mod tests {
    use super::*;
    use calculet_calrt::{BufferDirection, TensorInfo};
    use std::path::Path;

    #[test]
    fn finds_the_captured_qwen_prefill_and_decode_contracts() {
        let root = Path::new(env!("CARGO_MANIFEST_DIR")).join(
            "../../npu_features/snapshot_20260801/remote/data/models/\
             Qwen3-30B-A3B-dynamic-W8A8-W4AF16-full_layers_merged_2_chips_40960_fa_2026-05-22",
        );
        let calbin = Calbin::load(root).unwrap();
        let plan = CalrtModelPlan::from_calbin(&calbin).unwrap();
        assert!(plan.prefill_model.contains("prefill"));
        assert!(plan.decode_model.contains("decode"));
        assert_eq!(plan.vocabulary_size, 151_936);
        assert_eq!(plan.maximum_sequence_length, 40_960);
    }

    #[test]
    fn converts_calrt_bf16_logits_to_a_candle_tensor() {
        let mut tensor = Tensor::new(
            TensorInfo {
                name: LOGITS_OUTPUT.into(),
                shape: vec![1, 1, 3],
                data_type: PrimitiveType::Bf16,
                ping_address: 0x1000,
                pong_address: 0x2000,
                byte_size: 6,
            },
            BufferDirection::DeviceToHost,
        )
        .unwrap();
        let values = [1.0_f32, -2.0, 3.5];
        for (target, value) in tensor.data_mut().chunks_exact_mut(2).zip(values) {
            target.copy_from_slice(&((value.to_bits() >> 16) as u16).to_le_bytes());
        }
        let logits = decode_logits(
            &tensor,
            &ModelInput {
                sequence_id: 0,
                phase: ModelPhase::Decode,
                past_tokens: 0,
                token_ids: vec![1],
                positions: vec![0],
            },
            3,
        )
        .unwrap();
        assert_eq!(logits.to_vec1::<f32>().unwrap(), values);
    }
}