iron-shell 14.0.3

use crate::ai_api::models::*;
use anyhow::Result;
use chrono::{DateTime, Utc};
use ironcrypto::core_types::traits::Digest;
use ironcrypto::hash::Sha256;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};

/// The SHA-256 of a file, read in 1 MiB chunks so its size does not matter.
pub fn sha256_file(path: &Path) -> Result<String> {
    use std::io::Read;
    let mut file = fs::File::open(path)?;
    let mut hasher = Sha256::new();
    let mut buf = vec![0u8; 1 << 20];
    loop {
        let n = file.read(&mut buf)?;
        if n == 0 {
            break;
        }
        hasher.update(&buf[..n]);
    }
    Ok(hex::encode(hasher.finalize()?))
}

/// The longest model id accepted. Hugging Face repository names are far shorter.
pub const MODEL_ID_MAX: usize = 256;

/// Check that a model id can be used as a path, without escaping the model
/// directories. Ids come from requests, so they are not trusted.
///
/// Accepted: ASCII letters, digits, `-`, `_`, `.` and `/`, in non-empty
/// segments that do not start with `.`, with no `..` anywhere. That refuses
/// `\` (a path separator on Windows), `:` (a drive prefix or an alternate
/// data stream), NUL and other control characters, and a leading `/`.
pub fn validate_model_id(id: &str) -> Result<()> {
    if id.is_empty() || id.len() > MODEL_ID_MAX {
        anyhow::bail!("model id must be 1 to {MODEL_ID_MAX} bytes long");
    }
    if !id
        .bytes()
        .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b'/'))
    {
        anyhow::bail!(
            "model id {id:?} has a character outside letters, digits, '-', '_', '.' and '/'"
        );
    }
    if id.contains("..") {
        anyhow::bail!("model id {id:?} contains '..'");
    }
    if id
        .split('/')
        .any(|seg| seg.is_empty() || seg.starts_with('.'))
    {
        anyhow::bail!("model id {id:?} has an empty segment or one starting with '.'");
    }
    Ok(())
}

/// The file-name stem for a model id: the id with `/` replaced by `_`, after
/// [`validate_model_id`] has passed. Every path built from an id goes through
/// here, so no path can carry an id that was not checked.
pub fn model_file_stem(id: &str) -> Result<String> {
    validate_model_id(id)?;
    Ok(id.replace('/', "_"))
}

/// A model id for a model fetched from a direct URL: the URL's host and path,
/// with every character outside the accepted set replaced by `_`, and empty,
/// `.` and `..` segments dropped. The result always passes
/// [`validate_model_id`], or this returns an error.
pub fn url_model_id(url: &str) -> Result<String> {
    let parsed = url::Url::parse(url)?;
    let raw = format!("{}{}", parsed.host_str().unwrap_or(""), parsed.path());
    let mut segments = Vec::new();
    for seg in raw.split('/') {
        let clean: String = seg
            .chars()
            .map(|c| {
                if c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.') {
                    c
                } else {
                    '_'
                }
            })
            .collect::<String>()
            .replace("..", "_");
        let clean = clean.trim_start_matches('.');
        if !clean.is_empty() {
            segments.push(clean.to_string());
        }
    }
    let id = segments.join("/");
    validate_model_id(&id)?;
    Ok(id)
}

