//! High-level builder API for fine-tuning and inference.
//!
//! This module provides ergonomic one-liner APIs that wrap the CLI's orchestration
//! logic into composable builder structs.
//!
//! # Fine-tuning
//!
//! ```rust,no_run
//! # async fn example() -> pmetal_core::Result<()> {
//! let result = pmetal::easy::finetune("Qwen/Qwen3-0.6B", "data.jsonl")
//!     .lora(16, 32.0)
//!     .epochs(3)
//!     .learning_rate(2e-4)
//!     .output("./output")
//!     .run()
//!     .await?;
//! println!("Final loss: {:.4}", result.final_loss);
//! # Ok(())
//! # }
//! ```
//!
//! # Inference
//!
//! ```rust,no_run
//! # async fn example() -> pmetal_core::Result<()> {
//! let result = pmetal::easy::infer("Qwen/Qwen3-0.6B")
//!     .temperature(0.7)
//!     .generate("What is 2+2?")
//!     .await?;
//! println!("{}", result.text);
//! # Ok(())
//! # }
//! ```

use std::path::{Path, PathBuf};

use pmetal_core::{LoraConfig, PMetalError, Result, TrainingCallback, TrainingConfig};
use pmetal_data::{DataLoaderConfig, DatasetFormat, Tokenizer, TrainingDataset};
use pmetal_lora::{DynamicLoraModel, TrainableModel};
use pmetal_mlx::Builder as _;
use pmetal_models::{
    DynamicModel, GenerationConfig, generate_cached_async, generate_cached_async_streaming,
};
use pmetal_trainer::{CheckpointManager, MetricsJsonCallback, TrainingLoop, TrainingLoopConfig};

/// Map any Display error to `PMetalError::Training`.
fn training_err(e: impl std::fmt::Display) -> PMetalError {
    PMetalError::Training(e.to_string())
}

/// Map any Display error to `PMetalError::Mlx`.
fn mlx_err(e: impl std::fmt::Display) -> PMetalError {
    PMetalError::Mlx(e.to_string())
}

/// Map any Display error to `PMetalError::ModelLoad`.
fn model_err(e: impl std::fmt::Display) -> PMetalError {
    PMetalError::ModelLoad(e.to_string())
}

/// Check if a string looks like a HuggingFace model ID (e.g., "org/model")
/// rather than a local filesystem path.
fn is_hf_model_id(s: &str) -> bool {
    !s.starts_with('/') && !s.starts_with('.') && s.contains('/')
}

/// Result of a fine-tuning run.
#[derive(Debug, Clone)]
pub struct FinetuneResult {
    /// Final training loss.
    pub final_loss: f64,
    /// Total training steps completed.
    pub total_steps: usize,
    /// Total tokens processed.
    pub total_tokens: usize,
    /// Output directory path.
    pub output_dir: PathBuf,
    /// Path to saved LoRA weights.
    pub lora_weights_path: PathBuf,
}

/// Result of an inference run.
#[derive(Debug, Clone)]
pub struct InferResult {
    /// Generated text.
    pub text: String,
    /// Number of tokens generated.
    pub tokens_generated: usize,
    /// Generation throughput.
    pub tokens_per_sec: f64,
}

/// Create a fine-tuning builder for the given model and dataset.
///
/// # Arguments
/// * `model` - HuggingFace model ID (e.g., `"Qwen/Qwen3-0.6B"`) or local path
/// * `dataset` - Path to JSONL training dataset
pub fn finetune(model: impl Into<String>, dataset: impl Into<String>) -> FinetuneBuilder {
    FinetuneBuilder::new(model.into(), dataset.into())
}

/// Create a DPO fine-tuning builder.
pub fn dpo(model: impl Into<String>, dataset: impl Into<String>) -> PreferenceTuneBuilder {
    PreferenceTuneBuilder::new(
        model.into(),
        dataset.into(),
        PreferenceMethod::Dpo(pmetal_trainer::DpoConfig::default()),
    )
}

/// Create a SimPO fine-tuning builder.
pub fn simpo(model: impl Into<String>, dataset: impl Into<String>) -> PreferenceTuneBuilder {
    PreferenceTuneBuilder::new(
        model.into(),
        dataset.into(),
        PreferenceMethod::Simpo(pmetal_trainer::SimpoConfig::default()),
    )
}

/// Create an ORPO fine-tuning builder.
pub fn orpo(model: impl Into<String>, dataset: impl Into<String>) -> PreferenceTuneBuilder {
    PreferenceTuneBuilder::new(
        model.into(),
        dataset.into(),
        PreferenceMethod::Orpo(pmetal_trainer::OrpoConfig::default()),
    )
}

/// Create a KTO fine-tuning builder.
pub fn kto(model: impl Into<String>, dataset: impl Into<String>) -> PreferenceTuneBuilder {
    PreferenceTuneBuilder::new(
        model.into(),
        dataset.into(),
        PreferenceMethod::Kto(pmetal_trainer::KtoConfig::default()),
    )
}

/// Create an inference builder for the given model.
///
/// # Arguments
/// * `model` - HuggingFace model ID (e.g., `"Qwen/Qwen3-0.6B"`) or local path
pub fn infer(model: impl Into<String>) -> InferBuilder {
    InferBuilder::new(model.into())
}

enum PreferenceMethod {
    Dpo(pmetal_trainer::DpoConfig),
    Simpo(pmetal_trainer::SimpoConfig),
    Orpo(pmetal_trainer::OrpoConfig),
    Kto(pmetal_trainer::KtoConfig),
}

/// Builder for fine-tuning a model with LoRA.
///
/// Wraps the full training pipeline: model download, tokenizer loading,
/// dataset preparation, LoRA initialization, training loop, and weight saving.
pub struct FinetuneBuilder {
    model_id: String,
    dataset_path: String,
    eval_dataset_path: Option<String>,
    lora_r: usize,
    lora_alpha: f32,
    learning_rate: f64,
    batch_size: usize,
    num_epochs: usize,
    max_seq_len: usize,
    output_dir: String,
    lora_dropout: f32,
    use_rslora: bool,
    use_dora: bool,
    flash_attention: bool,
    sequence_packing: bool,
    gradient_checkpointing: bool,
    gradient_checkpointing_layers: usize,
    metal_fused_optimizer: bool,
    embedding_lr: Option<f32>,
    callbacks: Vec<Box<dyn TrainingCallback>>,
    metrics_path: Option<PathBuf>,
}

impl FinetuneBuilder {
    fn new(model_id: String, dataset_path: String) -> Self {
        Self {
            model_id,
            dataset_path,
            eval_dataset_path: None,
            lora_r: 16,
            lora_alpha: 32.0,
            learning_rate: 2e-4,
            batch_size: 4,
            num_epochs: 3,
            max_seq_len: 2048,
            output_dir: "./output".to_string(),
            lora_dropout: 0.0,
            use_rslora: false,
            use_dora: false,
            flash_attention: true,
            sequence_packing: true,
            gradient_checkpointing: false,
            gradient_checkpointing_layers: 4,
            metal_fused_optimizer: false,
            embedding_lr: None,
            callbacks: Vec::new(),
            metrics_path: None,
        }
    }

