//! HuggingFace Hub compatibility layer.
//!
//! This module provides seamless integration with HuggingFace Hub, enabling:
//! - Loading models from HuggingFace Hub
//! - Converting HuggingFace configs to RustyTorch format
//! - Auto-downloading and caching models
//! - Support for private models with authentication

use crate::{HubError, HubResult, LocalCache, PretrainedConfig};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use tokio::fs;
use tracing::{debug, info, warn};

/// HuggingFace Hub base URL.
const HF_HUB_URL: &str = "https://huggingface.co";

/// HuggingFace model configuration (config.json).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HFModelConfig {
    /// Model architecture (e.g., "LlamaForCausalLM", "GPT2LMHeadModel")
    #[serde(default)]
    pub architectures: Vec<String>,
    /// Model type (e.g., "llama", "gpt2", "bert")
    #[serde(default)]
    pub model_type: Option<String>,
    /// Hidden size
    #[serde(default)]
    pub hidden_size: Option<usize>,
    /// Number of hidden layers
    #[serde(default)]
    pub num_hidden_layers: Option<usize>,
    /// Number of attention heads
    #[serde(default)]
    pub num_attention_heads: Option<usize>,
    /// Number of key-value heads (for GQA)
    #[serde(default)]
    pub num_key_value_heads: Option<usize>,
    /// Intermediate size (FFN)
    #[serde(default)]
    pub intermediate_size: Option<usize>,
    /// Vocabulary size
    #[serde(default)]
    pub vocab_size: Option<usize>,
    /// Maximum position embeddings
    #[serde(default)]
    pub max_position_embeddings: Option<usize>,
    /// RMS norm epsilon
    #[serde(default)]
    pub rms_norm_eps: Option<f64>,
    /// Layer norm epsilon
    #[serde(default)]
    pub layer_norm_eps: Option<f64>,
    /// Rope theta
    #[serde(default)]
    pub rope_theta: Option<f64>,
    /// Rope scaling configuration
    #[serde(default)]
    pub rope_scaling: Option<RopeScalingConfig>,
    /// Tie word embeddings
    #[serde(default)]
    pub tie_word_embeddings: Option<bool>,
    /// Hidden activation function
    #[serde(default)]
    pub hidden_act: Option<String>,
    /// Torch dtype (e.g., "float16", "bfloat16")
    #[serde(default)]
    pub torch_dtype: Option<String>,
    /// Use cache for generation
    #[serde(default)]
    pub use_cache: Option<bool>,
    /// Beginning of sequence token ID
    #[serde(default)]
    pub bos_token_id: Option<u32>,
    /// End of sequence token ID
    #[serde(default)]
    pub eos_token_id: Option<EosTokenId>,
    /// Padding token ID
    #[serde(default)]
    pub pad_token_id: Option<u32>,
    /// Additional configuration fields
    #[serde(flatten)]
    pub extra: HashMap<String, serde_json::Value>,
}

/// EOS token ID can be a single value or list.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum EosTokenId {
    /// Single EOS token
    Single(u32),
    /// Multiple EOS tokens
    Multiple(Vec<u32>),
}

impl EosTokenId {
    /// Get the primary EOS token ID.
    pub fn primary(&self) -> u32 {
        match self {
            EosTokenId::Single(id) => *id,
            EosTokenId::Multiple(ids) => ids.first().copied().unwrap_or(0),
        }
    }

    /// Get all EOS token IDs.
    pub fn all(&self) -> Vec<u32> {
        match self {
            EosTokenId::Single(id) => vec![*id],
            EosTokenId::Multiple(ids) => ids.clone(),
        }
    }
}

/// Rope scaling configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RopeScalingConfig {
    /// Scaling type (e.g., "linear", "dynamic")
    #[serde(rename = "type")]
    pub scaling_type: String,
    /// Scaling factor
    pub factor: f64,
    /// Original max position embeddings (for dynamic scaling)
    #[serde(default)]
    pub original_max_position_embeddings: Option<usize>,
}

