#![recursion_limit = "256"]
#![feature(iter_array_chunks)]

use burn::{prelude::*, tensor::activation::quiet_softmax};
use tokenizers::Tokenizer;

mod data {
    #![allow(clippy::all)] // have massive codegen here.
    extern crate alloc;

    include!(concat!(env!("OUT_DIR"), "/burned/ai_text/model.rs"));

    pub static TOKENIZER: &[u8] =
        include_bytes!(concat!(env!("OUT_DIR"), "/burned/ai_text/tokenizer_d.json"));
}

#[derive(Debug, thiserror::Error)]
pub enum ModelErr {
    #[error("Failed to load TOKENIZER")]
    Tokenizer(Box<dyn core::error::Error + Send + Sync>),
    #[error("Invalid TOKENIZER_CONFIG")]
    TokenizerConfig(serde_json::Error),
    #[error("Invalid CONFIG")]
    Config(serde_json::Error),
    #[error("Missing required token `{0}` (this is name, not the token itself)")]
    MissingToken(&'static str),
    #[error("Failed to apply truncation to the tokenizer")]
    TokenizerWithTruncation(Box<dyn core::error::Error + Send + Sync>),
}

#[derive(Debug, thiserror::Error)]
pub enum ModelInferenceErr {
    #[error("Failed to encode inputs")]
    InputEncode(Box<dyn core::error::Error + Send + Sync>),
}

#[derive(Debug, Clone)]
pub struct ModelInference {
    pub prob_ai: f32,
    pub prob_human: f32,
}

#[derive(Debug, Clone)]
pub struct ModelStats {
    pub tokenizer_max_length: Option<usize>,
}

pub struct Model<B: Backend> {
    model: data::Model<B>,
    tokenizer: Tokenizer,
    device: &'static <B as Backend>::Device,
}

impl<B: Backend> Model<B> {
    pub fn try_load(device: &'static <B as Backend>::Device) -> Result<Self, ModelErr> {
        tracing::debug!(device = ?device, "Requested to load model and tokenizer");

        let tokenizer = Tokenizer::from_bytes(data::TOKENIZER).map_err(ModelErr::Tokenizer)?;
        let model = data::Model::from_embedded(device);

        Ok(Self {
            device,
            model,
            tokenizer,
        })
    }

    pub fn stats(&self) -> ModelStats {
        let tokenizer_max_length = self.tokenizer.get_truncation().map(|t| t.max_length);

        ModelStats {
            tokenizer_max_length,
        }
    }

    #[tracing::instrument(skip_all)]
    pub fn batch_inference(
        &self,
        texts: Vec<&str>,
    ) -> Result<Vec<ModelInference>, ModelInferenceErr> {
        tracing::debug!(len = texts.len(), "Requested batch inference");

        if texts.is_empty() {
            return Ok(Vec::new());
        }

        let encodings = self
            .tokenizer
            .encode_batch(texts, true)
            .map_err(ModelInferenceErr::InputEncode)?;

        let batch_len = encodings.len();
        let encoding_len = encodings.first().map(|f| f.len()).unwrap_or_default();

        let (ids, attn, types) = {
            let bulk_size = batch_len * encoding_len;

            let mut buf = vec![0u32; bulk_size * 3]; // 3 bulks: ids, attn, types

            let ids_i = 0;
            let attn_i = ids_i + bulk_size;
            let types_i = attn_i + bulk_size;

            for (i, e) in encodings.iter().enumerate() {
                let off = i * encoding_len;
                buf[ids_i + off..ids_i + off + encoding_len].copy_from_slice(e.get_ids());
                buf[attn_i + off..attn_i + off + encoding_len]
                    .copy_from_slice(e.get_attention_mask());
                buf[types_i + off..types_i + off + encoding_len].copy_from_slice(e.get_type_ids());
            }

            let shape = Shape::new([batch_len, encoding_len]);

            let ids = &buf[ids_i..ids_i + bulk_size];
            let attn = &buf[attn_i..attn_i + bulk_size];
            let types = &buf[types_i..types_i + bulk_size];

            let ids: Tensor<B, 2, Int> =
                Tensor::<B, 1, Int>::from_ints(ids, self.device).reshape(shape.clone());

            let attn: Tensor<B, 2, Int> =
                Tensor::<B, 1, Int>::from_ints(attn, self.device).reshape(shape.clone());

            let types: Tensor<B, 2, Int> =
                Tensor::<B, 1, Int>::from_ints(types, self.device).reshape(shape);

            (ids, attn, types)
        };

        tracing::debug!("Batch inference ready, encoded");
        let logits = cfg_select! {
            all(feature = "_has_attention", feature = "_has_types") => self.model.forward(ids, attn, types),
            feature = "_has_types" => self.model.forward(ids, types),
            feature = "_has_attention" => self.model.forward(ids, attn),
            _ => self.model.forward(ids)
        };

        tracing::debug!(logits = ?logits, "Batch inference complete");
        let probs = quiet_softmax(logits, 1).into_data();

        Ok(Vec::from_iter(probs.iter::<f32>().array_chunks::<2>().map(
            |[prob_human, prob_ai]| ModelInference {
                prob_ai,
                prob_human,
            },
        )))
    }
}