    /// Set LoRA rank and alpha.
    pub fn lora(mut self, r: usize, alpha: f32) -> Self {
        self.lora_r = r;
        self.lora_alpha = alpha;
        self
    }

    /// Set number of training epochs.
    pub fn epochs(mut self, n: usize) -> Self {
        self.num_epochs = n;
        self
    }

    /// Set learning rate.
    pub fn learning_rate(mut self, lr: f64) -> Self {
        self.learning_rate = lr;
        self
    }

    /// Set batch size.
    pub fn batch_size(mut self, n: usize) -> Self {
        self.batch_size = n;
        self
    }

    /// Set maximum sequence length.
    pub fn max_seq_len(mut self, n: usize) -> Self {
        self.max_seq_len = n;
        self
    }

    /// Set output directory for checkpoints and final weights.
    pub fn output(mut self, path: impl Into<String>) -> Self {
        self.output_dir = path.into();
        self
    }

    /// Set LoRA dropout.
    pub fn lora_dropout(mut self, dropout: f32) -> Self {
        self.lora_dropout = dropout;
        self
    }

    /// Enable rank-stabilized LoRA.
    pub fn use_rslora(mut self, enabled: bool) -> Self {
        self.use_rslora = enabled;
        self
    }

    /// Enable DoRA adapters.
    pub fn use_dora(mut self, enabled: bool) -> Self {
        self.use_dora = enabled;
        self
    }

    /// Set evaluation dataset path.
    pub fn eval_dataset(mut self, path: impl Into<String>) -> Self {
        self.eval_dataset_path = Some(path.into());
        self
    }

    /// Enable or disable flash attention.
    pub fn flash_attention(mut self, enabled: bool) -> Self {
        self.flash_attention = enabled;
        self
    }

    /// Enable or disable sequence packing.
    pub fn sequence_packing(mut self, enabled: bool) -> Self {
        self.sequence_packing = enabled;
        self
    }

    /// Enable or disable gradient checkpointing.
    pub fn gradient_checkpointing(mut self, enabled: bool) -> Self {
        self.gradient_checkpointing = enabled;
        self
    }

    /// Set layers per gradient checkpoint block.
    pub fn gradient_checkpointing_layers(mut self, layers: usize) -> Self {
        self.gradient_checkpointing_layers = layers;
        self
    }

    /// Enable or disable Metal fused optimizer.
    pub fn metal_fused_optimizer(mut self, enabled: bool) -> Self {
        self.metal_fused_optimizer = enabled;
        self
    }

    /// Set separate learning rate for embedding layers.
    pub fn embedding_lr(mut self, lr: f32) -> Self {
        self.embedding_lr = Some(lr);
        self
    }

    /// Add a training callback.
    pub fn callback(mut self, callback: Box<dyn TrainingCallback>) -> Self {
        self.callbacks.push(callback);
        self
    }

    /// Write training metrics to JSONL at the given path.
    pub fn metrics_path(mut self, path: impl Into<PathBuf>) -> Self {
        self.metrics_path = Some(path.into());
        self
    }

    /// Run the fine-tuning pipeline.
    pub async fn run(self) -> Result<FinetuneResult> {
        // Resolve model path (download if HuggingFace ID)
        let model_path = resolve_model_path(&self.model_id).await?;

        // Load tokenizer (reads special_tokens_map.json and tokenizer_config.json too)
        let tokenizer = Tokenizer::from_model_dir(&model_path).map_err(|e| {
            PMetalError::ModelLoad(format!("Failed to load tokenizer from {model_path:?}: {e}"))
        })?;

        // Detect chat template
        let chat_template =
            pmetal_data::chat_templates::detect_chat_template(&model_path, &self.model_id);

        // Load and tokenize training dataset
        let train_dataset = TrainingDataset::from_jsonl_tokenized(
            &self.dataset_path,
            &tokenizer,
            DatasetFormat::Auto,
            self.max_seq_len,
            Some(&chat_template),
        )?;

        // Load evaluation dataset if provided
        let eval_dataset = if let Some(ref eval_path) = self.eval_dataset_path {
            Some(TrainingDataset::from_jsonl_tokenized(
                eval_path,
                &tokenizer,
                DatasetFormat::Auto,
                self.max_seq_len,
                Some(&chat_template),
            )?)
        } else {
            None
        };

        // Create LoRA config
        let lora_config = LoraConfig {
            r: self.lora_r,
            alpha: self.lora_alpha,
            dropout: self.lora_dropout,
            use_rslora: self.use_rslora,
            use_dora: self.use_dora,
            ..Default::default()
        };

        // Initialize model with LoRA adapters
        let model =
            DynamicLoraModel::from_pretrained(&model_path, lora_config).map_err(model_err)?;

        // Set up checkpoint manager
        let output_dir = PathBuf::from(&self.output_dir);
        let checkpoint_dir = output_dir.join("checkpoints");
        let checkpoint_manager = CheckpointManager::new(&checkpoint_dir)
            .map_err(training_err)?
            .with_max_checkpoints(3);

        // Build training config
        let training_config = TrainingConfig {
            learning_rate: self.learning_rate,
            batch_size: self.batch_size,
            num_epochs: self.num_epochs,
            max_seq_len: self.max_seq_len,
            output_dir: self.output_dir.clone(),
            ..Default::default()
        };

        let dataloader_config = DataLoaderConfig {
            batch_size: self.batch_size,
            max_seq_len: self.max_seq_len,
            shuffle: true,
            seed: 42,
            pad_token_id: tokenizer.pad_token_id().unwrap_or(0),
            drop_last: false,
            ..Default::default()
        };

        let training_loop_config = TrainingLoopConfig {
            training: training_config,
            dataloader: dataloader_config,
            use_metal_flash_attention: self.flash_attention,
            log_every: 10,
            checkpoint_every: 500,
            eval_every: if eval_dataset.is_some() { 100 } else { 0 },
            use_jit_compilation: false,
            use_sequence_packing: self.sequence_packing,
            gradient_checkpointing: self.gradient_checkpointing,
            gradient_checkpointing_layers: self.gradient_checkpointing_layers,
            embedding_lr: self.embedding_lr,
            eager_evaluation: false,
            use_metal_fused_optimizer: self.metal_fused_optimizer,
            loraplus_lr_ratio: None,
            neftune_noise_alpha: None,
            use_cut_cross_entropy: false,
            #[cfg(feature = "distributed")]
            distributed: None,
        };

        let mut training_loop = TrainingLoop::new(training_loop_config);
        let use_sequence_packing = self.sequence_packing;

        if let Some(metrics_path) = &self.metrics_path {
            let callback = MetricsJsonCallback::new(metrics_path)
                .map_err(training_err)?
                .with_run_name(self.model_id.replace('/', "-"));
            training_loop.add_callback(Box::new(callback));
        }

        for callback in self.callbacks {
            training_loop.add_callback(callback);
        }

        // Run training (sequence packing by default)
        let model = if use_sequence_packing {
            training_loop
                .run_packed(
                    model,
                    train_dataset,
                    eval_dataset,
                    Some(&checkpoint_manager),
                )
                .map_err(training_err)?
        } else {
            let mut model = model;
            training_loop
                .run(
                    &mut model,
                    train_dataset,
                    eval_dataset,
                    Some(&checkpoint_manager),
                )
                .map_err(training_err)?;
            model
        };

        // Save final LoRA weights
        let lora_weights_path = output_dir.join("lora_weights.safetensors");
        model
            .save_lora_weights(&lora_weights_path)
            .map_err(training_err)?;

        Ok(FinetuneResult {
            final_loss: training_loop.current_loss(),
            total_steps: training_loop.current_step(),
            total_tokens: training_loop.total_tokens(),
            output_dir,
            lora_weights_path,
        })
    }
}