/// XDG Base Directory compliant storage for AI models
///
/// Directory structure:
/// ~/.local/share/ai-models/
/// ├── models/                 # Model files
/// │   ├── gguf/              # GGUF format models
/// │   ├── safetensors/       # SafeTensors format models
/// │   ├── pytorch/           # PyTorch format models
/// │   └── onnx/              # ONNX format models
/// ├── metadata/              # Model metadata files
/// ├── cache/                 # Temporary files and downloads
/// └── index.json             # Global model index
///
/// ~/.config/ai-models/
/// ├── config.toml            # Global configuration
/// ├── providers.toml         # Provider configurations
/// └── aliases.toml           # Model aliases and shortcuts
pub struct ModelStorage {
    data_dir: PathBuf,
    config_dir: PathBuf,
    cache_dir: PathBuf,
    index: ModelIndex,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelIndex {
    pub version: String,
    pub last_updated: DateTime<Utc>,
    pub models: HashMap<String, LocalModelMetadata>,
    pub aliases: HashMap<String, String>,
    pub total_size: u64,
}

impl ModelStorage {
    pub fn new(config: &StorageConfig) -> Result<Self> {
        let data_dir = if let Some(custom_path) = &config.custom_data_dir {
            custom_path.clone()
        } else {
            Self::get_xdg_data_dir()?
        };

        let config_dir = if let Some(custom_path) = &config.custom_config_dir {
            custom_path.clone()
        } else {
            Self::get_xdg_config_dir()?
        };

        let cache_dir = data_dir.join("cache");

        // Create necessary directories
        fs::create_dir_all(&data_dir)?;
        fs::create_dir_all(&config_dir)?;
        fs::create_dir_all(&cache_dir)?;
        fs::create_dir_all(data_dir.join("models"))?;
        fs::create_dir_all(data_dir.join("models/gguf"))?;
        fs::create_dir_all(data_dir.join("models/safetensors"))?;
        fs::create_dir_all(data_dir.join("models/pytorch"))?;
        fs::create_dir_all(data_dir.join("models/onnx"))?;
        fs::create_dir_all(data_dir.join("metadata"))?;

        let mut storage = Self {
            data_dir,
            config_dir,
            cache_dir,
            index: ModelIndex {
                version: "1.0".to_string(),
                last_updated: Utc::now(),
                models: HashMap::new(),
                aliases: HashMap::new(),
                total_size: 0,
            },
        };

        storage.load_index()?;
        Ok(storage)
    }

    /// Get XDG data directory for AI models
    fn get_xdg_data_dir() -> Result<PathBuf> {
        if let Ok(xdg_data_home) = std::env::var("XDG_DATA_HOME") {
            Ok(PathBuf::from(xdg_data_home).join("ai-models"))
        } else if let Ok(home) = std::env::var("HOME") {
            Ok(PathBuf::from(home).join(".local/share/ai-models"))
        } else {
            // Windows fallback
            if let Ok(appdata) = std::env::var("APPDATA") {
                Ok(PathBuf::from(appdata).join("ai-models"))
            } else {
                Err(anyhow::anyhow!("Cannot determine data directory"))
            }
        }
    }

    /// Get XDG config directory for AI models
    fn get_xdg_config_dir() -> Result<PathBuf> {
        if let Ok(xdg_config_home) = std::env::var("XDG_CONFIG_HOME") {
            Ok(PathBuf::from(xdg_config_home).join("ai-models"))
        } else if let Ok(home) = std::env::var("HOME") {
            Ok(PathBuf::from(home).join(".config/ai-models"))
        } else {
            // Windows fallback
            if let Ok(appdata) = std::env::var("APPDATA") {
                Ok(PathBuf::from(appdata).join("ai-models"))
            } else {
                Err(anyhow::anyhow!("Cannot determine config directory"))
            }
        }
    }

    /// Load the model index from disk
    fn load_index(&mut self) -> Result<()> {
        let index_path = self.data_dir.join("index.json");
        if index_path.exists() {
            let content = fs::read_to_string(&index_path)?;
            self.index = serde_json::from_str(&content)?;
        }
        Ok(())
    }

    /// Save the model index to disk
    fn save_index(&self) -> Result<()> {
        let index_path = self.data_dir.join("index.json");
        let content = serde_json::to_string_pretty(&self.index)?;
        fs::write(&index_path, content)?;
        Ok(())
    }

    /// List all locally stored models
    pub fn list_local_models(&self) -> Result<Vec<ModelInfo>> {
        let mut models = Vec::new();

        for (id, metadata) in &self.index.models {
            let model_info = ModelInfo {
                id: id.clone(),
                object: "model".to_string(),
                created: metadata.downloaded_at.timestamp(),
                owned_by: "local".to_string(),
                provider: "local".to_string(),
                context_length: metadata
                    .parameters
                    .get("context_length")
                    .and_then(|v| v.as_u64())
                    .map(|v| v as u32),
                max_output: metadata
                    .parameters
                    .get("max_output")
                    .and_then(|v| v.as_u64())
                    .map(|v| v as u32),
                per_request_limits: None,
                pricing: None, // Local models are free
                capabilities: metadata.capabilities.clone(),
                local_path: Some(metadata.file_path.clone()),
                format: metadata.format.clone(),
                size_bytes: Some(metadata.size_bytes),
                metadata: metadata.parameters.clone(),
            };
            models.push(model_info);
        }

        Ok(models)
    }

