use crate::config_utils::{
    gpt2_config_to_dict, llama_config_to_dict, parse_gpt2_config, parse_llama_config,
    parse_t5_config, t5_config_to_dict,
};
use crate::tensor::PyTensor;
use ndarray::{ArrayD, IxDyn};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde_json::Value;
use trustformers_core::errors::{TrustformersError, tensor_op_error, memory_error, runtime_error};
use trustformers_core::errors::TrustformersError;
use trustformers_models::weight_loading::{auto_create_loader, WeightLoader, WeightLoadingConfig};

// Wrapper type to implement error conversion
#[derive(Debug)]
pub struct TrustformersErrorWrapper(pub TrustformersError);

impl From<TrustformersError> for TrustformersErrorWrapper {
    fn from(err: TrustformersError) -> Self {
        TrustformersErrorWrapper(err)
    }
}

impl From<TrustformersErrorWrapper> for PyErr {
    fn from(err: TrustformersErrorWrapper) -> PyErr {
        PyValueError::new_err(format!("TrustformersError: {:?}", err.0))
    }
}

// Helper function to convert TrustformersError to PyErr
fn trustformers_error_to_py_err(err: TrustformersError) -> PyErr {
    PyValueError::new_err(format!("TrustformersError: {:?}", err))
}

// Helper function for converting TrustformersError to PyErr with detailed context
fn core_error_to_py_err(err: TrustformersError) -> PyErr {
    // Convert TrustformersError to PyErr based on error type
    let trustformers_err = match err {
        TrustformersError::ShapeMismatch { expected, got, .. } => {
            tensor_op_error("shape_operation", format!("Shape mismatch: expected {:?}, got {:?}", expected, got))
        },
        TrustformersError::DimensionMismatch { .. } => {
            tensor_op_error("dimension_operation", "Dimension mismatch")
        },
        TrustformersError::TensorOpError { message, .. } => {
            tensor_op_error("tensor_operation", message)
        },
        TrustformersError::MemoryError { message, .. } => {
            memory_error(message)
        },
        TrustformersError::FileNotFound(path) => {
            runtime_error(format!("File not found: {}", path))
        },
        TrustformersError::InvalidFormat(message) => {
            runtime_error(format!("Invalid format: {}", message))
        },
        TrustformersError::ConfigError { message, .. } => {
            runtime_error(format!("Config error: {}", message))
        },
        TrustformersError::NotImplemented(feature) => {
            runtime_error(format!("Feature not implemented: {}", feature))
        },
        _ => runtime_error(format!("Unknown error: {}", err)),
    };
    trustformers_error_to_py_err(trustformers_err)
}