/// Builder for preference optimization methods such as DPO, SimPO, ORPO, and KTO.
pub struct PreferenceTuneBuilder {
    model_id: String,
    dataset_path: String,
    reference_model_id: Option<String>,
    method: PreferenceMethod,
    lora_r: usize,
    lora_alpha: f32,
    lora_dropout: f32,
    use_rslora: bool,
    use_dora: bool,
    learning_rate: f64,
    batch_size: usize,
    num_epochs: usize,
    output_dir: String,
    gradient_checkpointing: bool,
    gradient_checkpointing_layers: usize,
    callbacks: Vec<Box<dyn TrainingCallback>>,
    metrics_path: Option<PathBuf>,
}

impl PreferenceTuneBuilder {
    fn new(model_id: String, dataset_path: String, method: PreferenceMethod) -> Self {
        Self {
            model_id,
            dataset_path,
            reference_model_id: None,
            method,
            lora_r: 16,
            lora_alpha: 32.0,
            lora_dropout: 0.0,
            use_rslora: false,
            use_dora: false,
            learning_rate: 5e-6,
            batch_size: 1,
            num_epochs: 1,
            output_dir: "./output".to_string(),
            gradient_checkpointing: false,
            gradient_checkpointing_layers: 4,
            callbacks: Vec::new(),
            metrics_path: None,
        }
    }

    /// Set the reference model to use for DPO/KTO or SimPO+CPO.
    pub fn reference_model(mut self, model: impl Into<String>) -> Self {
        self.reference_model_id = Some(model.into());
        self
    }

    /// Set LoRA rank and alpha.
    pub fn lora(mut self, r: usize, alpha: f32) -> Self {
        self.lora_r = r;
        self.lora_alpha = alpha;
        self
    }

    /// Set LoRA dropout.
    pub fn lora_dropout(mut self, dropout: f32) -> Self {
        self.lora_dropout = dropout;
        self
    }

    /// Enable rank-stabilized LoRA.
    pub fn use_rslora(mut self, enabled: bool) -> Self {
        self.use_rslora = enabled;
        self
    }

    /// Enable DoRA adapters.
    pub fn use_dora(mut self, enabled: bool) -> Self {
        self.use_dora = enabled;
        self
    }

    /// Set learning rate.
    pub fn learning_rate(mut self, lr: f64) -> Self {
        self.learning_rate = lr;
        self
    }

    /// Set batch size.
    pub fn batch_size(mut self, n: usize) -> Self {
        self.batch_size = n;
        self
    }

    /// Set number of epochs.
    pub fn epochs(mut self, n: usize) -> Self {
        self.num_epochs = n;
        self
    }

    /// Set output directory.
    pub fn output(mut self, path: impl Into<String>) -> Self {
        self.output_dir = path.into();
        self
    }

    /// Enable or disable gradient checkpointing.
    pub fn gradient_checkpointing(mut self, enabled: bool) -> Self {
        self.gradient_checkpointing = enabled;
        self
    }

    /// Set layers per gradient checkpoint block.
    pub fn gradient_checkpointing_layers(mut self, layers: usize) -> Self {
        self.gradient_checkpointing_layers = layers;
        self
    }

    /// Add a training callback.
    pub fn callback(mut self, callback: Box<dyn TrainingCallback>) -> Self {
        self.callbacks.push(callback);
        self
    }

    /// Write training metrics to JSONL at the given path.
    pub fn metrics_path(mut self, path: impl Into<PathBuf>) -> Self {
        self.metrics_path = Some(path.into());
        self
    }

    /// Set DPO beta.
    pub fn dpo_beta(mut self, beta: f64) -> Self {
        if let PreferenceMethod::Dpo(config) = &mut self.method {
            config.beta = beta;
        }
        self
    }

    /// Set DPO loss type.
    pub fn dpo_loss_type(mut self, loss_type: pmetal_trainer::DpoLossType) -> Self {
        if let PreferenceMethod::Dpo(config) = &mut self.method {
            config.loss_type = loss_type;
        }
        self
    }

    /// Enable DPO reference-free mode.
    pub fn dpo_reference_free(mut self, enabled: bool) -> Self {
        if let PreferenceMethod::Dpo(config) = &mut self.method {
            config.reference_free = enabled;
        }
        self
    }

    /// Set SimPO beta.
    pub fn simpo_beta(mut self, beta: f64) -> Self {
        if let PreferenceMethod::Simpo(config) = &mut self.method {
            config.beta = beta;
        }
        self
    }

    /// Set SimPO gamma.
    pub fn simpo_gamma(mut self, gamma: f64) -> Self {
        if let PreferenceMethod::Simpo(config) = &mut self.method {
            config.gamma = gamma;
        }
        self
    }

    /// Enable SimPO CPO regularization.
    pub fn simpo_cpo(mut self, alpha: f64) -> Self {
        if let PreferenceMethod::Simpo(config) = &mut self.method {
            config.cpo_alpha = alpha;
        }
        self
    }

    /// Set ORPO beta.
    pub fn orpo_beta(mut self, beta: f64) -> Self {
        if let PreferenceMethod::Orpo(config) = &mut self.method {
            config.beta = beta;
        }
        self
    }

    /// Set KTO beta.
    pub fn kto_beta(mut self, beta: f64) -> Self {
        if let PreferenceMethod::Kto(config) = &mut self.method {
            config.beta = beta;
        }
        self
    }

    /// Set KTO desirable weight.
    pub fn kto_desirable_weight(mut self, weight: f64) -> Self {
        if let PreferenceMethod::Kto(config) = &mut self.method {
            config.desirable_weight = weight;
        }
        self
    }

    /// Set KTO undesirable weight.
    pub fn kto_undesirable_weight(mut self, weight: f64) -> Self {
        if let PreferenceMethod::Kto(config) = &mut self.method {
            config.undesirable_weight = weight;
        }
        self
    }

    /// Enable KTO reference-free mode.
    pub fn kto_reference_free(mut self, enabled: bool) -> Self {
        if let PreferenceMethod::Kto(config) = &mut self.method {
            config.reference_free = enabled;
        }
        self
    }

