use burn_onnx::ModelGen;
use serde::{Deserialize, Deserializer};
use std::{
    env,
    path::{Path, PathBuf},
};
use tokenizers::Tokenizer;

#[derive(Deserialize)]
pub struct Config {
    hidden_size: usize,
    num_attention_heads: usize,
}

#[derive(Deserialize)]
pub struct TokenizerConfig {
    pad_token: String,
    #[serde(
        alias = "model_max_length",
        default = "TokenizerConfig::max_length_default",
        deserialize_with = "TokenizerConfig::max_length_de"
    )]
    max_length: usize,
}

impl TokenizerConfig {
    const fn max_length_default() -> usize {
        512
    }

    /// If value is far above `LENGTH_CAP` - that means broken config. We're going to
    /// play safe and replace it with `max_length_default`.
    ///
    /// Yeah yeah, I know, "(in)sane defaults".
    ///
    /// Yet our current model BTW (!) has insanely large `model_max_length` while
    /// in reality its cap is 512.
    fn max_length_de<'de, D>(deserializer: D) -> Result<usize, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct Visitor;
        const LENGTH_CAP: usize = 8192;

        impl<'de> serde::de::Visitor<'de> for Visitor {
            type Value = usize;

            fn expecting(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
                write!(
                    fmt,
                    "unsigned integer. Values are capped by {LENGTH_CAP}. Float is accepted, and also capped to {LENGTH_CAP}."
                )
            }

            fn visit_u64<E>(self, v: u64) -> Result<Self::Value, E>
            where
                E: serde::de::Error,
            {
                // This is configuration sanity check.
                if v > (LENGTH_CAP * 3) as u64 {
                    eprintln!(
                        "WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
                        TokenizerConfig::max_length_default()
                    );
                    Ok(TokenizerConfig::max_length_default())
                } else {
                    Ok(v.clamp(0, LENGTH_CAP as u64) as usize)
                }
            }

            fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E>
            where
                E: serde::de::Error,
            {
                // This is configuration sanity check.
                if v > (LENGTH_CAP * 3) as i64 {
                    eprintln!(
                        "WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
                        TokenizerConfig::max_length_default()
                    );
                    Ok(TokenizerConfig::max_length_default())
                } else {
                    Ok(v.clamp(0, LENGTH_CAP as i64) as usize)
                }
            }

            fn visit_f64<E>(self, v: f64) -> Result<Self::Value, E>
            where
                E: serde::de::Error,
            {
                // This is configuration sanity check.
                if v > (LENGTH_CAP * 3) as f64 {
                    eprintln!(
                        "WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
                        TokenizerConfig::max_length_default()
                    );
                    Ok(TokenizerConfig::max_length_default())
                } else {
                    Ok(v.clamp(0.0, LENGTH_CAP as f64) as usize)
                }
            }
        }

        deserializer.deserialize_any(Visitor)
    }
}