    /// Store a model file and metadata
    pub async fn store_model(
        &mut self,
        model_data: &[u8],
        metadata: LocalModelMetadata,
    ) -> Result<()> {
        // Determine storage path based on format
        let format_dir = match metadata.format {
            ModelFormat::GGUF => "gguf",
            ModelFormat::SafeTensors => "safetensors",
            ModelFormat::PyTorch => "pytorch",
            ModelFormat::ONNX => "onnx",
            _ => "other",
        };

        let model_dir = self.data_dir.join("models").join(format_dir);
        fs::create_dir_all(&model_dir)?;

        // Generate filename from model ID
        let filename = format!(
            "{}.{}",
            model_file_stem(&metadata.id)?,
            self.get_file_extension(&metadata.format)
        );
        let file_path = model_dir.join(&filename);

        // Write model file
        fs::write(&file_path, model_data)?;

        // Calculate and verify checksum
        let calculated_hash = hex::encode(Sha256::digest(model_data)?);
        self.finish_store(
            file_path,
            calculated_hash,
            model_data.len() as u64,
            metadata,
        )
    }

    /// [`Self::store_model`] for a model already on disk at `src`: it is copied
    /// and hashed in chunks, never held in memory. A downloaded model was read
    /// whole to be hashed and stored, and model files run to many gigabytes.
    pub async fn store_model_file(
        &mut self,
        src: &Path,
        metadata: LocalModelMetadata,
    ) -> Result<()> {
        let format_dir = match metadata.format {
            ModelFormat::GGUF => "gguf",
            ModelFormat::SafeTensors => "safetensors",
            ModelFormat::PyTorch => "pytorch",
            ModelFormat::ONNX => "onnx",
            _ => "other",
        };
        let model_dir = self.data_dir.join("models").join(format_dir);
        fs::create_dir_all(&model_dir)?;
        let filename = format!(
            "{}.{}",
            model_file_stem(&metadata.id)?,
            self.get_file_extension(&metadata.format)
        );
        let file_path = model_dir.join(&filename);
        let size = fs::copy(src, &file_path)?;
        let calculated_hash = sha256_file(&file_path)?;
        self.finish_store(file_path, calculated_hash, size, metadata)
    }

    /// Verify a stored model's checksum, then record its metadata and index it.
    fn finish_store(
        &mut self,
        file_path: PathBuf,
        calculated_hash: String,
        size_bytes: u64,
        metadata: LocalModelMetadata,
    ) -> Result<()> {
        if calculated_hash != metadata.sha256 {
            fs::remove_file(&file_path)?;
            return Err(anyhow::anyhow!(
                "Checksum mismatch: expected {}, got {}",
                metadata.sha256,
                calculated_hash
            ));
        }

        // Update metadata with actual file path
        let mut updated_metadata = metadata;
        updated_metadata.file_path = file_path.to_string_lossy().to_string();
        updated_metadata.size_bytes = size_bytes;

        // Store metadata file
        let metadata_path = self
            .data_dir
            .join("metadata")
            .join(format!("{}.json", model_file_stem(&updated_metadata.id)?));
        let metadata_content = serde_json::to_string_pretty(&updated_metadata)?;
        fs::write(&metadata_path, metadata_content)?;

        // Update index
        self.index
            .models
            .insert(updated_metadata.id.clone(), updated_metadata);
        self.index.last_updated = Utc::now();
        self.index.total_size = self.index.models.values().map(|m| m.size_bytes).sum();
        self.save_index()?;

        Ok(())
    }

    /// Remove a model from storage.
    ///
    /// Every file is attempted. The index entry is dropped only when all of
    /// them are gone, so a model that could not be deleted stays listed and
    /// the call can be retried.
    pub fn remove_model(&mut self, model_id: &str) -> Result<()> {
        let metadata = self
            .index
            .models
            .get(model_id)
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("Model {} not found", model_id))?;

        let metadata_path = self
            .data_dir
            .join("metadata")
            .join(format!("{}.json", model_file_stem(model_id)?));
        let mut paths = vec![PathBuf::from(&metadata.file_path), metadata_path];
        if let Some(config_path) = &metadata.config_path {
            paths.push(PathBuf::from(config_path));
        }
        if let Some(tokenizer_path) = &metadata.tokenizer_path {
            paths.push(PathBuf::from(tokenizer_path));
        }