    /// Run the preference optimization pipeline.
    pub async fn run(self) -> Result<FinetuneResult> {
        let Self {
            model_id,
            dataset_path,
            reference_model_id,
            method,
            lora_r,
            lora_alpha,
            lora_dropout,
            use_rslora,
            use_dora,
            learning_rate,
            batch_size,
            num_epochs,
            output_dir,
            gradient_checkpointing,
            gradient_checkpointing_layers,
            callbacks,
            metrics_path,
        } = self;

        let model_path = resolve_model_path(&model_id).await?;
        let tokenizer = Tokenizer::from_model_dir(&model_path).map_err(|e| {
            PMetalError::ModelLoad(format!("Failed to load tokenizer from {model_path:?}: {e}"))
        })?;

        let lora_config = LoraConfig {
            r: lora_r,
            alpha: lora_alpha,
            dropout: lora_dropout,
            use_rslora,
            use_dora,
            ..Default::default()
        };
        let mut model = DynamicLoraModel::from_pretrained(&model_path, lora_config.clone())
            .map_err(model_err)?;
        if gradient_checkpointing && model.supports_gradient_checkpointing() {
            model.enable_gradient_checkpointing(gradient_checkpointing_layers);
        }

        let output_dir = PathBuf::from(&output_dir);
        std::fs::create_dir_all(&output_dir).map_err(training_err)?;
        let metrics_path = metrics_path.unwrap_or_else(|| output_dir.join("metrics.jsonl"));
        let mut extra_callbacks = Some(callbacks);

        let (final_loss, total_steps, total_tokens) = match method {
            PreferenceMethod::Dpo(config) => {
                let dataset = load_dpo_dataset(
                    &dataset_path,
                    &tokenizer,
                    config.max_prompt_length,
                    config.max_completion_length,
                    config.truncate_prompt_left,
                )?;
                let total_steps = dataset.len().div_ceil(batch_size.max(1)) * num_epochs.max(1);
                let total_tokens = dataset
                    .iter()
                    .map(|pair| pair.chosen_ids.len() + pair.rejected_ids.len())
                    .sum();
                let reference_path = reference_model_id
                    .clone()
                    .unwrap_or_else(|| model_id.clone());
                let mut reference_model =
                    if config.reference_free || config.use_stop_gradient_reference {
                        None
                    } else {
                        Some(
                            DynamicLoraModel::from_pretrained(
                                &resolve_model_path(&reference_path).await?,
                                LoraConfig {
                                    r: 0,
                                    ..Default::default()
                                },
                            )
                            .map_err(model_err)?,
                        )
                    };
                let mut trainer = pmetal_trainer::DpoTrainer::new(
                    config,
                    TrainingConfig {
                        learning_rate,
                        batch_size,
                        num_epochs,
                        output_dir: output_dir.display().to_string(),
                        ..Default::default()
                    },
                )
                .map_err(training_err)?;
                let metrics_cb = MetricsJsonCallback::new(&metrics_path)
                    .map_err(training_err)?
                    .with_run_name(format!("dpo-{}", model_id.replace('/', "-")));
                trainer.add_callback(Box::new(metrics_cb));
                if let Some(callbacks) = extra_callbacks.take() {
                    for callback in callbacks {
                        trainer.add_callback(callback);
                    }
                }
                let mut optimizer = mlx_rs::optimizers::AdamWBuilder::new(learning_rate as f32)
                    .build()
                    .map_err(training_err)?;
                let history = trainer
                    .train(
                        &mut model,
                        reference_model.as_mut(),
                        &dataset,
                        &mut optimizer,
                    )
                    .map_err(training_err)?;
                (
                    history.last().map(|m| m.loss as f64).unwrap_or(0.0),
                    total_steps,
                    total_tokens,
                )
            }
            PreferenceMethod::Simpo(config) => {
                let dataset = load_simpo_dataset(
                    &dataset_path,
                    &tokenizer,
                    config.max_prompt_length,
                    config
                        .max_seq_length
                        .saturating_sub(config.max_prompt_length),
                )?;
                let total_steps = dataset.len().div_ceil(batch_size.max(1)) * num_epochs.max(1);
                let total_tokens = dataset
                    .iter()
                    .map(|pair| {
                        pair.prompt_ids.len() + pair.chosen_ids.len() + pair.rejected_ids.len()
                    })
                    .sum();
                let reference_path = reference_model_id
                    .clone()
                    .unwrap_or_else(|| model_id.clone());
                let mut reference_model = if config.cpo_alpha > 0.0 {
                    Some(
                        DynamicLoraModel::from_pretrained(
                            &resolve_model_path(&reference_path).await?,
                            LoraConfig {
                                r: 0,
                                ..Default::default()
                            },
                        )
                        .map_err(model_err)?,
                    )
                } else {
                    None
                };
                let mut trainer = pmetal_trainer::SimpoTrainer::new(
                    config,
                    TrainingConfig {
                        learning_rate,
                        batch_size,
                        num_epochs,
                        output_dir: output_dir.display().to_string(),
                        ..Default::default()
                    },
                )
                .map_err(training_err)?;
                let metrics_cb = MetricsJsonCallback::new(&metrics_path)
                    .map_err(training_err)?
                    .with_run_name(format!("simpo-{}", model_id.replace('/', "-")));
                trainer.add_callback(Box::new(metrics_cb));
                if let Some(callbacks) = extra_callbacks.take() {
                    for callback in callbacks {
                        trainer.add_callback(callback);
                    }
                }
                let mut optimizer = mlx_rs::optimizers::AdamWBuilder::new(learning_rate as f32)
                    .build()
                    .map_err(training_err)?;
                let history = trainer
                    .train(
                        &mut model,
                        reference_model.as_mut(),
                        &dataset,
                        &mut optimizer,
                    )
                    .map_err(training_err)?;
                (
                    history.last().map(|m| m.loss as f64).unwrap_or(0.0),
                    total_steps,
                    total_tokens,
                )
            }
            PreferenceMethod::Orpo(config) => {
                let dataset = load_dpo_dataset(
                    &dataset_path,
                    &tokenizer,
                    config.max_prompt_length,
                    config.max_completion_length,
                    config.truncate_prompt_left,
                )?;
                let total_steps = dataset.len().div_ceil(batch_size.max(1)) * num_epochs.max(1);
                let total_tokens = dataset
                    .iter()
                    .map(|pair| pair.chosen_ids.len() + pair.rejected_ids.len())
                    .sum();
                let mut trainer = pmetal_trainer::OrpoTrainer::new(
                    config,
                    TrainingConfig {
                        learning_rate,
                        batch_size,
                        num_epochs,
                        output_dir: output_dir.display().to_string(),
                        ..Default::default()
                    },
                )
                .map_err(training_err)?;
                let metrics_cb = MetricsJsonCallback::new(&metrics_path)
                    .map_err(training_err)?
                    .with_run_name(format!("orpo-{}", model_id.replace('/', "-")));
                trainer.add_callback(Box::new(metrics_cb));
                if let Some(callbacks) = extra_callbacks.take() {
                    for callback in callbacks {
                        trainer.add_callback(callback);
                    }
                }
                let mut optimizer = mlx_rs::optimizers::AdamWBuilder::new(learning_rate as f32)
                    .build()
                    .map_err(training_err)?;
                let history = trainer
                    .train(&mut model, &dataset, &mut optimizer)
                    .map_err(training_err)?;
                (
                    history.last().map(|m| m.loss as f64).unwrap_or(0.0),
                    total_steps,
                    total_tokens,
                )
            }
            PreferenceMethod::Kto(config) => {
                let dataset = load_kto_dataset(
                    &dataset_path,
                    &tokenizer,
                    config.max_prompt_length,
                    config.max_completion_length,
                    config.truncate_prompt_left,
                )?;
                let total_steps = dataset.len().div_ceil(batch_size.max(1)) * num_epochs.max(1);
                let total_tokens = dataset.iter().map(|sample| sample.response_ids.len()).sum();
                let reference_path = reference_model_id
                    .clone()
                    .unwrap_or_else(|| model_id.clone());
                let mut reference_model = if config.reference_free {
                    None
                } else {
                    Some(
                        DynamicLoraModel::from_pretrained(
                            &resolve_model_path(&reference_path).await?,
                            LoraConfig {
                                r: 0,
                                ..Default::default()
                            },
                        )
                        .map_err(model_err)?,
                    )
                };
                let mut trainer = pmetal_trainer::KtoTrainer::new(
                    config,
                    TrainingConfig {
                        learning_rate,
                        batch_size,
                        num_epochs,
                        output_dir: output_dir.display().to_string(),
                        ..Default::default()
                    },
                )
                .map_err(training_err)?;
                let metrics_cb = MetricsJsonCallback::new(&metrics_path)
                    .map_err(training_err)?
                    .with_run_name(format!("kto-{}", model_id.replace('/', "-")));
                trainer.add_callback(Box::new(metrics_cb));
                if let Some(callbacks) = extra_callbacks.take() {
                    for callback in callbacks {
                        trainer.add_callback(callback);
                    }
                }
                let mut optimizer = mlx_rs::optimizers::AdamWBuilder::new(learning_rate as f32)
                    .build()
                    .map_err(training_err)?;
                let history = trainer
                    .train(
                        &mut model,
                        reference_model.as_mut(),
                        &dataset,
                        &mut optimizer,
                    )
                    .map_err(training_err)?;
                (
                    history.last().map(|m| m.loss as f64).unwrap_or(0.0),
                    total_steps,
                    total_tokens,
                )
            }
        };

        let lora_weights_path = output_dir.join("lora_weights.safetensors");
        model
            .save_lora_weights(&lora_weights_path)
            .map_err(training_err)?;
        let adapter_config = serde_json::json!({
            "r": lora_config.r,
            "alpha": lora_config.alpha,
            "target_modules": lora_config.target_modules,
            "use_rslora": lora_config.use_rslora,
        });
        std::fs::write(
            output_dir.join("adapter_config.json"),
            serde_json::to_string_pretty(&adapter_config).map_err(training_err)?,
        )
        .map_err(training_err)?;

        Ok(FinetuneResult {
            final_loss,
            total_steps,
            total_tokens,
            output_dir,
            lora_weights_path,
        })
    }
}