/// HuggingFace tokenizer configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HFTokenizerConfig {
    /// Tokenizer class name
    #[serde(default)]
    pub tokenizer_class: Option<String>,
    /// Model max length
    #[serde(default)]
    pub model_max_length: Option<usize>,
    /// Padding side
    #[serde(default)]
    pub padding_side: Option<String>,
    /// Truncation side
    #[serde(default)]
    pub truncation_side: Option<String>,
    /// Chat template
    #[serde(default)]
    pub chat_template: Option<String>,
    /// Beginning of sequence token
    #[serde(default)]
    pub bos_token: Option<TokenValue>,
    /// End of sequence token
    #[serde(default)]
    pub eos_token: Option<TokenValue>,
    /// Padding token
    #[serde(default)]
    pub pad_token: Option<TokenValue>,
    /// Unknown token
    #[serde(default)]
    pub unk_token: Option<TokenValue>,
    /// Additional configuration fields
    #[serde(flatten)]
    pub extra: HashMap<String, serde_json::Value>,
}

/// Token value (can be string or object with content).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum TokenValue {
    /// Simple string token
    Simple(String),
    /// Token with metadata
    WithMetadata {
        content: String,
        #[serde(default)]
        lstrip: bool,
        #[serde(default)]
        rstrip: bool,
        #[serde(default)]
        single_word: bool,
        #[serde(default)]
        normalized: bool,
    },
}

impl TokenValue {
    /// Get the token content.
    pub fn content(&self) -> &str {
        match self {
            TokenValue::Simple(s) => s,
            TokenValue::WithMetadata { content, .. } => content,
        }
    }
}

/// HuggingFace generation configuration.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct HFGenerationConfig {
    /// Maximum new tokens to generate
    #[serde(default)]
    pub max_new_tokens: Option<usize>,
    /// Maximum total length
    #[serde(default)]
    pub max_length: Option<usize>,
    /// Temperature for sampling
    #[serde(default)]
    pub temperature: Option<f64>,
    /// Top-p (nucleus) sampling
    #[serde(default)]
    pub top_p: Option<f64>,
    /// Top-k sampling
    #[serde(default)]
    pub top_k: Option<usize>,
    /// Repetition penalty
    #[serde(default)]
    pub repetition_penalty: Option<f64>,
    /// Do sample
    #[serde(default)]
    pub do_sample: Option<bool>,
    /// Number of beams
    #[serde(default)]
    pub num_beams: Option<usize>,
    /// Early stopping
    #[serde(default)]
    pub early_stopping: Option<bool>,
    /// EOS token ID
    #[serde(default)]
    pub eos_token_id: Option<EosTokenId>,
    /// PAD token ID
    #[serde(default)]
    pub pad_token_id: Option<u32>,
    /// BOS token ID
    #[serde(default)]
    pub bos_token_id: Option<u32>,
    /// Additional configuration fields
    #[serde(flatten)]
    pub extra: HashMap<String, serde_json::Value>,
}

/// RustyTorch model configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RTXModelConfig {
    /// Model architecture type
    pub architecture: RTXArchitecture,
    /// Hidden dimension size
    pub hidden_size: usize,
    /// Number of transformer layers
    pub num_layers: usize,
    /// Number of attention heads
    pub num_heads: usize,
    /// Number of key-value heads (for GQA, defaults to num_heads)
    pub num_kv_heads: usize,
    /// FFN intermediate size
    pub intermediate_size: usize,
    /// Vocabulary size
    pub vocab_size: usize,
    /// Maximum sequence length
    pub max_seq_len: usize,
    /// RoPE theta parameter
    pub rope_theta: f64,
    /// Layer norm epsilon
    pub norm_eps: f64,
    /// Activation function
    pub activation: RTXActivation,
    /// Data type for inference
    pub dtype: RTXDType,
    /// Whether to tie input/output embeddings
    pub tie_embeddings: bool,
    /// RoPE scaling configuration
    pub rope_scaling: Option<RTXRopeScaling>,
    /// Beginning of sequence token ID
    pub bos_token_id: u32,
    /// End of sequence token IDs
    pub eos_token_ids: Vec<u32>,
    /// Padding token ID
    pub pad_token_id: Option<u32>,
}