fn main() {
    println!("cargo:rerun-if-changed=build.rs");
    println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT");
    println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT_ONNX");
    println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT_ATTENTION");
    println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT_TYPES");

    let model_repo = std::env::var("_MK_MODEL_AI_TEXT")
        .inspect_err(|_| {
            eprintln!("_MK_MODEL_AI_TEXT env variable required to point to HF model repository")
        })
        .expect("_MK_MODEL_AI_TEXT is set to huggingface model repo");
    let model_file =
        std::env::var("_MK_MODEL_AI_TEXT_ONNX").inspect_err(|_| {
            eprintln!("_MK_MODEL_AI_TEXT_ONNX env variable required to point to onnx file under HF model repository");
        }).expect("_MK_MODEL_AI_TEXT_FILE is set to .onnx path at huggingface model repo");

    let model_name = model_file
        .rsplit_once('/')
        .map(|(_, m)| m)
        .unwrap_or(model_file.as_str())
        .rsplit_once('.')
        .map(|(name, _)| name)
        .unwrap_or(model_file.as_str());

    let has_attention = std::env::var("_MK_MODEL_AI_TEXT_ATTENTION")
        .is_ok_and(|v| v == "true" || v == "yes" || v == "1");
    let has_types = std::env::var("_MK_MODEL_AI_TEXT_TYPES")
        .is_ok_and(|v| v == "true" || v == "yes" || v == "1");

    if has_attention {
        println!("cargo:rustc-cfg=feature=\"_has_attention\"");
    }
    if has_types {
        println!("cargo:rustc-cfg=feature=\"_has_types\"");
    }

    let api = hf_hub::api::sync::Api::new().expect("Huggingface API");

    // dwnl model
    let repo = api.model(model_repo);
    let model = repo.get(&model_file).expect("ai_text: model.onnx fetch");
    let config_json = repo.get("config.json").expect("ai_text: config.json fetch");
    let tokenizer_json = repo
        .get("tokenizer.json")
        .expect("ai_text: tokenizer.json fetch");
    let tokenizer_config_json = repo
        .get("tokenizer_config.json")
        .expect("ai_text: tokenizer_config.json fetch");

    let burned_dir = Path::new("burned").join("ai_text");
    let out_dir = Path::new(&env::var("OUT_DIR").expect("OUT_DIR is env var")).join(&burned_dir);

    // handle model embedding
    ModelGen::new()
        .load_strategy(burn_onnx::LoadStrategy::Embedded)
        .simplify(true)
        .input(model.to_str().unwrap())
        .out_dir(burned_dir.to_str().unwrap())
        .run_from_script();

    // They put model into <filename>.rs, which depends on env variables set, to we move
    // model into well-known location.
    let loc = PathBuf::from(std::env::var("OUT_DIR").expect("OUT_DIR"))
        .join(&burned_dir)
        .join(model_name)
        .with_extension("rs");
    let to = PathBuf::from(std::env::var("OUT_DIR").expect("OUT_DIR"))
        .join(&burned_dir)
        .join("model.rs");

    std::fs::copy(&loc, &to)
        .and_then(|_| std::fs::remove_file(&loc))
        .inspect_err(|e| eprintln!("While moving {loc:?} to {to:?}: {e:?}"))
        .expect("To move <model_name>.rs to model.rs");

    // construct tokenizer AoT
    let tokenizer = {
        let mut tokenizer = Tokenizer::from_file(tokenizer_json)
            .inspect_err(|e| eprintln!("While loading tokenizer.json: {e:?}"))
            .expect("Tokenizer JSON at tokenizer.json");

        let tokenizer_config: TokenizerConfig = serde_json::from_reader(
            std::fs::File::open(tokenizer_config_json)
                .inspect_err(|e| eprintln!("While opening tokenizer_config.json: {e:?}"))
                .expect("File tokenizer_config.json"),
        )
        .inspect_err(|e| eprint!("While deserializing tokenizer_config.json: {e:?}"))
        .expect("Valid JSON at tokenizer_config.json");

        let pad_id = tokenizer
            .token_to_id(&tokenizer_config.pad_token)
            .expect("Tokenizer to know PAD token");

        let config: Config = serde_json::from_reader(
            std::fs::File::open(config_json)
                .inspect_err(|e| eprintln!("While opening config.json: {e:?}"))
                .expect("Open file config.json"),
        )
        .inspect_err(|e| eprintln!("While deserializing config.json: {e:?}"))
        .expect("Valid JSON at config.json");

        tokenizer
            .with_truncation(Some(tokenizers::TruncationParams {
                direction: tokenizers::TruncationDirection::Right,
                max_length: tokenizer_config.max_length,
                strategy: Default::default(),
                stride: Default::default(),
            }))
            .inspect_err(|e| eprintln!("While setting tokenizer truncation: {e:?}"))
            .expect("Set tokenizer truncation");

        tokenizer.with_padding(Some(tokenizers::PaddingParams {
            strategy: tokenizers::PaddingStrategy::BatchLongest,
            direction: tokenizers::PaddingDirection::Right,
            pad_to_multiple_of: Some(config.hidden_size / config.num_attention_heads),
            pad_id,
            pad_type_id: 0,
            pad_token: tokenizer_config.pad_token,
        }));

        tokenizer
    };

    tokenizer
        .save(out_dir.join("tokenizer_d.json"), false)
        .inspect_err(|e| eprintln!("While saving tokenizer: {e:?}"))
        .expect("Save tokenizer");
}