/// Builder for running inference on a model.
///
/// Wraps the full inference pipeline: model download, tokenizer loading,
/// chat template detection, KV-cached generation, and decoding.
pub struct InferBuilder {
    model_id: String,
    lora_path: Option<String>,
    temperature: f32,
    top_k: usize,
    top_p: f32,
    min_p: f32,
    max_tokens: usize,
    repetition_penalty: f32,
    seed: Option<u64>,
    fp8: bool,
    #[cfg(feature = "ane")]
    device: Option<pmetal_core::Device>,
}

impl InferBuilder {
    fn new(model_id: String) -> Self {
        Self {
            model_id,
            lora_path: None,
            temperature: 0.7,
            top_k: 50,
            top_p: 0.9,
            min_p: 0.05,
            max_tokens: 256,
            repetition_penalty: 1.0,
            seed: None,
            fp8: false,
            #[cfg(feature = "ane")]
            device: Some(pmetal_core::Device::Ane),
        }
    }

    /// Set path to LoRA adapter weights.
    pub fn lora(mut self, path: impl Into<String>) -> Self {
        self.lora_path = Some(path.into());
        self
    }

    /// Set sampling temperature.
    pub fn temperature(mut self, t: f32) -> Self {
        self.temperature = t;
        self
    }

    /// Set top-k sampling.
    pub fn top_k(mut self, k: usize) -> Self {
        self.top_k = k;
        self
    }

    /// Set top-p (nucleus) sampling.
    pub fn top_p(mut self, p: f32) -> Self {
        self.top_p = p;
        self
    }

    /// Set min-p sampling threshold.
    pub fn min_p(mut self, p: f32) -> Self {
        self.min_p = p;
        self
    }

    /// Set max tokens to generate.
    pub fn max_tokens(mut self, n: usize) -> Self {
        self.max_tokens = n;
        self
    }

    /// Set repetition penalty.
    pub fn repetition_penalty(mut self, penalty: f32) -> Self {
        self.repetition_penalty = penalty;
        self
    }

    /// Enable FP8 quantization for inference.
    pub fn fp8(mut self, enabled: bool) -> Self {
        self.fp8 = enabled;
        self
    }

    /// Set random seed for reproducible generation.
    pub fn seed(mut self, s: u64) -> Self {
        self.seed = Some(s);
        self
    }

    /// Set the compute device (ANE inference requires `ane` feature).
    #[cfg(feature = "ane")]
    pub fn device(mut self, device: pmetal_core::Device) -> Self {
        self.device = Some(device);
        self
    }

    /// Use GPU (Metal) instead of ANE for inference.
    #[cfg(feature = "ane")]
    pub fn gpu(mut self) -> Self {
        self.device = Some(pmetal_core::Device::Gpu);
        self
    }

    /// Generate text from the given prompt.
    pub async fn generate(self, prompt: &str) -> Result<InferResult> {
        // Resolve model path
        let model_path = resolve_model_path(&self.model_id).await?;

        // Load tokenizer (reads special_tokens_map.json and tokenizer_config.json too)
        let tokenizer = Tokenizer::from_model_dir(&model_path).map_err(|e| {
            PMetalError::ModelLoad(format!("Failed to load tokenizer from {model_path:?}: {e}"))
        })?;

        // Branch: ANE inference
        #[cfg(feature = "ane")]
        if matches!(self.device, Some(pmetal_core::Device::Ane)) {
            return self.generate_ane(&model_path, &tokenizer, prompt);
        }

        // Branch: LoRA inference vs standard inference
        if let Some(lora_path) = &self.lora_path {
            return self.generate_with_lora(&model_path, lora_path, &tokenizer, prompt);
        }

        // Load model with auto-detected architecture
        let mut model = DynamicModel::load(&model_path).map_err(mlx_err)?;

        // Apply FP8 quantization if requested
        if self.fp8 {
            model.quantize_fp8().map_err(mlx_err)?;
        }

        // Tokenize
        let input_ids = tokenizer.encode(prompt)?;

        // Configure generation
        let gen_config = self.build_gen_config(&tokenizer);

        // Create KV cache
        let max_seq_len = input_ids.len() + self.max_tokens + 64;
        let mut cache = model.create_cache(max_seq_len);

        // Generate
        let start = std::time::Instant::now();
        let output = generate_cached_async(
            |input, cache| model.forward_with_hybrid_cache(input, None, Some(cache), None),
            &input_ids,
            gen_config,
            &mut cache,
        )
        .map_err(mlx_err)?;
        let elapsed = start.elapsed();

        // Decode generated tokens
        let prompt_len = input_ids.len();
        let generated_tokens = &output.token_ids[prompt_len..];
        let text = tokenizer.decode(generated_tokens)?;

        Ok(InferResult {
            text,
            tokens_generated: output.num_generated,
            tokens_per_sec: output.num_generated as f64 / elapsed.as_secs_f64(),
        })
    }