/// RustyTorch model architecture types.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RTXArchitecture {
    /// LLaMA family (LLaMA, LLaMA 2, LLaMA 3)
    Llama,
    /// Mistral family
    Mistral,
    /// Mixtral (MoE)
    Mixtral,
    /// GPT-2
    Gpt2,
    /// GPT-NeoX / Pythia
    GptNeoX,
    /// Falcon
    Falcon,
    /// Phi family
    Phi,
    /// Qwen family
    Qwen,
    /// BERT
    Bert,
    /// RoBERTa
    Roberta,
    /// T5
    T5,
    /// Mamba (state space model)
    Mamba,
    /// Gemma
    Gemma,
    /// Other/unknown
    Other,
}

impl RTXArchitecture {
    /// Detect architecture from HuggingFace model type.
    pub fn from_hf_model_type(model_type: &str) -> Self {
        match model_type.to_lowercase().as_str() {
            "llama" => RTXArchitecture::Llama,
            "mistral" => RTXArchitecture::Mistral,
            "mixtral" => RTXArchitecture::Mixtral,
            "gpt2" => RTXArchitecture::Gpt2,
            "gpt_neox" | "gptneox" => RTXArchitecture::GptNeoX,
            "falcon" => RTXArchitecture::Falcon,
            "phi" | "phi3" => RTXArchitecture::Phi,
            "qwen" | "qwen2" => RTXArchitecture::Qwen,
            "bert" => RTXArchitecture::Bert,
            "roberta" => RTXArchitecture::Roberta,
            "t5" => RTXArchitecture::T5,
            "mamba" => RTXArchitecture::Mamba,
            "gemma" | "gemma2" => RTXArchitecture::Gemma,
            _ => RTXArchitecture::Other,
        }
    }
}

/// Activation function types.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RTXActivation {
    /// SiLU / Swish
    SiLU,
    /// GELU
    GELU,
    /// GELUApprox (tanh approximation)
    GELUApprox,
    /// ReLU
    ReLU,
    /// Mish
    Mish,
}

impl RTXActivation {
    /// Convert from HuggingFace activation name.
    pub fn from_hf(name: &str) -> Self {
        match name.to_lowercase().as_str() {
            "silu" | "swish" | "silu_and_mul" => RTXActivation::SiLU,
            "gelu" | "gelu_new" => RTXActivation::GELU,
            "gelu_pytorch_tanh" | "gelu_fast" => RTXActivation::GELUApprox,
            "relu" => RTXActivation::ReLU,
            "mish" => RTXActivation::Mish,
            _ => RTXActivation::SiLU, // Default
        }
    }
}

/// Data type for model weights.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RTXDType {
    /// 32-bit float
    F32,
    /// 16-bit float
    F16,
    /// Brain float 16
    BF16,
    /// 8-bit float (E4M3)
    F8E4M3,
    /// 8-bit float (E5M2)
    F8E5M2,
}

impl RTXDType {
    /// Convert from HuggingFace dtype string.
    pub fn from_hf(dtype: &str) -> Self {
        match dtype.to_lowercase().as_str() {
            "float32" | "f32" => RTXDType::F32,
            "float16" | "f16" => RTXDType::F16,
            "bfloat16" | "bf16" => RTXDType::BF16,
            "float8_e4m3fn" | "fp8_e4m3" => RTXDType::F8E4M3,
            "float8_e5m2" | "fp8_e5m2" => RTXDType::F8E5M2,
            _ => RTXDType::F16, // Default
        }
    }
}

/// RoPE scaling configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RTXRopeScaling {
    /// Scaling type
    pub scaling_type: RopeScalingType,
    /// Scaling factor
    pub factor: f64,
    /// Original max position (for dynamic)
    pub original_max_position: Option<usize>,
}

/// RoPE scaling types.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RopeScalingType {
    /// Linear interpolation
    Linear,
    /// Dynamic NTK-aware scaling
    Dynamic,
    /// YaRN scaling
    Yarn,
    /// LongRoPE
    LongRope,
}