        let mut failures = Vec::new();
        for path in &paths {
            match fs::remove_file(path) {
                Ok(()) => {}
                Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
                Err(e) => failures.push(format!("{}: {}", path.display(), e)),
            }
        }
        if !failures.is_empty() {
            return Err(anyhow::anyhow!(
                "Model {} was not removed; its entry stays in the index so it can be retried. Failed: {}",
                model_id,
                failures.join("; ")
            ));
        }

        self.index.models.remove(model_id);
        self.index.last_updated = Utc::now();
        self.index.total_size = self.index.models.values().map(|m| m.size_bytes).sum();
        self.save_index()?;

        Ok(())
    }

    /// Get model metadata
    pub fn get_model_metadata(&self, model_id: &str) -> Option<&LocalModelMetadata> {
        self.index.models.get(model_id)
    }

    /// Update model usage statistics
    pub fn update_model_usage(&mut self, model_id: &str) -> Result<()> {
        if let Some(metadata) = self.index.models.get_mut(model_id) {
            metadata.last_used = Some(Utc::now());
            metadata.usage_count += 1;
            self.save_index()?;
        }
        Ok(())
    }

    /// Get storage statistics
    pub fn get_storage_stats(&self) -> StorageStats {
        let model_count = self.index.models.len();
        let total_size = self.index.total_size;

        let mut format_breakdown = HashMap::new();
        for metadata in self.index.models.values() {
            *format_breakdown.entry(metadata.format.clone()).or_insert(0) += 1;
        }

        StorageStats {
            model_count,
            total_size,
            format_breakdown,
            cache_size: self.get_cache_size().unwrap_or(0),
            data_dir: self.data_dir.clone(),
            config_dir: self.config_dir.clone(),
        }
    }

    /// Clean up cache directory
    pub fn cleanup_cache(&self, max_age_days: u64) -> Result<u64> {
        let cache_dir = &self.cache_dir;
        let cutoff_time = Utc::now() - chrono::Duration::days(max_age_days as i64);
        let mut cleaned_size = 0u64;

        if cache_dir.exists() {
            for entry in fs::read_dir(cache_dir)? {
                let entry = entry?;
                let metadata = entry.metadata()?;

                if let Ok(modified) = metadata.modified() {
                    let modified_time: DateTime<Utc> = modified.into();
                    if modified_time < cutoff_time && metadata.is_file() {
                        cleaned_size += metadata.len();
                        fs::remove_file(entry.path())?;
                    }
                }
            }
        }

        Ok(cleaned_size)
    }

    /// Add model alias
    pub fn add_alias(&mut self, alias: String, model_id: String) -> Result<()> {
        self.index.aliases.insert(alias, model_id);
        self.save_index()
    }

    /// Remove model alias
    pub fn remove_alias(&mut self, alias: &str) -> Result<()> {
        self.index.aliases.remove(alias);
        self.save_index()
    }

    /// Resolve model ID from alias
    pub fn resolve_alias(&self, id_or_alias: &str) -> String {
        self.index
            .aliases
            .get(id_or_alias)
            .cloned()
            .unwrap_or_else(|| id_or_alias.to_string())
    }

    /// Get cache directory path
    pub fn get_cache_dir(&self) -> &Path {
        &self.cache_dir
    }

    fn get_file_extension(&self, format: &ModelFormat) -> &str {
        match format {
            ModelFormat::GGUF => "gguf",
            ModelFormat::SafeTensors => "safetensors",
            ModelFormat::PyTorch => "pt",
            ModelFormat::ONNX => "onnx",
            ModelFormat::TensorFlow => "pb",
            _ => "bin",
        }
    }