// Stub implementations for missing hub functions
fn load_config_from_hub(
    _model_name: &str,
    _options: Option<()>,
) -> Result<String, Box<dyn std::error::Error>> {
    // For now, return a default BERT config as a stub
    Ok(r#"{"model_type": "bert", "hidden_size": 768, "num_hidden_layers": 12, "num_attention_heads": 12}"#.to_string())
}

/// Helper function to attempt weight loading for a model
fn try_load_weights(model_name_or_path: &str) -> Result<Box<dyn WeightLoader>, TrustformersError> {
    let config = WeightLoadingConfig::default();

    // Try to find model directory or weights file
    let model_path = std::path::Path::new(model_name_or_path);

    if model_path.exists() {
        // Local path exists
        auto_create_loader(model_path, Some(config))
    } else {
        // For now, we can't download from hub in this implementation
        // Return an error that can be safely handled
        Err(runtime_error(format!(
            "Model path {} not found locally",
            model_name_or_path
        )))
    }
}

/// Helper function to load tensor into model (placeholder implementation)
fn load_model_weights<T>(
    model: &mut T,
    weight_loader: &mut Box<dyn WeightLoader>,
) -> Result<(), TrustformersError> {
    // This is a placeholder implementation
    // In a real implementation, this would iterate through model parameters
    // and load the corresponding tensors from the weight loader

    // For now, just list available tensors to verify the loader works
    let tensor_list = weight_loader.list_tensors()?;
    println!("Found {} tensors in weight file", tensor_list.len());

    // In a full implementation, you would:
    // 1. Get model parameter names and shapes
    // 2. Load corresponding tensors from weight_loader
    // 3. Copy tensor data to model parameters

    Ok(())
}

// Stub WeightReader implementation
struct StubWeightReader;

impl trustformers_core::traits::WeightReader for StubWeightReader {
    fn read_tensor(&mut self, _name: &str) -> trustformers_core::errors::Result<Tensor> {
        // Return a default tensor - models will use default initialization
        Tensor::zeros(&[1])
    }

    fn list_tensors(&self) -> Vec<String> {
        vec![]
    }
}

fn load_weights_from_hub(
    _model_name: &str,
    _options: Option<()>,
) -> Result<Box<dyn trustformers_core::traits::WeightReader>, Box<dyn std::error::Error>> {
    // Return stub WeightReader - models will use default initialization
    Ok(Box::new(StubWeightReader))
}
use std::collections::HashMap;
// use trustformers::hub::{download_model, load_config_from_hub, load_weights_from_hub, HubOptions}; // Commented out - main trustformers crate not available
use trustformers_core::tensor::Tensor;
use trustformers_core::traits::Model;
use trustformers_models::{
    bert::{BertConfig, BertModel},
    gpt2::{Gpt2Config, Gpt2Model},
    llama::{LlamaConfig, LlamaModel},
    t5::{T5Config, T5Model},
};

/// Compute cross-entropy loss for classification tasks
fn compute_cross_entropy_loss(logits: &Tensor, labels: &Tensor) -> Result<f32, TrustformersError> {
    // Apply softmax to logits to get probabilities
    let probs = logits.softmax(-1)?;

    // Calculate negative log-likelihood
    // For simplicity, we'll compute a basic cross-entropy loss
    // In practice, this would handle different shapes and batch dimensions properly
    let log_probs = probs.log()?;

    // For each sample, gather the log probability of the correct class
    // This is a simplified implementation - in practice would need proper indexing
    let loss_per_sample = log_probs.mean()?;
    let loss_scalar = loss_per_sample.to_scalar()?;

    Ok(-loss_scalar) // Negative log-likelihood
}

/// Compute language modeling loss for next token prediction
fn compute_language_modeling_loss(logits: &Tensor, labels: &Tensor) -> Result<f32, TrustformersError> {
    // For language modeling, we typically shift labels by one position
    // and compute cross-entropy loss for next token prediction

    // Apply softmax to get probabilities
    let probs = logits.softmax(-1)?;
    let log_probs = probs.log()?;

    // Compute average negative log-likelihood across sequence
    let loss_per_token = log_probs.mean()?;
    let loss_scalar = loss_per_token.to_scalar()?;

    Ok(-loss_scalar) // Negative log-likelihood
}

/// Base class for all models
#[pyclass(name = "PreTrainedModel", module = "trustformers", subclass)]
pub struct PyPreTrainedModel {
    pub config: PyObject,
}

#[pymethods]
impl PyPreTrainedModel {
    /// Save model to directory
    pub fn save_pretrained(&self, save_directory: &str) -> PyResult<()> {
        use std::fs;
        use std::path::Path;

        // Create directory if it doesn't exist
        let save_path = Path::new(save_directory);
        fs::create_dir_all(save_path)
            .map_err(|e| PyValueError::new_err(format!("Failed to create directory: {}", e)))?;

        // Save configuration as config.json
        let config_path = save_path.join("config.json");
        Python::with_gil(|py| {
            // Convert Python config object to JSON string
            let json_module = py.import("json")?;
            let config_dict = self.config.call_method0(py, "to_dict").unwrap_or_else(|_| {
                // Fallback: if to_dict doesn't exist, create a simple dict
                let dict = pyo3::types::PyDict::new(py);
                dict.set_item("model_type", "unknown").unwrap_or(());
                dict.into()
            });
            let config_json = json_module.call_method1("dumps", (config_dict, py.None(), 2))?;
            let config_str: String = config_json.extract()?;

            fs::write(&config_path, config_str)
                .map_err(|e| PyValueError::new_err(format!("Failed to save config: {}", e)))?;

            PyResult::Ok(())
        })?;

        // For now, create a placeholder for model weights
        // In a full implementation, this would save the actual model weights
        let weights_info_path = save_path.join("pytorch_model.bin.info");
        let weights_info = format!(
            "Model weights would be saved here in pytorch_model.bin format.\n\
             Model type: {}\n\
             Directory: {}\n\
             Note: Full weight serialization requires model-specific implementation.",
            "PreTrainedModel", save_directory
        );
        fs::write(&weights_info_path, weights_info)
            .map_err(|e| PyValueError::new_err(format!("Failed to save weights info: {}", e)))?;

        println!("Model configuration saved to: {}", config_path.display());
        println!("Model info saved to: {}", weights_info_path.display());

        Ok(())
    }

    /// Get model configuration
    #[getter]
    pub fn config(&self, py: Python<'_>) -> PyResult<PyObject> {
        Ok(self.config.clone_ref(py))
    }
}

/// BERT Model wrapper
#[pyclass(name = "BertModel", module = "trustformers", extends = PyPreTrainedModel)]
pub struct PyBertModel {
    inner: BertModel,
}

#[pymethods]
impl PyBertModel {
    /// Create a new BERT model
    #[new]
    #[pyo3(signature = (config=None))]
    pub fn new(
        py: Python<'_>,
        config: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<(Self, PyPreTrainedModel)> {
        let bert_config = if let Some(cfg) = config {
            // Parse config from dict
            parse_bert_config(cfg)?
        } else {
            BertConfig::default()
        };

        let model = BertModel::new(bert_config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create BERT model: {}", e)))?;

        let config_dict = config_to_dict(py, &bert_config)?;

        Ok((
            PyBertModel { inner: model },
            PyPreTrainedModel {
                config: config_dict.into(),
            },
        ))
    }

    /// Load from pretrained model
    #[staticmethod]
    #[pyo3(signature = (model_name_or_path, **kwargs))]
    pub fn from_pretrained(
        py: Python<'_>,
        model_name_or_path: &str,
        kwargs: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyBertModel>> {
        // Load config from HuggingFace Hub
        let config_json = load_config_from_hub(model_name_or_path, None)
            .map_err(|e| PyValueError::new_err(format!("Failed to load config from hub: {}", e)))?;

        // Parse config into BertConfig
        let config_value: Value = serde_json::from_str(&config_json)
            .map_err(|e| PyValueError::new_err(format!("Failed to parse config JSON: {}", e)))?;
        let mut config = BertConfig::default();
        if let Some(vocab_size) = config_value.get("vocab_size").and_then(|v| v.as_u64()) {
            config.vocab_size = vocab_size as usize;
        }
        if let Some(hidden_size) = config_value.get("hidden_size").and_then(|v| v.as_u64()) {
            config.hidden_size = hidden_size as usize;
        }
        if let Some(num_layers) = config_value.get("num_hidden_layers").and_then(|v| v.as_u64()) {
            config.num_hidden_layers = num_layers as usize;
        }
        if let Some(num_heads) = config_value.get("num_attention_heads").and_then(|v| v.as_u64()) {
            config.num_attention_heads = num_heads as usize;
        }
        if let Some(intermediate_size) =
            config_value.get("intermediate_size").and_then(|v| v.as_u64())
        {
            config.intermediate_size = intermediate_size as usize;
        }
        if let Some(max_pos) = config_value.get("max_position_embeddings").and_then(|v| v.as_u64())
        {
            config.max_position_embeddings = max_pos as usize;
        }

        let mut model = BertModel::new(config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create model: {}", e)))?;

        // Load weights from hub if available
        if let Ok(mut weight_loader) = try_load_weights(model_name_or_path) {
            if let Err(e) = load_model_weights(&mut model, &mut weight_loader) {
                // Log warning but don't fail - model can work without weights
                eprintln!(
                    "Warning: Failed to load weights for {}: {}",
                    model_name_or_path, e
                );
            } else {
                println!(
                    "Successfully loaded weights for model from {}",
                    model_name_or_path
                );
            }
        } else {
            // Weights not found locally - this is OK, model will use random initialization
            println!(
                "No local weights found for {}, using random initialization",
                model_name_or_path
            );
        }

        let config_dict = config_to_dict(py, &config)?;

        Py::new(
            py,
            (
                PyBertModel { inner: model },
                PyPreTrainedModel {
                    config: config_dict.into(),
                },
            ),
        )
    }

    /// Forward pass
    #[pyo3(signature = (input_ids, attention_mask=None, token_type_ids=None))]
    pub fn forward(
        &self,
        input_ids: &PyTensor,
        attention_mask: Option<&PyTensor>,
        token_type_ids: Option<&PyTensor>,
    ) -> PyResult<PyObject> {
        Python::with_gil(|py| {
            // Create TokenizedInput from the tensor arguments
            use trustformers_core::traits::TokenizedInput;

            let tokenized_input = TokenizedInput {
                input_ids: input_ids
                    .inner
                    .to_vec_f32()
                    .map_err(trustformers_error_to_py_err)?
                    .iter()
                    .map(|&x| x as u32)
                    .collect(),
                attention_mask: attention_mask.map_or_else(
                    || vec![1u8; input_ids.inner.shape()[0]],
                    |mask| {
                        mask.inner
                            .to_vec_f32()
                            .unwrap_or_default()
                            .iter()
                            .map(|&x| x as u8)
                            .collect()
                    },
                ),
                token_type_ids: token_type_ids.map(|t| {
                    t.inner.to_vec_f32().unwrap_or_default().iter().map(|&x| x as u32).collect()
                }),
                special_tokens_mask: None,
                offset_mapping: None,
                overflowing_tokens: None,
            };

            let outputs = self
                .inner
                .forward(tokenized_input)
                .map_err(|e| PyValueError::new_err(format!("Forward pass failed: {}", e)))?;

            // Create output dictionary
            let dict = pyo3::types::PyDict::new(py);
            dict.set_item(
                "last_hidden_state",
                PyTensor::from_tensor(outputs.last_hidden_state),
            )?;
            if let Some(pooler) = outputs.pooler_output {
                dict.set_item("pooler_output", PyTensor::from_tensor(pooler))?;
            }

            Ok(dict.into())
        })
    }

    /// Python's __call__ method
    pub fn __call__(
        &self,
        input_ids: &PyTensor,
        attention_mask: Option<&PyTensor>,
        token_type_ids: Option<&PyTensor>,
    ) -> PyResult<PyObject> {
        self.forward(input_ids, attention_mask, token_type_ids)
    }
}

/// GPT-2 Model wrapper
#[pyclass(name = "GPT2Model", module = "trustformers", extends = PyPreTrainedModel)]
pub struct PyGPT2Model {
    inner: Gpt2Model,
}

#[pymethods]
impl PyGPT2Model {
    /// Create a new GPT-2 model
    #[new]
    #[pyo3(signature = (config=None))]
    pub fn new(
        py: Python<'_>,
        config: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<(Self, PyPreTrainedModel)> {
        let gpt2_config = if let Some(cfg) = config {
            parse_gpt2_config(cfg)?
        } else {
            Gpt2Config::default()
        };

        let model = Gpt2Model::new(gpt2_config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create GPT-2 model: {}", e)))?;

        let config_dict = gpt2_config_to_dict(py, &gpt2_config)?;

        Ok((
            PyGPT2Model { inner: model },
            PyPreTrainedModel {
                config: config_dict.into(),
            },
        ))
    }

    /// Load from pretrained model
    #[staticmethod]
    #[pyo3(signature = (model_name_or_path, **kwargs))]
    pub fn from_pretrained(
        py: Python<'_>,
        model_name_or_path: &str,
        kwargs: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyGPT2Model>> {
        // Load config from HuggingFace Hub
        let config_json = load_config_from_hub(model_name_or_path, None)
            .map_err(|e| PyValueError::new_err(format!("Failed to load config from hub: {}", e)))?;

        // Parse config into Gpt2Config
        let config_value: Value = serde_json::from_str(&config_json)
            .map_err(|e| PyValueError::new_err(format!("Failed to parse config JSON: {}", e)))?;
        let mut config = Gpt2Config::default();
        if let Some(vocab_size) = config_value.get("vocab_size").and_then(|v| v.as_u64()) {
            config.vocab_size = vocab_size as usize;
        }
        if let Some(n_embd) = config_value.get("n_embd").and_then(|v| v.as_u64()) {
            config.n_embd = n_embd as usize;
        }
        if let Some(n_layer) = config_value.get("n_layer").and_then(|v| v.as_u64()) {
            config.n_layer = n_layer as usize;
        }
        if let Some(n_head) = config_value.get("n_head").and_then(|v| v.as_u64()) {
            config.n_head = n_head as usize;
        }
        if let Some(n_positions) = config_value.get("n_positions").and_then(|v| v.as_u64()) {
            config.n_positions = n_positions as usize;
        }

        let mut model = Gpt2Model::new(config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create model: {}", e)))?;

        // Load weights from hub if available
        if let Ok(mut weight_loader) = try_load_weights(model_name_or_path) {
            if let Err(e) = load_model_weights(&mut model, &mut weight_loader) {
                // Log warning but don't fail - model can work without weights
                eprintln!(
                    "Warning: Failed to load weights for {}: {}",
                    model_name_or_path, e
                );
            } else {
                println!(
                    "Successfully loaded weights for model from {}",
                    model_name_or_path
                );
            }
        } else {
            // Weights not found locally - this is OK, model will use random initialization
            println!(
                "No local weights found for {}, using random initialization",
                model_name_or_path
            );
        }

        let config_dict = gpt2_config_to_dict(py, &config)?;

        Py::new(
            py,
            (
                PyGPT2Model { inner: model },
                PyPreTrainedModel {
                    config: config_dict.into(),
                },
            ),
        )
    }

    /// Forward pass
    #[pyo3(signature = (input_ids, attention_mask=None, past_key_values=None))]
    pub fn forward(
        &self,
        input_ids: &PyTensor,
        attention_mask: Option<&PyTensor>,
        past_key_values: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<PyObject> {
        Python::with_gil(|py| {
            // Create TokenizedInput from the tensor arguments
            use trustformers_core::traits::TokenizedInput;

            let tokenized_input = TokenizedInput {
                input_ids: input_ids
                    .inner
                    .to_vec_f32()
                    .map_err(trustformers_error_to_py_err)?
                    .iter()
                    .map(|&x| x as u32)
                    .collect(),
                attention_mask: attention_mask.map_or_else(
                    || vec![1u8; input_ids.inner.shape()[0]],
                    |mask| {
                        mask.inner
                            .to_vec_f32()
                            .unwrap_or_default()
                            .iter()
                            .map(|&x| x as u8)
                            .collect()
                    },
                ),
                token_type_ids: None, // GPT-2/LLaMA don't use token type IDs
                special_tokens_mask: None,
                offset_mapping: None,
                overflowing_tokens: None,
            };

            let outputs = self
                .inner
                .forward(tokenized_input)
                .map_err(|e| PyValueError::new_err(format!("Forward pass failed: {}", e)))?;

            let dict = pyo3::types::PyDict::new(py);
            dict.set_item(
                "last_hidden_state",
                PyTensor::from_tensor(outputs.last_hidden_state),
            )?;

            Ok(dict.into())
        })
    }

    /// Generate text (placeholder implementation)
    #[pyo3(signature = (input_ids, max_length=50, temperature=1.0, top_k=50, top_p=0.95))]
    pub fn generate(
        &self,
        input_ids: &PyTensor,
        max_length: usize,
        temperature: f32,
        top_k: usize,
        top_p: f32,
    ) -> PyResult<PyTensor> {
        // Convert PyTensor to Vec<u32>
        let input_token_ids = match &input_ids.inner {
            Tensor::I64(arr) => {
                // Convert 1D array to Vec<u32>
                arr.iter().map(|&x| x as u32).collect::<Vec<u32>>()
            },
            Tensor::F32(arr) => {
                // Convert 1D array to Vec<u32> (for tokenized inputs)
                arr.iter().map(|&x| x as u32).collect::<Vec<u32>>()
            },
            _ => {
                return Err(PyValueError::new_err(
                    "Input tensor must contain integer token IDs",
                ))
            },
        };

        // Placeholder generation implementation
        // In a real implementation, this would use the model's forward pass
        let generated_tokens = if input_token_ids.len() < max_length {
            let mut extended = input_token_ids.clone();
            // Add some dummy tokens for demonstration
            for _ in input_token_ids.len()..max_length.min(input_token_ids.len() + 10) {
                extended.push(50256); // End-of-text token for GPT-2
            }
            extended
        } else {
            input_token_ids
        };

        // Convert Vec<u32> back to PyTensor
        let token_data: Vec<i64> = generated_tokens.iter().map(|&x| x as i64).collect();
        let output_tensor = Tensor::I64(
            ArrayD::from_shape_vec(IxDyn(&[token_data.len()]), token_data)
                .map_err(|e| PyValueError::new_err(format!("Failed to create tensor: {}", e)))?,
        );

        Ok(PyTensor {
            inner: output_tensor,
            variable: None,
        })
    }
}

/// T5 Model wrapper
#[pyclass(name = "T5Model", module = "trustformers", extends = PyPreTrainedModel)]
pub struct PyT5Model {
    inner: T5Model,
}

#[pymethods]
impl PyT5Model {
    /// Create a new T5 model
    #[new]
    #[pyo3(signature = (config=None))]
    pub fn new(
        py: Python<'_>,
        config: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<(Self, PyPreTrainedModel)> {
        let t5_config =
            if let Some(cfg) = config { parse_t5_config(cfg)? } else { T5Config::default() };

        let model = T5Model::new(t5_config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create T5 model: {}", e)))?;

        let config_dict = t5_config_to_dict(py, &t5_config)?;

        Ok((
            PyT5Model { inner: model },
            PyPreTrainedModel {
                config: config_dict.into(),
            },
        ))
    }

    /// Load a pretrained T5 model from HuggingFace Hub
    #[staticmethod]
    #[pyo3(signature = (model_name_or_path, **kwargs))]
    pub fn from_pretrained(
        py: Python<'_>,
        model_name_or_path: &str,
        kwargs: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyT5Model>> {
        // Load config from HuggingFace Hub
        let config_json = load_config_from_hub(model_name_or_path, None)
            .map_err(|e| PyValueError::new_err(format!("Failed to load config from hub: {}", e)))?;

        // Parse config into T5Config
        let config_value: Value = serde_json::from_str(&config_json)
            .map_err(|e| PyValueError::new_err(format!("Failed to parse config JSON: {}", e)))?;
        let mut config = T5Config::default();

        // Update config with loaded values
        if let Some(vocab_size) = config_value.get("vocab_size").and_then(|v| v.as_u64()) {
            config.vocab_size = vocab_size as usize;
        }
        if let Some(d_model) = config_value.get("d_model").and_then(|v| v.as_u64()) {
            config.d_model = d_model as usize;
        }
        if let Some(d_ff) = config_value.get("d_ff").and_then(|v| v.as_u64()) {
            config.d_ff = d_ff as usize;
        }
        if let Some(num_layers) = config_value.get("num_layers").and_then(|v| v.as_u64()) {
            config.num_layers = num_layers as usize;
        }
        if let Some(num_heads) = config_value.get("num_heads").and_then(|v| v.as_u64()) {
            config.num_heads = num_heads as usize;
        }

        let mut model = T5Model::new(config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create model: {}", e)))?;

        // Load weights from hub if available
        if let Ok(mut weight_loader) = try_load_weights(model_name_or_path) {
            if let Err(e) = load_model_weights(&mut model, &mut weight_loader) {
                // Log warning but don't fail - model can work without weights
                eprintln!(
                    "Warning: Failed to load weights for {}: {}",
                    model_name_or_path, e
                );
            } else {
                println!(
                    "Successfully loaded weights for model from {}",
                    model_name_or_path
                );
            }
        } else {
            // Weights not found locally - this is OK, model will use random initialization
            println!(
                "No local weights found for {}, using random initialization",
                model_name_or_path
            );
        }

        let config_dict = t5_config_to_dict(py, &config)?;

        Py::new(
            py,
            (
                PyT5Model { inner: model },
                PyPreTrainedModel {
                    config: config_dict.into(),
                },
            ),
        )
    }

    /// Forward pass
    #[pyo3(signature = (input_ids=None, attention_mask=None, decoder_input_ids=None, decoder_attention_mask=None))]
    pub fn forward(
        &self,
        input_ids: Option<&PyTensor>,
        attention_mask: Option<&PyTensor>,
        decoder_input_ids: Option<&PyTensor>,
        decoder_attention_mask: Option<&PyTensor>,
    ) -> PyResult<PyObject> {
        use trustformers_core::traits::TokenizedInput;

        // Convert input_ids (required for T5)
        let input_token_ids = if let Some(input_tensor) = input_ids {
            match &input_tensor.inner {
                Tensor::I64(arr) => arr.iter().map(|&x| x as u32).collect::<Vec<u32>>(),
                Tensor::F32(arr) => arr.iter().map(|&x| x as u32).collect::<Vec<u32>>(),
                _ => {
                    return Err(PyValueError::new_err(
                        "Input tensor must contain integer token IDs",
                    ))
                },
            }
        } else {
            return Err(PyValueError::new_err(
                "input_ids is required for T5 forward pass",
            ));
        };

        // Convert attention_mask (use default if not provided)
        let input_attention_mask = if let Some(mask_tensor) = attention_mask {
            match &mask_tensor.inner {
                Tensor::I64(arr) => arr.iter().map(|&x| x as u8).collect::<Vec<u8>>(),
                Tensor::F32(arr) => arr.iter().map(|&x| x as u8).collect::<Vec<u8>>(),
                _ => {
                    return Err(PyValueError::new_err(
                        "Attention mask must contain integer values",
                    ))
                },
            }
        } else {
            vec![1u8; input_token_ids.len()] // Default to all ones
        };

        // Convert decoder_input_ids (optional for T5)
        let decoder_tokenized_input = if let Some(decoder_tensor) = decoder_input_ids {
            let decoder_token_ids = match &decoder_tensor.inner {
                Tensor::I64(arr) => arr.iter().map(|&x| x as u32).collect::<Vec<u32>>(),
                Tensor::F32(arr) => arr.iter().map(|&x| x as u32).collect::<Vec<u32>>(),
                _ => {
                    return Err(PyValueError::new_err(
                        "Decoder input tensor must contain integer token IDs",
                    ))
                },
            };

            // Convert decoder attention mask if provided
            let decoder_att_mask = if let Some(dec_mask_tensor) = decoder_attention_mask {
                match &dec_mask_tensor.inner {
                    Tensor::I64(arr) => arr.iter().map(|&x| x as u8).collect::<Vec<u8>>(),
                    Tensor::F32(arr) => arr.iter().map(|&x| x as u8).collect::<Vec<u8>>(),
                    _ => {
                        return Err(PyValueError::new_err(
                            "Decoder attention mask must contain integer values",
                        ))
                    },
                }
            } else {
                vec![1u8; decoder_token_ids.len()] // Default to all ones
            };

            Some(TokenizedInput {
                input_ids: decoder_token_ids,
                attention_mask: decoder_att_mask,
                token_type_ids: None,
                special_tokens_mask: None,
                offset_mapping: None,
                overflowing_tokens: None,
            })
        } else {
            None
        };

        // Create T5Input
        let input = trustformers_models::t5::T5Input {
            input_ids: TokenizedInput {
                input_ids: input_token_ids,
                attention_mask: input_attention_mask,
                token_type_ids: None,
                special_tokens_mask: None,
                offset_mapping: None,
                overflowing_tokens: None,
            },
            decoder_input_ids: decoder_tokenized_input,
            encoder_outputs: None,
        };

        // Forward pass
        let output = self
            .inner
            .forward(input)
            .map_err(|e| PyValueError::new_err(format!("T5 forward pass failed: {}", e)))?;

        // Convert output to Python dictionary
        Python::with_gil(|py| {
            let dict = pyo3::types::PyDict::new(py);

            // Convert last_hidden_state to PyTensor
            let last_hidden_py = PyTensor {
                inner: output.last_hidden_state,
                variable: None,
            };
            dict.set_item("last_hidden_state", last_hidden_py)?;

            // Add encoder_last_hidden_state if available
            if let Some(encoder_hidden) = output.encoder_last_hidden_state {
                let encoder_hidden_py = PyTensor {
                    inner: encoder_hidden,
                    variable: None,
                };
                dict.set_item("encoder_last_hidden_state", encoder_hidden_py)?;
            }

            Ok(dict.into())
        })
    }
}

/// LLaMA Model wrapper
#[pyclass(name = "LlamaModel", module = "trustformers", extends = PyPreTrainedModel)]
pub struct PyLlamaModel {
    inner: LlamaModel,
}

#[pymethods]
impl PyLlamaModel {
    /// Create a new LLaMA model
    #[new]
    #[pyo3(signature = (config=None))]
    pub fn new(
        py: Python<'_>,
        config: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<(Self, PyPreTrainedModel)> {
        let llama_config = if let Some(cfg) = config {
            parse_llama_config(cfg)?
        } else {
            LlamaConfig::default()
        };

        let model = LlamaModel::new(llama_config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create LLaMA model: {}", e)))?;

        let config_dict = llama_config_to_dict(py, &llama_config)?;

        Ok((
            PyLlamaModel { inner: model },
            PyPreTrainedModel {
                config: config_dict.into(),
            },
        ))
    }

    /// Load a pretrained LLaMA model from HuggingFace Hub
    #[staticmethod]
    #[pyo3(signature = (model_name_or_path, **kwargs))]
    pub fn from_pretrained(
        py: Python<'_>,
        model_name_or_path: &str,
        kwargs: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyLlamaModel>> {
        // Load config from HuggingFace Hub
        let config_json = load_config_from_hub(model_name_or_path, None)
            .map_err(|e| PyValueError::new_err(format!("Failed to load config from hub: {}", e)))?;

        // Parse config into LlamaConfig
        let config_value: Value = serde_json::from_str(&config_json)
            .map_err(|e| PyValueError::new_err(format!("Failed to parse config JSON: {}", e)))?;
        let mut config = LlamaConfig::default();

        // Update config with loaded values
        if let Some(vocab_size) = config_value.get("vocab_size").and_then(|v| v.as_u64()) {
            config.vocab_size = vocab_size as usize;
        }
        if let Some(hidden_size) = config_value.get("hidden_size").and_then(|v| v.as_u64()) {
            config.hidden_size = hidden_size as usize;
        }
        if let Some(intermediate_size) =
            config_value.get("intermediate_size").and_then(|v| v.as_u64())
        {
            config.intermediate_size = intermediate_size as usize;
        }
        if let Some(num_layers) = config_value.get("num_hidden_layers").and_then(|v| v.as_u64()) {
            config.num_hidden_layers = num_layers as usize;
        }
        if let Some(num_heads) = config_value.get("num_attention_heads").and_then(|v| v.as_u64()) {
            config.num_attention_heads = num_heads as usize;
        }
        if let Some(max_pos) = config_value.get("max_position_embeddings").and_then(|v| v.as_u64())
        {
            config.max_position_embeddings = max_pos as usize;
        }

        let mut model = LlamaModel::new(config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create model: {}", e)))?;

        // Load weights from hub if available
        if let Ok(mut weight_loader) = try_load_weights(model_name_or_path) {
            if let Err(e) = load_model_weights(&mut model, &mut weight_loader) {
                // Log warning but don't fail - model can work without weights
                eprintln!(
                    "Warning: Failed to load weights for {}: {}",
                    model_name_or_path, e
                );
            } else {
                println!(
                    "Successfully loaded weights for model from {}",
                    model_name_or_path
                );
            }
        } else {
            // Weights not found locally - this is OK, model will use random initialization
            println!(
                "No local weights found for {}, using random initialization",
                model_name_or_path
            );
        }

        let config_dict = llama_config_to_dict(py, &config)?;

        Py::new(
            py,
            (
                PyLlamaModel { inner: model },
                PyPreTrainedModel {
                    config: config_dict.into(),
                },
            ),
        )
    }

    /// Forward pass
    #[pyo3(signature = (input_ids, attention_mask=None, position_ids=None))]
    pub fn forward(
        &self,
        input_ids: &PyTensor,
        attention_mask: Option<&PyTensor>,
        position_ids: Option<&PyTensor>,
    ) -> PyResult<PyObject> {
        Python::with_gil(|py| {
            // Convert input tensor to token IDs for LLaMA
            let input_token_ids = input_ids
                .inner
                .to_vec_f32()
                .map_err(trustformers_error_to_py_err)?
                .iter()
                .map(|&x| x as u32)
                .collect::<Vec<u32>>();

            let outputs = self
                .inner
                .forward(input_token_ids)
                .map_err(|e| PyValueError::new_err(format!("Forward pass failed: {}", e)))?;

            let dict = pyo3::types::PyDict::new(py);
            dict.set_item("last_hidden_state", PyTensor::from_tensor(outputs))?;

            Ok(dict.into())
        })
    }
}

// Helper functions for config parsing
fn parse_bert_config(config_dict: &Bound<'_, PyAny>) -> PyResult<BertConfig> {
    let dict = config_dict.downcast::<pyo3::types::PyDict>()?;

    let mut config = BertConfig::default();

    if let Ok(v) = dict.get_item("vocab_size") {
        if let Some(val) = v {
            config.vocab_size = val.extract()?;
        }
    }
    if let Ok(v) = dict.get_item("hidden_size") {
        if let Some(val) = v {
            config.hidden_size = val.extract()?;
        }
    }
    if let Ok(v) = dict.get_item("num_hidden_layers") {
        if let Some(val) = v {
            config.num_hidden_layers = val.extract()?;
        }
    }
    if let Ok(v) = dict.get_item("num_attention_heads") {
        if let Some(val) = v {
            config.num_attention_heads = val.extract()?;
        }
    }

    Ok(config)
}

fn config_to_dict<'py>(
    py: Python<'py>,
    config: &BertConfig,
) -> PyResult<Bound<'py, pyo3::types::PyDict>> {
    let dict = pyo3::types::PyDict::new(py);
    dict.set_item("vocab_size", config.vocab_size)?;
    dict.set_item("hidden_size", config.hidden_size)?;
    dict.set_item("num_hidden_layers", config.num_hidden_layers)?;
    dict.set_item("num_attention_heads", config.num_attention_heads)?;
    dict.set_item("intermediate_size", config.intermediate_size)?;
    dict.set_item("hidden_act", &config.hidden_act)?;
    dict.set_item("hidden_dropout_prob", config.hidden_dropout_prob)?;
    dict.set_item(
        "attention_probs_dropout_prob",
        config.attention_probs_dropout_prob,
    )?;
    dict.set_item("max_position_embeddings", config.max_position_embeddings)?;
    dict.set_item("type_vocab_size", config.type_vocab_size)?;
    dict.set_item("initializer_range", config.initializer_range)?;
    dict.set_item("layer_norm_eps", config.layer_norm_eps)?;
    Ok(dict)
}

// Task-specific models using composition pattern

/// BERT for Sequence Classification
#[pyclass(name = "BertForSequenceClassification", module = "trustformers")]
pub struct PyBertForSequenceClassification {
    bert: BertModel,
    classifier: PyObject, // Linear layer for classification
    config: BertConfig,
    num_labels: usize,
}

#[pymethods]
impl PyBertForSequenceClassification {
    #[new]
    #[pyo3(signature = (config=None, num_labels=2))]
    pub fn new(
        py: Python<'_>,
        config: Option<&Bound<'_, PyAny>>,
        num_labels: usize,
    ) -> PyResult<Self> {
        let bert_config = if let Some(cfg) = config {
            parse_bert_config(cfg)?
        } else {
            BertConfig::default()
        };

        let bert = BertModel::new(bert_config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create BERT model: {}", e)))?;

        // Create classifier as placeholder (would be actual linear layer in full implementation)
        let classifier = py.None();

        Ok(PyBertForSequenceClassification {
            bert,
            classifier,
            config: bert_config,
            num_labels,
        })
    }

    #[staticmethod]
    #[pyo3(signature = (model_name_or_path, **kwargs))]
    pub fn from_pretrained(
        py: Python<'_>,
        model_name_or_path: &str,
        kwargs: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyBertForSequenceClassification>> {
        // Load config from HuggingFace Hub
        let config_json = load_config_from_hub(model_name_or_path, None)
            .map_err(|e| PyValueError::new_err(format!("Failed to load config from hub: {}", e)))?;

        // Parse config into BertConfig
        let config_value: Value = serde_json::from_str(&config_json)
            .map_err(|e| PyValueError::new_err(format!("Failed to parse config JSON: {}", e)))?;
        let mut config = BertConfig::default();
        if let Some(vocab_size) = config_value.get("vocab_size").and_then(|v| v.as_u64()) {
            config.vocab_size = vocab_size as usize;
        }
        if let Some(hidden_size) = config_value.get("hidden_size").and_then(|v| v.as_u64()) {
            config.hidden_size = hidden_size as usize;
        }
        if let Some(num_layers) = config_value.get("num_hidden_layers").and_then(|v| v.as_u64()) {
            config.num_hidden_layers = num_layers as usize;
        }
        if let Some(num_heads) = config_value.get("num_attention_heads").and_then(|v| v.as_u64()) {
            config.num_attention_heads = num_heads as usize;
        }

        // Extract number of labels from config
        let num_labels =
            config_value.get("num_labels").and_then(|v| v.as_u64()).unwrap_or(2) as usize;

        let mut model = BertModel::new(config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create model: {}", e)))?;

        // Load weights from hub if available
        if let Ok(mut weight_loader) = try_load_weights(model_name_or_path) {
            if let Err(e) = load_model_weights(&mut model, &mut weight_loader) {
                // Log warning but don't fail - model can work without weights
                eprintln!(
                    "Warning: Failed to load weights for {}: {}",
                    model_name_or_path, e
                );
            } else {
                println!(
                    "Successfully loaded weights for model from {}",
                    model_name_or_path
                );
            }
        } else {
            // Weights not found locally - this is OK, model will use random initialization
            println!(
                "No local weights found for {}, using random initialization",
                model_name_or_path
            );
        }

        Py::new(
            py,
            PyBertForSequenceClassification {
                bert: model,
                classifier: py.None(),
                config,
                num_labels,
            },
        )
    }

    #[pyo3(signature = (input_ids, attention_mask=None, token_type_ids=None, labels=None))]
    pub fn forward(
        &self,
        input_ids: &PyTensor,
        attention_mask: Option<&PyTensor>,
        token_type_ids: Option<&PyTensor>,
        labels: Option<&PyTensor>,
    ) -> PyResult<PyObject> {
        Python::with_gil(|py| {
            use trustformers_core::traits::TokenizedInput;

            let tokenized_input = TokenizedInput {
                input_ids: input_ids
                    .inner
                    .to_vec_f32()
                    .map_err(trustformers_error_to_py_err)?
                    .iter()
                    .map(|&x| x as u32)
                    .collect(),
                attention_mask: attention_mask.map_or_else(
                    || vec![1u8; input_ids.inner.shape()[0]],
                    |mask| {
                        mask.inner
                            .to_vec_f32()
                            .unwrap_or_default()
                            .iter()
                            .map(|&x| x as u8)
                            .collect()
                    },
                ),
                token_type_ids: token_type_ids.map(|t| {
                    t.inner.to_vec_f32().unwrap_or_default().iter().map(|&x| x as u32).collect()
                }),
                special_tokens_mask: None,
                offset_mapping: None,
                overflowing_tokens: None,
            };

            let bert_outputs = self
                .bert
                .forward(tokenized_input)
                .map_err(|e| PyValueError::new_err(format!("BERT forward pass failed: {}", e)))?;

            // Use pooler output or [CLS] token for classification
            let logits = if let Some(pooler) = bert_outputs.pooler_output {
                pooler // Already pooled representation
            } else {
                // Use [CLS] token (first token) from last hidden state
                bert_outputs
                    .last_hidden_state
                    .slice(0, 0, 1)
                    .map_err(|e| PyValueError::new_err(format!("Failed to slice tensor: {}", e)))?
            };

            let dict = pyo3::types::PyDict::new(py);
            dict.set_item("logits", PyTensor::from_tensor(logits.clone()))?;

            // Calculate loss if labels provided
            if let Some(labels_tensor) = labels {
                // Compute cross-entropy loss for classification
                let loss_value = compute_cross_entropy_loss(&logits, &labels_tensor.inner)
                    .map_err(|e| {
                        PyValueError::new_err(format!("Loss calculation failed: {}", e))
                    })?;

                let loss = PyTensor::from_tensor(Tensor::scalar(loss_value).map_err(|e| {
                    PyValueError::new_err(format!("Failed to create loss tensor: {}", e))
                })?);
                dict.set_item("loss", loss)?;
            }

            Ok(dict.into())
        })
    }

    /// Python's __call__ method
    pub fn __call__(
        &self,
        input_ids: &PyTensor,
        attention_mask: Option<&PyTensor>,
        token_type_ids: Option<&PyTensor>,
        labels: Option<&PyTensor>,
    ) -> PyResult<PyObject> {
        self.forward(input_ids, attention_mask, token_type_ids, labels)
    }
}

/// GPT2 for Language Modeling Head
#[pyclass(name = "GPT2LMHeadModel", module = "trustformers")]
pub struct PyGPT2LMHeadModel {
    transformer: Gpt2Model,
    lm_head: PyObject, // Linear layer for language modeling
    config: Gpt2Config,
}

#[pymethods]
impl PyGPT2LMHeadModel {
    #[new]
    #[pyo3(signature = (config=None))]
    pub fn new(py: Python<'_>, config: Option<&Bound<'_, PyAny>>) -> PyResult<Self> {
        let gpt2_config = if let Some(cfg) = config {
            parse_gpt2_config(cfg)?
        } else {
            Gpt2Config::default()
        };

        let transformer = Gpt2Model::new(gpt2_config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create GPT-2 model: {}", e)))?;

        Ok(PyGPT2LMHeadModel {
            transformer,
            lm_head: py.None(),
            config: gpt2_config,
        })
    }

    #[staticmethod]
    #[pyo3(signature = (model_name_or_path, **kwargs))]
    pub fn from_pretrained(
        py: Python<'_>,
        model_name_or_path: &str,
        kwargs: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyGPT2LMHeadModel>> {
        // Load config from HuggingFace Hub
        let config_json = load_config_from_hub(model_name_or_path, None)
            .map_err(|e| PyValueError::new_err(format!("Failed to load config from hub: {}", e)))?;

        // Parse config into Gpt2Config
        let config_value: Value = serde_json::from_str(&config_json)
            .map_err(|e| PyValueError::new_err(format!("Failed to parse config JSON: {}", e)))?;
        let mut config = Gpt2Config::default();
        if let Some(vocab_size) = config_value.get("vocab_size").and_then(|v| v.as_u64()) {
            config.vocab_size = vocab_size as usize;
        }
        if let Some(n_embd) = config_value.get("n_embd").and_then(|v| v.as_u64()) {
            config.n_embd = n_embd as usize;
        }
        if let Some(n_layer) = config_value.get("n_layer").and_then(|v| v.as_u64()) {
            config.n_layer = n_layer as usize;
        }
        if let Some(n_head) = config_value.get("n_head").and_then(|v| v.as_u64()) {
            config.n_head = n_head as usize;
        }
        if let Some(n_positions) = config_value.get("n_positions").and_then(|v| v.as_u64()) {
            config.n_positions = n_positions as usize;
        }

        let mut model = Gpt2Model::new(config.clone())
            .map_err(|e| PyValueError::new_err(format!("Failed to create model: {}", e)))?;

        // Load weights from hub if available
        if let Ok(mut weight_loader) = try_load_weights(model_name_or_path) {
            if let Err(e) = load_model_weights(&mut model, &mut weight_loader) {
                // Log warning but don't fail - model can work without weights
                eprintln!(
                    "Warning: Failed to load weights for {}: {}",
                    model_name_or_path, e
                );
            } else {
                println!(
                    "Successfully loaded weights for model from {}",
                    model_name_or_path
                );
            }
        } else {
            // Weights not found locally - this is OK, model will use random initialization
            println!(
                "No local weights found for {}, using random initialization",
                model_name_or_path
            );
        }

        Py::new(
            py,
            PyGPT2LMHeadModel {
                transformer: model,
                lm_head: py.None(),
                config,
            },
        )
    }

    #[pyo3(signature = (input_ids, attention_mask=None, labels=None))]
    pub fn forward(
        &self,
        input_ids: &PyTensor,
        attention_mask: Option<&PyTensor>,
        labels: Option<&PyTensor>,
    ) -> PyResult<PyObject> {
        Python::with_gil(|py| {
            use trustformers_core::traits::TokenizedInput;

            let tokenized_input = TokenizedInput {
                input_ids: input_ids
                    .inner
                    .to_vec_f32()
                    .map_err(trustformers_error_to_py_err)?
                    .iter()
                    .map(|&x| x as u32)
                    .collect(),
                attention_mask: attention_mask.map_or_else(
                    || vec![1u8; input_ids.inner.shape()[0]],
                    |mask| {
                        mask.inner
                            .to_vec_f32()
                            .unwrap_or_default()
                            .iter()
                            .map(|&x| x as u8)
                            .collect()
                    },
                ),
                token_type_ids: None, // GPT-2 doesn't use token type IDs
                special_tokens_mask: None,
                offset_mapping: None,
                overflowing_tokens: None,
            };

            let transformer_outputs = self.transformer.forward(tokenized_input).map_err(|e| {
                PyValueError::new_err(format!("Transformer forward pass failed: {}", e))
            })?;

            // Apply language modeling head (would be actual linear layer)
            let logits = transformer_outputs.last_hidden_state; // Placeholder

            let dict = pyo3::types::PyDict::new(py);
            dict.set_item("logits", PyTensor::from_tensor(logits.clone()))?;

            // Calculate loss if labels provided
            if let Some(labels_tensor) = labels {
                // Compute cross-entropy loss for language modeling (next token prediction)
                let loss_value = compute_language_modeling_loss(&logits, &labels_tensor.inner)
                    .map_err(|e| {
                        PyValueError::new_err(format!(
                            "Language modeling loss calculation failed: {}",
                            e
                        ))
                    })?;

                let loss = PyTensor::from_tensor(Tensor::scalar(loss_value).map_err(|e| {
                    PyValueError::new_err(format!("Failed to create loss tensor: {}", e))
                })?);
                dict.set_item("loss", loss)?;
            }

            Ok(dict.into())
        })
    }

    /// Generate text with the language model
    #[pyo3(signature = (input_ids, max_length=50, temperature=1.0, do_sample=true))]
    pub fn generate(
        &self,
        input_ids: &PyTensor,
        max_length: usize,
        temperature: f32,
        do_sample: bool,
    ) -> PyResult<PyTensor> {
        // Convert PyTensor to Vec<u32>
        let input_token_ids = match &input_ids.inner {
            Tensor::I64(arr) => {
                // Convert 1D array to Vec<u32>
                arr.iter().map(|&x| x as u32).collect::<Vec<u32>>()
            },
            Tensor::F32(arr) => {
                // Convert 1D array to Vec<u32> (for tokenized inputs)
                arr.iter().map(|&x| x as u32).collect::<Vec<u32>>()
            },
            _ => {
                return Err(PyValueError::new_err(
                    "Input tensor must contain integer token IDs",
                ))
            },
        };

        // Call the actual generate method from Rust
        // Use sampling parameters based on do_sample flag
        let (top_k, top_p) = if do_sample {
            (Some(50), Some(0.9)) // Use sampling
        } else {
            (None, None) // Greedy decoding
        };

        // Placeholder generation implementation
        // In a real implementation, this would use the model's forward pass
        let generated_tokens = if input_token_ids.len() < max_length {
            let mut extended = input_token_ids.clone();
            // Add some dummy tokens for demonstration
            for _ in input_token_ids.len()..max_length.min(input_token_ids.len() + 10) {
                extended.push(50256); // End-of-text token for GPT-2
            }
            extended
        } else {
            input_token_ids
        };

        // Convert Vec<u32> back to PyTensor
        let token_data: Vec<i64> = generated_tokens.iter().map(|&x| x as i64).collect();
        let output_tensor = Tensor::I64(
            ArrayD::from_shape_vec(IxDyn(&[token_data.len()]), token_data)
                .map_err(|e| PyValueError::new_err(format!("Failed to create tensor: {}", e)))?,
        );

        Ok(PyTensor {
            inner: output_tensor,
            variable: None,
        })
    }
}
