//! Candle host pipeline for CALCULET-compiled language models.
//!
//! This replaces the tokenizer, Qwen chat-template, sampling, generation loop,
//! and CALRT tensor glue from the captured llama.cpp fork. The compiled
//! prefill/decode graphs remain Calbin models executed by `calculet-calrt`.

mod backend;
mod calrt;
mod engine;
mod error;
mod template;
mod tokenizer;

pub use backend::CandleChatBackend;
pub use calrt::{CalrtCompiledModel, CalrtModelPlan};
pub use engine::{
    CandleEngine, CompiledModel, DEFAULT_SEQUENCE_ID, GenerationConfig, ModelInput, ModelPhase,
    ModelStatus,
};
pub use error::{Error, Result};
pub use template::{
    ParsedAssistantOutput, QWEN_END_OF_TEXT, QWEN_IM_END, parse_qwen3_output, render_qwen3,
};
pub use tokenizer::{HuggingFaceTokenizer, TokenCodec};

#[cfg(test)]
mod tests {
    use super::*;
    use agentos_protocol::{
        CHAT_PROTOCOL_VERSION, ChatMessage, ChatRequest, MessageRole, StopReason, ToolChoice,
    };
    use candle_core::{Device, Tensor};
    use serde_json::Map;
    use std::collections::BTreeMap;
    use std::sync::Arc;

    #[derive(Debug)]
    struct TestTokenizer {
        pieces: BTreeMap<u32, &'static str>,
    }

    impl TestTokenizer {
        fn new() -> Self {
            Self {
                pieces: BTreeMap::from([(0, "prompt"), (1, "done"), (2, QWEN_IM_END)]),
            }
        }
    }

    impl TokenCodec for TestTokenizer {
        fn encode(&self, _text: &str) -> Result<Vec<u32>> {
            Ok(vec![0])
        }

        fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result<String> {
            Ok(token_ids
                .iter()
                .filter_map(|token| self.pieces.get(token))
                .copied()
                .collect())
        }

        fn token_to_id(&self, token: &str) -> Option<u32> {
            self.pieces
                .iter()
                .find_map(|(id, piece)| (*piece == token).then_some(*id))
        }

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

    #[derive(Debug, Default)]
    struct TestModel {
        step: usize,
    }

    impl CompiledModel for TestModel {
        fn status(&self) -> ModelStatus {
            ModelStatus {
                ready: true,
                reason: "test model is ready".into(),
                details: serde_json::Value::Null,
            }
        }

        fn vocabulary_size(&self) -> usize {
            4
        }

        fn maximum_sequence_length(&self) -> usize {
            32
        }

        fn reset_sequence(&mut self, _sequence_id: u32) -> Result<()> {
            self.step = 0;
            Ok(())
        }

        fn forward(&mut self, _input: &ModelInput) -> Result<Tensor> {
            let token = if self.step == 0 { 1 } else { 2 };
            self.step += 1;
            let mut logits = vec![-100.0_f32; self.vocabulary_size()];
            logits[token] = 100.0;
            logits[3] = 200.0;
            Ok(Tensor::from_vec(
                logits,
                self.vocabulary_size(),
                &Device::Cpu,
            )?)
        }
    }

    #[test]
    fn runs_a_deterministic_candle_generation_loop_without_hardware() {
        let mut engine = CandleEngine::new(
            Arc::new(TestTokenizer::new()),
            TestModel::default(),
            GenerationConfig::default(),
        )
        .unwrap();
        let response = engine
            .generate(&ChatRequest {
                protocol_version: CHAT_PROTOCOL_VERSION.into(),
                messages: vec![ChatMessage::text(MessageRole::User, "run")],
                tools: Vec::new(),
                max_output_tokens: 4,
                tool_choice: ToolChoice::Auto,
                parallel_tool_calls: false,
                metadata: Map::new(),
            })
            .unwrap();
        assert_eq!(response.content, "done");
        assert_eq!(response.stop_reason, StopReason::EndTurn);
        assert_eq!(response.usage.output_tokens, 2);
    }
}