impl RopeScalingType {
    /// Convert from HuggingFace scaling type.
    pub fn from_hf(s: &str) -> Self {
        match s.to_lowercase().as_str() {
            "linear" => RopeScalingType::Linear,
            "dynamic" => RopeScalingType::Dynamic,
            "yarn" => RopeScalingType::Yarn,
            "longrope" => RopeScalingType::LongRope,
            _ => RopeScalingType::Linear,
        }
    }
}

/// Convert HuggingFace config to RustyTorch config.
pub fn convert_hf_config(hf_config: &HFModelConfig) -> HubResult<RTXModelConfig> {
    let model_type = hf_config
        .model_type
        .as_deref()
        .or_else(|| {
            hf_config
                .architectures
                .first()
                .map(std::string::String::as_str)
        })
        .unwrap_or("unknown");

    let architecture = RTXArchitecture::from_hf_model_type(model_type);

    let hidden_size = hf_config
        .hidden_size
        .ok_or_else(|| HubError::InvalidPackage {
            reason: "Missing hidden_size in config".to_string(),
        })?;

    let num_layers = hf_config
        .num_hidden_layers
        .ok_or_else(|| HubError::InvalidPackage {
            reason: "Missing num_hidden_layers in config".to_string(),
        })?;

    let num_heads = hf_config
        .num_attention_heads
        .ok_or_else(|| HubError::InvalidPackage {
            reason: "Missing num_attention_heads in config".to_string(),
        })?;

    let num_kv_heads = hf_config.num_key_value_heads.unwrap_or(num_heads);

    let intermediate_size =
        hf_config
            .intermediate_size
            .ok_or_else(|| HubError::InvalidPackage {
                reason: "Missing intermediate_size in config".to_string(),
            })?;

    let vocab_size = hf_config
        .vocab_size
        .ok_or_else(|| HubError::InvalidPackage {
            reason: "Missing vocab_size in config".to_string(),
        })?;

    let max_seq_len = hf_config.max_position_embeddings.unwrap_or(4096);

    let rope_theta = hf_config.rope_theta.unwrap_or(10000.0);

    let norm_eps = hf_config
        .rms_norm_eps
        .or(hf_config.layer_norm_eps)
        .unwrap_or(1e-5);

    let activation = hf_config
        .hidden_act
        .as_deref()
        .map_or(RTXActivation::SiLU, RTXActivation::from_hf);

    let dtype = hf_config
        .torch_dtype
        .as_deref()
        .map_or(RTXDType::F16, RTXDType::from_hf);

    let tie_embeddings = hf_config.tie_word_embeddings.unwrap_or(false);

    let rope_scaling = hf_config.rope_scaling.as_ref().map(|rs| RTXRopeScaling {
        scaling_type: RopeScalingType::from_hf(&rs.scaling_type),
        factor: rs.factor,
        original_max_position: rs.original_max_position_embeddings,
    });

    let bos_token_id = hf_config.bos_token_id.unwrap_or(1);
    let eos_token_ids = hf_config
        .eos_token_id
        .as_ref()
        .map_or_else(|| vec![2], EosTokenId::all);
    let pad_token_id = hf_config.pad_token_id;

    Ok(RTXModelConfig {
        architecture,
        hidden_size,
        num_layers,
        num_heads,
        num_kv_heads,
        intermediate_size,
        vocab_size,
        max_seq_len,
        rope_theta,
        norm_eps,
        activation,
        dtype,
        tie_embeddings,
        rope_scaling,
        bos_token_id,
        eos_token_ids,
        pad_token_id,
    })
}

/// HuggingFace Hub client for downloading models.
pub struct HFHubClient {
    /// HTTP client
    client: reqwest::Client,
    /// Authentication token
    token: Option<String>,
    /// Local cache
    cache: LocalCache,
    /// Base URL (can be changed for mirrors)
    base_url: String,
}

impl HFHubClient {
    /// Create a new HuggingFace Hub client.
    pub fn new(cache_dir: PathBuf) -> Self {
        Self {
            client: reqwest::Client::new(),
            token: std::env::var("HF_TOKEN")
                .or_else(|_| std::env::var("HUGGING_FACE_HUB_TOKEN"))
                .ok(),
            cache: LocalCache::new(cache_dir),
            base_url: HF_HUB_URL.to_string(),
        }
    }