    /// Generate text while streaming decoded deltas to the callback.
    ///
    /// Return `false` from the callback to stop generation early.
    pub async fn generate_streaming<F>(self, prompt: &str, mut on_delta: F) -> Result<InferResult>
    where
        F: FnMut(&str) -> bool,
    {
        // ANE currently uses the non-streaming engine. Emit the final result once.
        #[cfg(feature = "ane")]
        if matches!(self.device, Some(pmetal_core::Device::Ane)) {
            let result = self.generate(prompt).await?;
            if !result.text.is_empty() {
                let _ = on_delta(&result.text);
            }
            return Ok(result);
        }

        let model_path = resolve_model_path(&self.model_id).await?;
        let tokenizer = Tokenizer::from_model_dir(&model_path).map_err(|e| {
            PMetalError::ModelLoad(format!("Failed to load tokenizer from {model_path:?}: {e}"))
        })?;

        if let Some(lora_path) = &self.lora_path {
            return self.generate_with_lora_streaming(
                &model_path,
                lora_path,
                &tokenizer,
                prompt,
                &mut on_delta,
            );
        }

        let mut model = DynamicModel::load(&model_path).map_err(mlx_err)?;

        if self.fp8 {
            model.quantize_fp8().map_err(mlx_err)?;
        }

        let input_ids = tokenizer.encode(prompt)?;
        let gen_config = self.build_gen_config(&tokenizer);
        let max_seq_len = input_ids.len() + self.max_tokens + 64;
        let mut cache = model.create_cache(max_seq_len);

        let start = std::time::Instant::now();
        let mut token_buf: Vec<u32> = Vec::new();
        let mut streamed_text = String::new();

        let output = generate_cached_async_streaming(
            |input, cache| model.forward_with_hybrid_cache(input, None, Some(cache), None),
            &input_ids,
            gen_config,
            &mut cache,
            |token_id| {
                token_buf.push(token_id);
                if let Ok(text) = tokenizer.decode(&token_buf) {
                    if text.len() > streamed_text.len() {
                        let delta = &text[streamed_text.len()..];
                        if !delta.is_empty() && !on_delta(delta) {
                            return false;
                        }
                    }
                    streamed_text = text;
                }
                true
            },
        )
        .map_err(mlx_err)?;
        let elapsed = start.elapsed();

        Ok(InferResult {
            text: streamed_text,
            tokens_generated: output.num_generated,
            tokens_per_sec: output.num_generated as f64 / elapsed.as_secs_f64(),
        })
    }

    /// Inference with LoRA adapter.
    fn generate_with_lora(
        &self,
        model_path: &Path,
        lora_path: &str,
        tokenizer: &Tokenizer,
        prompt: &str,
    ) -> Result<InferResult> {
        let lora_config = LoraConfig::default();
        let mut model =
            DynamicLoraModel::from_pretrained(model_path, lora_config).map_err(model_err)?;
        model.load_lora_weights(lora_path).map_err(model_err)?;

        let input_ids = tokenizer.encode(prompt)?;
        let gen_config = self.build_gen_config(tokenizer);

        let max_seq_len = input_ids.len() + self.max_tokens + 64;
        let mut cache = model.create_cache(max_seq_len).ok_or_else(|| {
            PMetalError::NotImplemented("Model does not support KV cache".to_string())
        })?;

        let start = std::time::Instant::now();
        let output = generate_cached_async(
            |input, cache| {
                model
                    .forward_with_cache(input, None, Some(cache))
                    .map_err(|e| mlx_rs::error::Exception::custom(e.to_string()))
            },
            &input_ids,
            gen_config,
            &mut cache,
        )
        .map_err(mlx_err)?;
        let elapsed = start.elapsed();

        let prompt_len = input_ids.len();
        let generated_tokens = &output.token_ids[prompt_len..];
        let text = tokenizer.decode(generated_tokens)?;

        Ok(InferResult {
            text,
            tokens_generated: output.num_generated,
            tokens_per_sec: output.num_generated as f64 / elapsed.as_secs_f64(),
        })
    }

    fn generate_with_lora_streaming<F>(
        &self,
        model_path: &Path,
        lora_path: &str,
        tokenizer: &Tokenizer,
        prompt: &str,
        on_delta: &mut F,
    ) -> Result<InferResult>
    where
        F: FnMut(&str) -> bool,
    {
        let lora_config = LoraConfig::default();
        let mut model =
            DynamicLoraModel::from_pretrained(model_path, lora_config).map_err(model_err)?;
        model.load_lora_weights(lora_path).map_err(model_err)?;

        let input_ids = tokenizer.encode(prompt)?;
        let gen_config = self.build_gen_config(tokenizer);

        let max_seq_len = input_ids.len() + self.max_tokens + 64;
        let mut cache = model.create_cache(max_seq_len).ok_or_else(|| {
            PMetalError::NotImplemented("Model does not support KV cache".to_string())
        })?;

        let start = std::time::Instant::now();
        let mut token_buf: Vec<u32> = Vec::new();
        let mut streamed_text = String::new();

        let output = generate_cached_async_streaming(
            |input, cache| {
                model
                    .forward_with_cache(input, None, Some(cache))
                    .map_err(|e| mlx_rs::error::Exception::custom(e.to_string()))
            },
            &input_ids,
            gen_config,
            &mut cache,
            |token_id| {
                token_buf.push(token_id);
                if let Ok(text) = tokenizer.decode(&token_buf) {
                    if text.len() > streamed_text.len() {
                        let delta = &text[streamed_text.len()..];
                        if !delta.is_empty() && !on_delta(delta) {
                            return false;
                        }
                    }
                    streamed_text = text;
                }
                true
            },
        )
        .map_err(mlx_err)?;
        let elapsed = start.elapsed();

        Ok(InferResult {
            text: streamed_text,
            tokens_generated: output.num_generated,
            tokens_per_sec: output.num_generated as f64 / elapsed.as_secs_f64(),
        })
    }