    fn get_cache_size(&self) -> Result<u64> {
        let mut total_size = 0u64;

        if self.cache_dir.exists() {
            for entry in fs::read_dir(&self.cache_dir)? {
                let entry = entry?;
                let metadata = entry.metadata()?;
                if metadata.is_file() {
                    total_size += metadata.len();
                }
            }
        }

        Ok(total_size)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageConfig {
    pub custom_data_dir: Option<PathBuf>,
    pub custom_config_dir: Option<PathBuf>,
    pub max_cache_size_gb: Option<u64>,
    pub auto_cleanup_days: Option<u64>,
}

impl Default for StorageConfig {
    fn default() -> Self {
        Self {
            custom_data_dir: None,
            custom_config_dir: None,
            max_cache_size_gb: Some(10), // 10GB default cache limit
            auto_cleanup_days: Some(30), // Clean up cache files older than 30 days
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageStats {
    pub model_count: usize,
    pub total_size: u64,
    pub format_breakdown: HashMap<ModelFormat, usize>,
    pub cache_size: u64,
    pub data_dir: PathBuf,
    pub config_dir: PathBuf,
}

impl StorageStats {
    pub fn total_size_human(&self) -> String {
        human_bytes(self.total_size)
    }

    pub fn cache_size_human(&self) -> String {
        human_bytes(self.cache_size)
    }
}

fn human_bytes(bytes: u64) -> String {
    const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
    let mut size = bytes as f64;
    let mut unit_index = 0;

    while size >= 1024.0 && unit_index < UNITS.len() - 1 {
        size /= 1024.0;
        unit_index += 1;
    }

    if unit_index == 0 {
        format!("{} {}", bytes, UNITS[unit_index])
    } else {
        format!("{:.2} {}", size, UNITS[unit_index])
    }
}

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

    /// The chunked hash of a file that spans a chunk boundary matches the
    /// one-shot digest of its bytes.
    #[test]
    fn sha256_file_matches_the_one_shot_digest() {
        let path = std::env::temp_dir().join(format!("ae-sha256-file-{}", std::process::id()));
        let bytes: Vec<u8> = (0..(1usize << 20) + 3).map(|i| (i % 251) as u8).collect();
        fs::write(&path, &bytes).unwrap();
        let chunked = sha256_file(&path).unwrap();
        fs::remove_file(&path).unwrap();
        assert_eq!(chunked, hex::encode(Sha256::digest(&bytes).unwrap()));
    }

    fn sample_metadata(id: &str, file_path: &str) -> LocalModelMetadata {
        LocalModelMetadata {
            id: id.to_string(),
            name: id.to_string(),
            description: None,
            version: "test".to_string(),
            format: ModelFormat::GGUF,
            file_path: file_path.to_string(),
            config_path: None,
            tokenizer_path: None,
            size_bytes: 0,
            sha256: String::new(),
            downloaded_at: Utc::now(),
            last_used: None,
            usage_count: 0,
            capabilities: ModelCapabilities {
                chat: false,
                completions: false,
                embeddings: false,
                image_generation: false,
                image_understanding: false,
                audio_generation: false,
                audio_understanding: false,
                video_understanding: false,
                function_calling: false,
                streaming: false,
            },
            parameters: HashMap::new(),
            source: ModelSource {
                origin: "local".to_string(),
                url: None,
                repository: None,
                commit: None,
                license: None,
            },
        }
    }

    /// A model whose file cannot be deleted stays listed, and a retry after
    /// the obstacle is gone removes it. Uses a temp dir, never the real data dir.
    #[test]
    fn remove_model_keeps_entry_when_its_file_cannot_be_deleted() {
        let base = std::env::temp_dir().join(format!("ae-remove-model-{}", std::process::id()));
        let _ = fs::remove_dir_all(&base);
        let config = StorageConfig {
            custom_data_dir: Some(base.join("data")),
            custom_config_dir: Some(base.join("config")),
            ..StorageConfig::default()
        };
        let mut storage = ModelStorage::new(&config).unwrap();

        // A directory where the model file should be: remove_file fails on it.
        let blocked = base.join("data").join("blocked.gguf");
        fs::create_dir_all(&blocked).unwrap();
        let model_id = "local/blocked";
        storage.index.models.insert(
            model_id.to_string(),
            sample_metadata(model_id, &blocked.to_string_lossy()),
        );

        assert!(storage.remove_model(model_id).is_err());
        assert!(storage.get_model_metadata(model_id).is_some());

        fs::remove_dir(&blocked).unwrap();
        assert!(storage.remove_model(model_id).is_ok());
        assert!(storage.get_model_metadata(model_id).is_none());

        let _ = fs::remove_dir_all(&base);
    }
}