    /// Set authentication token.
    pub fn with_token(mut self, token: impl Into<String>) -> Self {
        self.token = Some(token.into());
        self
    }

    /// Set base URL (for mirrors).
    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
        self.base_url = url.into();
        self
    }

    /// Download a model from HuggingFace Hub.
    pub async fn download_model(
        &self,
        model_id: &str,
        config: &PretrainedConfig,
    ) -> HubResult<PathBuf> {
        let revision = config.revision.as_deref().unwrap_or("main");
        let cache_path = self.cache.cache_path(model_id, revision);

        // Check cache
        if !config.force_download && self.cache.is_cached(&cache_path).await {
            debug!("Using cached model: {}", cache_path.display());
            return Ok(cache_path);
        }

        if config.local_files_only {
            return Err(HubError::ModelNotFound {
                model_id: format!("{} (local files only)", model_id),
            });
        }

        info!("Downloading model from HuggingFace: {}", model_id);

        // Create cache directory
        fs::create_dir_all(&cache_path).await?;

        // Download essential files
        let files_to_download = [
            "config.json",
            "tokenizer.json",
            "tokenizer_config.json",
            "generation_config.json",
            "model.safetensors.index.json",
        ];

        for filename in &files_to_download {
            if let Err(e) = self
                .download_file(model_id, filename, revision, &cache_path)
                .await
            {
                debug!("Optional file {} not found: {}", filename, e);
            }
        }

        // Download weight files
        self.download_weights(model_id, revision, &cache_path)
            .await?;

        Ok(cache_path)
    }

    /// Download a single file.
    async fn download_file(
        &self,
        model_id: &str,
        filename: &str,
        revision: &str,
        output_dir: &Path,
    ) -> HubResult<PathBuf> {
        let url = format!(
            "{}/{}/resolve/{}/{}",
            self.base_url, model_id, revision, filename
        );

        let mut request = self.client.get(&url);

        if let Some(ref token) = self.token {
            request = request.header("Authorization", format!("Bearer {}", token));
        }

        let response = request.send().await.map_err(|e| HubError::DownloadFailed {
            url: url.clone(),
            reason: e.to_string(),
        })?;

        if !response.status().is_success() {
            return Err(HubError::DownloadFailed {
                url,
                reason: format!("HTTP {}", response.status()),
            });
        }

        let output_path = output_dir.join(filename);
        let content = response
            .bytes()
            .await
            .map_err(|e| HubError::DownloadFailed {
                url: url.clone(),
                reason: e.to_string(),
            })?;

        fs::write(&output_path, &content).await?;
        debug!("Downloaded: {}", filename);

        Ok(output_path)
    }

    /// Download model weights.
    async fn download_weights(
        &self,
        model_id: &str,
        revision: &str,
        output_dir: &Path,
    ) -> HubResult<()> {
        // Check for sharded model first
        let index_path = output_dir.join("model.safetensors.index.json");

        if index_path.exists() {
            // Sharded model - download all shards
            let index_content = fs::read_to_string(&index_path).await?;
            let index: crate::safetensors::SafeTensorsIndex = serde_json::from_str(&index_content)?;

            let shard_files: std::collections::HashSet<_> = index.weight_map.values().collect();

            for shard_file in shard_files {
                self.download_file(model_id, shard_file, revision, output_dir)
                    .await?;
            }
        } else {
            // Try single file model
            if let Err(e) = self
                .download_file(model_id, "model.safetensors", revision, output_dir)
                .await
            {
                warn!("SafeTensors not found, trying PyTorch format: {}", e);
                self.download_file(model_id, "pytorch_model.bin", revision, output_dir)
                    .await?;
            }
        }

        Ok(())
    }

    /// Load model configuration from cache.
    pub async fn load_config(&self, model_path: &Path) -> HubResult<HFModelConfig> {
        let config_path = model_path.join("config.json");
        let config_content = fs::read_to_string(&config_path).await?;
        let config: HFModelConfig = serde_json::from_str(&config_content)?;
        Ok(config)
    }

    /// Load tokenizer configuration from cache.
    pub async fn load_tokenizer_config(&self, model_path: &Path) -> HubResult<HFTokenizerConfig> {
        let config_path = model_path.join("tokenizer_config.json");
        let config_content = fs::read_to_string(&config_path).await?;
        let config: HFTokenizerConfig = serde_json::from_str(&config_content)?;
        Ok(config)
    }

    /// Load generation configuration from cache.
    pub async fn load_generation_config(&self, model_path: &Path) -> HubResult<HFGenerationConfig> {
        let config_path = model_path.join("generation_config.json");
        let config_content = fs::read_to_string(&config_path).await?;
        let config: HFGenerationConfig = serde_json::from_str(&config_content)?;
        Ok(config)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_architecture_detection() {
        assert_eq!(
            RTXArchitecture::from_hf_model_type("llama"),
            RTXArchitecture::Llama
        );
        assert_eq!(
            RTXArchitecture::from_hf_model_type("mistral"),
            RTXArchitecture::Mistral
        );
        assert_eq!(
            RTXArchitecture::from_hf_model_type("gpt2"),
            RTXArchitecture::Gpt2
        );
        assert_eq!(
            RTXArchitecture::from_hf_model_type("bert"),
            RTXArchitecture::Bert
        );
        assert_eq!(
            RTXArchitecture::from_hf_model_type("unknown"),
            RTXArchitecture::Other
        );
    }

    #[test]
    fn test_activation_conversion() {
        assert_eq!(RTXActivation::from_hf("silu"), RTXActivation::SiLU);
        assert_eq!(RTXActivation::from_hf("gelu"), RTXActivation::GELU);
        assert_eq!(RTXActivation::from_hf("relu"), RTXActivation::ReLU);
    }

    #[test]
    fn test_dtype_conversion() {
        assert_eq!(RTXDType::from_hf("float32"), RTXDType::F32);
        assert_eq!(RTXDType::from_hf("float16"), RTXDType::F16);
        assert_eq!(RTXDType::from_hf("bfloat16"), RTXDType::BF16);
    }

    #[test]
    fn test_hf_config_parsing() {
        let config_json = r#"{
            "architectures": ["LlamaForCausalLM"],
            "model_type": "llama",
            "hidden_size": 4096,
            "num_hidden_layers": 32,
            "num_attention_heads": 32,
            "num_key_value_heads": 8,
            "intermediate_size": 14336,
            "vocab_size": 128256,
            "max_position_embeddings": 8192,
            "rms_norm_eps": 1e-5,
            "rope_theta": 500000.0,
            "hidden_act": "silu",
            "torch_dtype": "bfloat16",
            "tie_word_embeddings": false,
            "bos_token_id": 128000,
            "eos_token_id": [128001, 128008, 128009]
        }"#;

        let config: HFModelConfig = serde_json::from_str(config_json).unwrap();

        assert_eq!(config.model_type, Some("llama".to_string()));
        assert_eq!(config.hidden_size, Some(4096));
        assert_eq!(config.num_hidden_layers, Some(32));
        assert_eq!(config.num_attention_heads, Some(32));
        assert_eq!(config.num_key_value_heads, Some(8));
        assert_eq!(config.vocab_size, Some(128256));
        assert_eq!(config.rope_theta, Some(500000.0));

        // Test EOS token ID
        match &config.eos_token_id {
            Some(EosTokenId::Multiple(ids)) => {
                assert_eq!(ids.len(), 3);
                assert_eq!(ids[0], 128001);
            }
            _ => panic!("Expected multiple EOS tokens"),
        }
    }

    #[test]
    fn test_convert_hf_config() {
        let hf_config = HFModelConfig {
            architectures: vec!["LlamaForCausalLM".to_string()],
            model_type: Some("llama".to_string()),
            hidden_size: Some(4096),
            num_hidden_layers: Some(32),
            num_attention_heads: Some(32),
            num_key_value_heads: Some(8),
            intermediate_size: Some(14336),
            vocab_size: Some(128256),
            max_position_embeddings: Some(8192),
            rms_norm_eps: Some(1e-5),
            layer_norm_eps: None,
            rope_theta: Some(500000.0),
            rope_scaling: None,
            tie_word_embeddings: Some(false),
            hidden_act: Some("silu".to_string()),
            torch_dtype: Some("bfloat16".to_string()),
            use_cache: Some(true),
            bos_token_id: Some(128000),
            eos_token_id: Some(EosTokenId::Single(128001)),
            pad_token_id: None,
            extra: HashMap::new(),
        };

        let rtx_config = convert_hf_config(&hf_config).unwrap();

        assert_eq!(rtx_config.architecture, RTXArchitecture::Llama);
        assert_eq!(rtx_config.hidden_size, 4096);
        assert_eq!(rtx_config.num_layers, 32);
        assert_eq!(rtx_config.num_heads, 32);
        assert_eq!(rtx_config.num_kv_heads, 8);
        assert_eq!(rtx_config.intermediate_size, 14336);
        assert_eq!(rtx_config.vocab_size, 128256);
        assert_eq!(rtx_config.max_seq_len, 8192);
        assert_eq!(rtx_config.rope_theta, 500000.0);
        assert_eq!(rtx_config.activation, RTXActivation::SiLU);
        assert_eq!(rtx_config.dtype, RTXDType::BF16);
    }

    #[test]
    fn test_tokenizer_config_parsing() {
        let config_json = r#"{
            "tokenizer_class": "PreTrainedTokenizerFast",
            "model_max_length": 8192,
            "padding_side": "left",
            "chat_template": "{% for message in messages %}..."
        }"#;

        let config: HFTokenizerConfig = serde_json::from_str(config_json).unwrap();

        assert_eq!(
            config.tokenizer_class,
            Some("PreTrainedTokenizerFast".to_string())
        );
        assert_eq!(config.model_max_length, Some(8192));
        assert_eq!(config.padding_side, Some("left".to_string()));
    }

    #[test]
    fn test_generation_config_parsing() {
        let config_json = r#"{
            "max_new_tokens": 512,
            "temperature": 0.7,
            "top_p": 0.9,
            "do_sample": true
        }"#;

        let config: HFGenerationConfig = serde_json::from_str(config_json).unwrap();

        assert_eq!(config.max_new_tokens, Some(512));
        assert_eq!(config.temperature, Some(0.7));
        assert_eq!(config.top_p, Some(0.9));
        assert_eq!(config.do_sample, Some(true));
    }

    #[test]
    fn test_eos_token_id_variants() {
        // Single
        let single: EosTokenId = serde_json::from_str("128001").unwrap();
        assert_eq!(single.primary(), 128001);
        assert_eq!(single.all(), vec![128001]);

        // Multiple
        let multiple: EosTokenId = serde_json::from_str("[128001, 128008]").unwrap();
        assert_eq!(multiple.primary(), 128001);
        assert_eq!(multiple.all(), vec![128001, 128008]);
    }

    #[test]
    fn test_token_value_variants() {
        // Simple
        let simple: TokenValue = serde_json::from_str(r#""<|begin_of_text|>""#).unwrap();
        assert_eq!(simple.content(), "<|begin_of_text|>");

        // With metadata
        let with_meta: TokenValue = serde_json::from_str(
            r#"{"content": "<|end_of_text|>", "lstrip": false, "rstrip": false}"#,
        )
        .unwrap();
        assert_eq!(with_meta.content(), "<|end_of_text|>");
    }

    #[test]
    fn test_rope_scaling() {
        let config_json = r#"{
            "model_type": "llama",
            "hidden_size": 4096,
            "num_hidden_layers": 32,
            "num_attention_heads": 32,
            "intermediate_size": 14336,
            "vocab_size": 128256,
            "rope_scaling": {
                "type": "dynamic",
                "factor": 2.0,
                "original_max_position_embeddings": 8192
            }
        }"#;

        let config: HFModelConfig = serde_json::from_str(config_json).unwrap();
        let rtx_config = convert_hf_config(&config).unwrap();

        assert!(rtx_config.rope_scaling.is_some());
        let scaling = rtx_config.rope_scaling.unwrap();
        assert_eq!(scaling.scaling_type, RopeScalingType::Dynamic);
        assert_eq!(scaling.factor, 2.0);
        assert_eq!(scaling.original_max_position, Some(8192));
    }
}