    /// Inference via ANE with SafeTensors/flat weight loading, LoRA fusion, KV cache.
    #[cfg(feature = "ane")]
    fn generate_ane(
        &self,
        model_path: &Path,
        tokenizer: &Tokenizer,
        prompt: &str,
    ) -> Result<InferResult> {
        use pmetal_metal::ane::inference::{AneInferenceConfig, AneInferenceEngine};

        // Read model config
        let config_path = model_path.join("config.json");
        let config_text = std::fs::read_to_string(&config_path).map_err(|e| {
            PMetalError::ModelLoad(format!(
                "Failed to read config.json at {config_path:?}: {e}"
            ))
        })?;

        fn extract_usize(json: &str, key: &str) -> std::result::Result<usize, PMetalError> {
            let needle = format!("\"{key}\"");
            let pos = json
                .find(&needle)
                .ok_or_else(|| PMetalError::ModelLoad(format!("config.json missing '{key}'")))?;
            let after_key = &json[pos + needle.len()..];
            let num_start = after_key
                .find(|c: char| c.is_ascii_digit())
                .ok_or_else(|| {
                    PMetalError::ModelLoad(format!("config.json: no value for '{key}'"))
                })?;
            let num_str = &after_key[num_start..];
            let num_end = num_str
                .find(|c: char| !c.is_ascii_digit())
                .unwrap_or(num_str.len());
            num_str[..num_end].parse::<usize>().map_err(|_| {
                PMetalError::ModelLoad(format!("config.json: invalid value for '{key}'"))
            })
        }

        fn extract_usize_optional(json: &str, key: &str) -> Option<usize> {
            extract_usize(json, key).ok()
        }

        let dim = extract_usize(&config_text, "hidden_size")?;
        let hidden_dim = extract_usize(&config_text, "intermediate_size")?;
        let n_heads = extract_usize(&config_text, "num_attention_heads")?;
        let n_layers = extract_usize(&config_text, "num_hidden_layers")?;
        let vocab_size = extract_usize(&config_text, "vocab_size")?;
        // GQA: read n_kv_heads, default to n_heads
        let n_kv_heads =
            extract_usize_optional(&config_text, "num_key_value_heads").unwrap_or(n_heads);

        // Parse rope_theta and rms_norm_eps from config.json
        fn extract_float(json: &str, key: &str) -> Option<f32> {
            let needle = format!("\"{key}\"");
            let pos = json.find(&needle)?;
            let after_key = &json[pos + needle.len()..];
            // Skip to colon, then whitespace
            let after_colon = after_key.find(':').map(|i| &after_key[i + 1..])?;
            let trimmed = after_colon.trim_start();
            // Extract the numeric value (may contain digits, '.', 'e', 'E', '+', '-')
            let end = trimmed
                .find(|c: char| {
                    !c.is_ascii_digit() && c != '.' && c != 'e' && c != 'E' && c != '+' && c != '-'
                })
                .unwrap_or(trimmed.len());
            trimmed[..end].parse::<f64>().ok().map(|v| v as f32)
        }
        let rope_theta = extract_float(&config_text, "rope_theta").unwrap_or(1_000_000.0);
        let rms_norm_eps = extract_float(&config_text, "rms_norm_eps").unwrap_or(1e-6);
        let head_dim = extract_usize_optional(&config_text, "head_dim");

        // Tokenize early so we know prompt length for ANE bucketing
        let input_ids_u32 = tokenizer.encode(prompt)?;
        let prompt_len = input_ids_u32.len();

        let ane_config = AneInferenceConfig {
            dim,
            hidden_dim,
            n_heads,
            n_kv_heads,
            n_layers,
            vocab_size,
            max_seq_len: prompt_len + self.max_tokens + 64,
            ane_seq_len: None, // auto-bucket from prompt length
            temperature: self.temperature,
            top_k: self.top_k,
            max_tokens: self.max_tokens,
            eos_token_id: tokenizer.eos_token_id(),
            rope_theta,
            rms_norm_eps,
            head_dim,
            ..Default::default()
        };

        let mut engine = AneInferenceEngine::new(ane_config, prompt_len)
            .map_err(|e| PMetalError::ModelLoad(format!("ANE engine init failed: {e}")))?;

        // Try SafeTensors first, fall back to model.bin
        let safetensors_single = model_path.join("model.safetensors");
        let safetensors_multi = model_path.join("model-00001-of-00002.safetensors");
        let safetensors_index = model_path.join("model.safetensors.index.json");
        let weights_bin = model_path.join("model.bin");

        if safetensors_single.exists() {
            engine
                .load_weights_safetensors(&safetensors_single)
                .map_err(|e| PMetalError::ModelLoad(format!("SafeTensors load failed: {e}")))?;
        } else if safetensors_index.exists() || safetensors_multi.exists() {
            engine
                .load_weights_safetensors(model_path)
                .map_err(|e| PMetalError::ModelLoad(format!("SafeTensors load failed: {e}")))?;
        } else if weights_bin.exists() {
            let weight_data = std::fs::read(&weights_bin)
                .map_err(|e| PMetalError::ModelLoad(format!("Failed to read model.bin: {e}")))?;
            if weight_data.len() % 4 != 0 {
                return Err(PMetalError::ModelLoad(
                    "model.bin size must be a multiple of 4 bytes".into(),
                ));
            }
            #[allow(unsafe_code)]
            let (prefix, weights, suffix) = unsafe { weight_data.align_to::<f32>() };
            if !prefix.is_empty() || !suffix.is_empty() {
                return Err(PMetalError::ModelLoad(
                    "model.bin data is not properly aligned for f32".into(),
                ));
            }
            engine.load_weights_flat(weights);
        } else {
            return Err(PMetalError::ModelLoad(format!(
                "No weight files found in {:?}. Expected model.safetensors or model.bin.",
                model_path
            )));
        }

        // Check for LoRA adapter and merge if present
        let adapter_config = model_path.join("adapter_config.json");
        let adapter_weights = model_path.join("adapter_model.safetensors");
        if adapter_config.exists() && adapter_weights.exists() {
            engine
                .load_lora_adapter(model_path)
                .map_err(|e| PMetalError::ModelLoad(format!("LoRA merge failed: {e}")))?;
        }

        // Also check user-specified LoRA path
        if let Some(ref lora_path) = self.lora_path {
            let lora_dir = Path::new(lora_path);
            engine
                .load_lora_adapter(lora_dir)
                .map_err(|e| PMetalError::ModelLoad(format!("LoRA merge failed: {e}")))?;
        }

        // Compile kernels (after LoRA merge)
        engine
            .compile_kernels()
            .map_err(|e| PMetalError::Training(format!("ANE compile failed: {e}")))?;

        // Generate with KV cache (input_ids_u32 already tokenized above)
        let start = std::time::Instant::now();
        let output_ids = engine
            .generate_cached(&input_ids_u32)
            .map_err(|e| PMetalError::Training(format!("ANE generate failed: {e}")))?;
        let elapsed = start.elapsed();
        let generated: Vec<u32> = output_ids[prompt_len..].to_vec();
        let text = tokenizer.decode(&generated)?;

        Ok(InferResult {
            text,
            tokens_generated: generated.len(),
            tokens_per_sec: generated.len() as f64 / elapsed.as_secs_f64(),
        })
    }

    /// Build a `GenerationConfig` from builder settings.
    fn build_gen_config(&self, tokenizer: &Tokenizer) -> GenerationConfig {
        let mut config = if self.temperature < 1e-6 {
            GenerationConfig::greedy(self.max_tokens)
        } else {
            GenerationConfig::sampling(self.max_tokens, self.temperature)
                .with_top_k(self.top_k)
                .with_top_p(self.top_p)
                .with_min_p(self.min_p)
                .with_repetition_penalty(self.repetition_penalty)
        };

        // Use actual EOS token, don't hardcode a fallback
        if let Some(eos) = tokenizer.eos_token_id() {
            config = config.with_stop_tokens(vec![eos]);
        }

        if let Some(seed) = self.seed {
            config = config.with_seed(seed);
        }

        config
    }
}

fn read_jsonl_objects(path: &str) -> Result<Vec<serde_json::Value>> {
    use std::io::BufRead;

    let file = std::fs::File::open(path).map_err(training_err)?;
    let reader = std::io::BufReader::new(file);
    let mut rows = Vec::new();
    for line in reader.lines() {
        let line = line.map_err(training_err)?;
        if line.trim().is_empty() {
            continue;
        }
        rows.push(serde_json::from_str(&line).map_err(training_err)?);
    }
    Ok(rows)
}

fn extract_first_string<'a>(value: &'a serde_json::Value, fields: &[&str]) -> Option<&'a str> {
    fields
        .iter()
        .find_map(|field| value.get(*field).and_then(serde_json::Value::as_str))
}

fn encode_segment(
    tokenizer: &Tokenizer,
    text: &str,
    max_len: usize,
    truncate_left: bool,
    append_eos: bool,
) -> Result<Vec<u32>> {
    let mut ids = tokenizer.encode(text)?;
    if append_eos {
        if let Some(eos) = tokenizer.eos_token_id() {
            if ids.last().copied() != Some(eos) {
                ids.push(eos);
            }
        }
    }
    if ids.len() > max_len {
        ids = if truncate_left {
            ids[ids.len() - max_len..].to_vec()
        } else {
            ids[..max_len].to_vec()
        };
    }
    Ok(ids)
}

fn load_dpo_dataset(
    path: &str,
    tokenizer: &Tokenizer,
    max_prompt_length: usize,
    max_completion_length: usize,
    truncate_prompt_left: bool,
) -> Result<Vec<pmetal_trainer::dpo::PreferencePair>> {
    let rows = read_jsonl_objects(path)?;
    let mut pairs = Vec::with_capacity(rows.len());
    for row in rows {
        let prompt = extract_first_string(&row, &["prompt", "instruction", "input"])
            .ok_or_else(|| PMetalError::Training("Preference row missing prompt".into()))?;
        let chosen = extract_first_string(
            &row,
            &[
                "chosen",
                "accepted",
                "preferred",
                "chosen_response",
                "output_chosen",
            ],
        )
        .ok_or_else(|| PMetalError::Training("Preference row missing chosen response".into()))?;
        let rejected = extract_first_string(
            &row,
            &[
                "rejected",
                "rejected_response",
                "dispreferred",
                "output_rejected",
            ],
        )
        .ok_or_else(|| PMetalError::Training("Preference row missing rejected response".into()))?;

        let prompt_ids = encode_segment(
            tokenizer,
            prompt,
            max_prompt_length,
            truncate_prompt_left,
            false,
        )?;
        let chosen_ids = encode_segment(tokenizer, chosen, max_completion_length, false, true)?;
        let rejected_ids = encode_segment(tokenizer, rejected, max_completion_length, false, true)?;
        pairs.push(pmetal_trainer::dpo::PreferencePair::new(
            prompt_ids,
            chosen_ids,
            rejected_ids,
        ));
    }
    Ok(pairs)
}

fn load_simpo_dataset(
    path: &str,
    tokenizer: &Tokenizer,
    max_prompt_length: usize,
    max_completion_length: usize,
) -> Result<Vec<pmetal_trainer::simpo::PreferencePair>> {
    let rows = read_jsonl_objects(path)?;
    let mut pairs = Vec::with_capacity(rows.len());
    for row in rows {
        let prompt = extract_first_string(&row, &["prompt", "instruction", "input"])
            .ok_or_else(|| PMetalError::Training("Preference row missing prompt".into()))?;
        let chosen = extract_first_string(
            &row,
            &[
                "chosen",
                "accepted",
                "preferred",
                "chosen_response",
                "output_chosen",
            ],
        )
        .ok_or_else(|| PMetalError::Training("Preference row missing chosen response".into()))?;
        let rejected = extract_first_string(
            &row,
            &[
                "rejected",
                "rejected_response",
                "dispreferred",
                "output_rejected",
            ],
        )
        .ok_or_else(|| PMetalError::Training("Preference row missing rejected response".into()))?;

        let prompt_ids = encode_segment(tokenizer, prompt, max_prompt_length, true, false)?;
        let chosen_ids = encode_segment(tokenizer, chosen, max_completion_length, false, true)?;
        let rejected_ids = encode_segment(tokenizer, rejected, max_completion_length, false, true)?;
        pairs.push(pmetal_trainer::simpo::PreferencePair::new(
            prompt_ids,
            chosen_ids,
            rejected_ids,
        ));
    }
    Ok(pairs)
}

fn load_kto_dataset(
    path: &str,
    tokenizer: &Tokenizer,
    max_prompt_length: usize,
    max_completion_length: usize,
    truncate_prompt_left: bool,
) -> Result<Vec<pmetal_trainer::KtoSample>> {
    let rows = read_jsonl_objects(path)?;
    let mut samples = Vec::with_capacity(rows.len());
    for row in rows {
        let prompt = extract_first_string(&row, &["prompt", "instruction", "input"])
            .ok_or_else(|| PMetalError::Training("KTO row missing prompt".into()))?;
        let response = extract_first_string(&row, &["completion", "response", "output", "answer"])
            .ok_or_else(|| PMetalError::Training("KTO row missing response".into()))?;
        let label = row
            .get("label")
            .or_else(|| row.get("rating"))
            .or_else(|| row.get("chosen"))
            .ok_or_else(|| PMetalError::Training("KTO row missing label".into()))?;

        let is_desirable = match label {
            serde_json::Value::Bool(value) => *value,
            serde_json::Value::Number(value) => value.as_i64().unwrap_or_default() > 0,
            serde_json::Value::String(value) => matches!(
                value.to_ascii_lowercase().as_str(),
                "desirable" | "good" | "preferred" | "true" | "yes" | "1"
            ),
            _ => false,
        };

        let prompt_ids = encode_segment(
            tokenizer,
            prompt,
            max_prompt_length,
            truncate_prompt_left,
            false,
        )?;
        let response_ids = encode_segment(tokenizer, response, max_completion_length, false, true)?;
        samples.push(pmetal_trainer::KtoSample::new(
            prompt_ids,
            response_ids,
            is_desirable,
        ));
    }
    Ok(samples)
}

/// Resolve a model ID to a local path, downloading from HuggingFace Hub if necessary.
async fn resolve_model_path(model_id: &str) -> Result<PathBuf> {
    if is_hf_model_id(model_id) {
        // HuggingFace model ID — download_model fetches all repo files
        let path = pmetal_hub::download_model(model_id, None, None).await?;
        Ok(path)
    } else {
        Ok(PathBuf::from(model_id))
    }
}
