//! Pre-trained model management and downloading.
//!
//! Provides a registry of available artifact detection models and
//! utilities for downloading and caching them.

use crate::error::{ArtifactError, ArtifactResult};
use crate::labels::ArtifactType;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};

/// Information about a pre-trained model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
    /// Unique model identifier
    pub id: String,
    /// Human-readable name
    pub name: String,
    /// Model description
    pub description: String,
    /// Model version
    pub version: String,
    /// Source URL for downloading
    pub source: ModelSource,
    /// Expected input channels
    pub n_channels: usize,
    /// Expected sampling frequency
    pub sfreq: f64,
    /// Window size the model was trained on
    pub window_size: usize,
    /// Artifact types this model can detect
    pub artifact_types: Vec<ArtifactType>,
    /// File size in bytes
    pub file_size: u64,
    /// SHA256 hash for verification
    pub sha256: Option<String>,
    /// License information
    pub license: String,
    /// Training dataset description
    pub training_data: String,
}

impl ModelInfo {
    /// Create a new model info
    pub fn new(id: impl Into<String>, name: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            name: name.into(),
            description: String::new(),
            version: "1.0.0".to_string(),
            source: ModelSource::Local(PathBuf::new()),
            n_channels: 64,
            sfreq: 1000.0,
            window_size: 1000,
            artifact_types: ArtifactType::all(),
            file_size: 0,
            sha256: None,
            license: "MIT".to_string(),
            training_data: String::new(),
        }
    }

    /// Set description
    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
        self.description = desc.into();
        self
    }

    /// Set source
    pub fn with_source(mut self, source: ModelSource) -> Self {
        self.source = source;
        self
    }

    /// Set input configuration
    pub fn with_input_config(mut self, n_channels: usize, sfreq: f64, window_size: usize) -> Self {
        self.n_channels = n_channels;
        self.sfreq = sfreq;
        self.window_size = window_size;
        self
    }

    /// Set supported artifact types
    pub fn with_artifact_types(mut self, types: Vec<ArtifactType>) -> Self {
        self.artifact_types = types;
        self
    }

    /// Check if model supports a specific artifact type
    pub fn supports(&self, artifact_type: ArtifactType) -> bool {
        self.artifact_types.contains(&artifact_type)
    }
}

/// Source location for a model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ModelSource {
    /// Local file path
    Local(PathBuf),
    /// Remote URL
    Url(String),
    /// Hugging Face Hub
    HuggingFace {
        /// Repository identifier (e.g., "username/model-name")
        repo: String,
        /// Name of the model file within the repository
        filename: String,
    },
    /// GitHub release
    GitHub {
        /// Repository owner (username or organization)
        owner: String,
        /// Repository name
        repo: String,
        /// Release tag (e.g., "v1.0.0")
        tag: String,
        /// Release asset filename
        asset: String,
    },
}

impl ModelSource {
    /// Create a local source
    pub fn local(path: impl Into<PathBuf>) -> Self {
        Self::Local(path.into())
    }

    /// Create a URL source
    pub fn url(url: impl Into<String>) -> Self {
        Self::Url(url.into())
    }

    /// Create a Hugging Face source
    pub fn huggingface(repo: impl Into<String>, filename: impl Into<String>) -> Self {
        Self::HuggingFace {
            repo: repo.into(),
            filename: filename.into(),
        }
    }

    /// Create a GitHub release source
    pub fn github(
        owner: impl Into<String>,
        repo: impl Into<String>,
        tag: impl Into<String>,
        asset: impl Into<String>,
    ) -> Self {
        Self::GitHub {
            owner: owner.into(),
            repo: repo.into(),
            tag: tag.into(),
            asset: asset.into(),
        }
    }

    /// Get the download URL
    pub fn to_url(&self) -> Option<String> {
        match self {
            Self::Local(_) => None,
            Self::Url(url) => Some(url.clone()),
            Self::HuggingFace { repo, filename } => Some(format!(
                "https://huggingface.co/{}/resolve/main/{}",
                repo, filename
            )),
            Self::GitHub {
                owner,
                repo,
                tag,
                asset,
            } => Some(format!(
                "https://github.com/{}/{}/releases/download/{}/{}",
                owner, repo, tag, asset
            )),
        }
    }

    /// Check if this is a local source
    pub fn is_local(&self) -> bool {
        matches!(self, Self::Local(_))
    }
}

/// Registry of available pre-trained models
#[derive(Debug, Clone)]
pub struct ModelRegistry {
    /// Available models by ID
    models: HashMap<String, ModelInfo>,
    /// Cache directory for downloaded models
    cache_dir: PathBuf,
}

impl Default for ModelRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl ModelRegistry {
    /// Create a new model registry with default models
    pub fn new() -> Self {
        let cache_dir = Self::default_cache_dir();
        let mut registry = Self {
            models: HashMap::new(),
            cache_dir,
        };

        // Register built-in models
        registry.register_builtin_models();

        registry
    }

    /// Create with custom cache directory
    pub fn with_cache_dir(cache_dir: impl Into<PathBuf>) -> Self {
        let cache_dir = cache_dir.into();
        let mut registry = Self {
            models: HashMap::new(),
            cache_dir,
        };
        registry.register_builtin_models();
        registry
    }

    /// Get default cache directory
    pub fn default_cache_dir() -> PathBuf {
        dirs::cache_dir()
            .unwrap_or_else(|| PathBuf::from("."))
            .join("rustytorch")
            .join("models")
            .join("neuro-artifacts")
    }

    /// Register built-in models
    fn register_builtin_models(&mut self) {
        // EEG artifact detector (general purpose)
        let eeg_general = ModelInfo::new("eeg-artifact-v1", "EEG Artifact Detector v1")
            .with_description(
                "General-purpose EEG artifact detection model trained on diverse datasets",
            )
            .with_source(ModelSource::huggingface(
                "rustytorch/neuro-artifacts",
                "eeg-artifact-v1.onnx",
            ))
            .with_input_config(64, 1000.0, 1000)
            .with_artifact_types(vec![
                ArtifactType::EyeBlink,
                ArtifactType::EyeMovement,
                ArtifactType::Muscle,
                ArtifactType::Heartbeat,
                ArtifactType::LineNoise,
                ArtifactType::Movement,
            ]);
        self.register(eeg_general);

        // MEG artifact detector
        let meg_general = ModelInfo::new("meg-artifact-v1", "MEG Artifact Detector v1")
            .with_description(
                "MEG artifact detection model optimized for magnetometer/gradiometer data",
            )
            .with_source(ModelSource::huggingface(
                "rustytorch/neuro-artifacts",
                "meg-artifact-v1.onnx",
            ))
            .with_input_config(306, 1000.0, 1000) // 306 channels for Elekta Neuromag
            .with_artifact_types(ArtifactType::all());
        self.register(meg_general);

        // Eye artifact specialist
        let eye_artifact = ModelInfo::new("eye-artifact-v1", "Eye Artifact Specialist v1")
            .with_description(
                "Specialized model for detecting eye-related artifacts with high precision",
            )
            .with_source(ModelSource::huggingface(
                "rustytorch/neuro-artifacts",
                "eye-artifact-v1.onnx",
            ))
            .with_input_config(64, 500.0, 500)
            .with_artifact_types(vec![ArtifactType::EyeBlink, ArtifactType::EyeMovement]);
        self.register(eye_artifact);

        // Muscle artifact specialist
        let muscle_artifact = ModelInfo::new("muscle-artifact-v1", "Muscle Artifact Specialist v1")
            .with_description("Specialized model for detecting EMG contamination")
            .with_source(ModelSource::huggingface(
                "rustytorch/neuro-artifacts",
                "muscle-artifact-v1.onnx",
            ))
            .with_input_config(64, 1000.0, 500)
            .with_artifact_types(vec![ArtifactType::Muscle]);
        self.register(muscle_artifact);

        // Cardiac artifact specialist
        let cardiac_artifact =
            ModelInfo::new("cardiac-artifact-v1", "Cardiac Artifact Specialist v1")
                .with_description("Specialized model for detecting heartbeat artifacts")
                .with_source(ModelSource::huggingface(
                    "rustytorch/neuro-artifacts",
                    "cardiac-artifact-v1.onnx",
                ))
                .with_input_config(64, 500.0, 1000)
                .with_artifact_types(vec![ArtifactType::Heartbeat]);
        self.register(cardiac_artifact);

        // Lightweight model for real-time
        let realtime = ModelInfo::new("realtime-artifact-v1", "Real-Time Artifact Detector v1")
            .with_description(
                "Lightweight model optimized for real-time artifact detection with minimal latency",
            )
            .with_source(ModelSource::huggingface(
                "rustytorch/neuro-artifacts",
                "realtime-artifact-v1.onnx",
            ))
            .with_input_config(32, 500.0, 250) // Smaller window for low latency
            .with_artifact_types(vec![
                ArtifactType::EyeBlink,
                ArtifactType::Muscle,
                ArtifactType::Heartbeat,
            ]);
        self.register(realtime);
    }

    /// Register a model
    pub fn register(&mut self, model: ModelInfo) {
        self.models.insert(model.id.clone(), model);
    }

    /// Get model info by ID
    pub fn get(&self, id: &str) -> Option<&ModelInfo> {
        self.models.get(id)
    }

    /// List all available models
    pub fn list(&self) -> Vec<&ModelInfo> {
        self.models.values().collect()
    }

    /// List models that support a specific artifact type
    pub fn list_for_artifact(&self, artifact_type: ArtifactType) -> Vec<&ModelInfo> {
        self.models
            .values()
            .filter(|m| m.supports(artifact_type))
            .collect()
    }

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

    /// Get cached path for a model
    pub fn cached_path(&self, model_id: &str) -> PathBuf {
        self.cache_dir.join(format!("{}.onnx", model_id))
    }

    /// Check if model is cached locally
    pub fn is_cached(&self, model_id: &str) -> bool {
        self.cached_path(model_id).exists()
    }

    /// Get model path (cached or local)
    pub fn get_path(&self, model_id: &str) -> ArtifactResult<PathBuf> {
        let model = self
            .get(model_id)
            .ok_or_else(|| ArtifactError::Model(format!("Unknown model: {}", model_id)))?;

        if let ModelSource::Local(path) = &model.source {
            if path.exists() {
                Ok(path.clone())
            } else {
                Err(ArtifactError::Model(format!(
                    "Model file not found: {}",
                    path.display()
                )))
            }
        } else {
            let cached = self.cached_path(model_id);
            if cached.exists() {
                Ok(cached)
            } else {
                Err(ArtifactError::Model(format!(
                    "Model not cached. Call download_model first: {}",
                    model_id
                )))
            }
        }
    }
}

/// Download a model to the cache
pub async fn download_model(registry: &ModelRegistry, model_id: &str) -> ArtifactResult<PathBuf> {
    let model = registry
        .get(model_id)
        .ok_or_else(|| ArtifactError::Model(format!("Unknown model: {}", model_id)))?;

    // Check if already cached
    let cache_path = registry.cached_path(model_id);
    if cache_path.exists() {
        return Ok(cache_path);
    }

    // Get download URL
    let url = model.source.to_url().ok_or_else(|| {
        ArtifactError::Download("Model source is local, cannot download".to_string())
    })?;

    // Ensure cache directory exists
    if let Some(parent) = cache_path.parent() {
        tokio::fs::create_dir_all(parent).await.map_err(|e| {
            ArtifactError::Download(format!("Failed to create cache directory: {}", e))
        })?;
    }

    // Download file
    download_file(&url, &cache_path).await?;

    // Verify hash if provided
    if let Some(ref expected_hash) = model.sha256 {
        let actual_hash = compute_file_hash(&cache_path).await?;
        if actual_hash != *expected_hash {
            // Remove corrupted file
            let _ = tokio::fs::remove_file(&cache_path).await;
            return Err(ArtifactError::Download(format!(
                "Hash mismatch: expected {}, got {}",
                expected_hash, actual_hash
            )));
        }
    }

    Ok(cache_path)
}

/// Download a file from URL
async fn download_file(url: &str, path: &Path) -> ArtifactResult<()> {
    // Use reqwest if available, otherwise fall back to simpler method
    #[cfg(feature = "download")]
    {
        let response = reqwest::get(url)
            .await
            .map_err(|e| ArtifactError::Download(format!("Failed to download: {}", e)))?;

        if !response.status().is_success() {
            return Err(ArtifactError::Download(format!(
                "Download failed with status: {}",
                response.status()
            )));
        }

        let bytes = response
            .bytes()
            .await
            .map_err(|e| ArtifactError::Download(format!("Failed to read response: {}", e)))?;

        tokio::fs::write(path, bytes)
            .await
            .map_err(|e| ArtifactError::Download(format!("Failed to write file: {}", e)))?;

        Ok(())
    }

    #[cfg(not(feature = "download"))]
    {
        // Without download feature, we can't download
        Err(ArtifactError::Download(format!(
            "Download feature not enabled. URL: {}, Path: {}",
            url,
            path.display()
        )))
    }
}

/// Compute SHA256 hash of a file
async fn compute_file_hash(path: &Path) -> ArtifactResult<String> {
    use std::io::Read;

    let path = path.to_path_buf();
    tokio::task::spawn_blocking(move || {
        let mut file = std::fs::File::open(&path).map_err(|e| {
            ArtifactError::Download(format!("Failed to open file for hashing: {}", e))
        })?;

        let mut hasher = Sha256::new();
        let mut buffer = [0u8; 8192];

        loop {
            let n = file
                .read(&mut buffer)
                .map_err(|e| ArtifactError::Download(format!("Failed to read file: {}", e)))?;
            if n == 0 {
                break;
            }
            hasher.update(&buffer[..n]);
        }

        let hash_bytes = hasher.finalize_inner();
        let hex: String = hash_bytes.iter().map(|b| format!("{:02x}", b)).collect();
        Ok(hex)
    })
    .await
    .map_err(|e| ArtifactError::Download(format!("Hash computation failed: {}", e)))?
}

/// Simple SHA256 implementation for verification
struct Sha256 {
    h: [u32; 8],
    block: [u8; 64],
    block_len: usize,
    total_len: u64,
}

impl Sha256 {
    fn new() -> Self {
        Self {
            h: [
                0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab,
                0x5be0cd19,
            ],
            block: [0; 64],
            block_len: 0,
            total_len: 0,
        }
    }

    fn update(&mut self, data: &[u8]) {
        self.total_len += data.len() as u64;
        let mut offset = 0;

        // Fill current block
        if self.block_len > 0 {
            let space = 64 - self.block_len;
            let to_copy = space.min(data.len());
            self.block[self.block_len..self.block_len + to_copy].copy_from_slice(&data[..to_copy]);
            self.block_len += to_copy;
            offset = to_copy;

            if self.block_len == 64 {
                self.process_block();
                self.block_len = 0;
            }
        }

        // Process full blocks
        while offset + 64 <= data.len() {
            self.block.copy_from_slice(&data[offset..offset + 64]);
            self.process_block();
            offset += 64;
        }

        // Save remaining
        if offset < data.len() {
            let remaining = data.len() - offset;
            self.block[..remaining].copy_from_slice(&data[offset..]);
            self.block_len = remaining;
        }
    }

    fn process_block(&mut self) {
        const K: [u32; 64] = [
            0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4,
            0xab1c5ed5, 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe,
            0x9bdc06a7, 0xc19bf174, 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f,
            0x4a7484aa, 0x5cb0a9dc, 0x76f988da, 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7,
            0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc,
            0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, 0xa2bfe8a1, 0xa81a664b,
            0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, 0x19a4c116,
            0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
            0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7,
            0xc67178f2,
        ];

        let mut w = [0u32; 64];

        // Prepare message schedule
        for i in 0..16 {
            w[i] = u32::from_be_bytes([
                self.block[i * 4],
                self.block[i * 4 + 1],
                self.block[i * 4 + 2],
                self.block[i * 4 + 3],
            ]);
        }

        for i in 16..64 {
            let s0 = w[i - 15].rotate_right(7) ^ w[i - 15].rotate_right(18) ^ (w[i - 15] >> 3);
            let s1 = w[i - 2].rotate_right(17) ^ w[i - 2].rotate_right(19) ^ (w[i - 2] >> 10);
            w[i] = w[i - 16]
                .wrapping_add(s0)
                .wrapping_add(w[i - 7])
                .wrapping_add(s1);
        }

        // Working variables
        let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = self.h;

        // Compression function
        for i in 0..64 {
            let s1 = e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25);
            let ch = (e & f) ^ ((!e) & g);
            let temp1 = h
                .wrapping_add(s1)
                .wrapping_add(ch)
                .wrapping_add(K[i])
                .wrapping_add(w[i]);
            let s0 = a.rotate_right(2) ^ a.rotate_right(13) ^ a.rotate_right(22);
            let maj = (a & b) ^ (a & c) ^ (b & c);
            let temp2 = s0.wrapping_add(maj);

            h = g;
            g = f;
            f = e;
            e = d.wrapping_add(temp1);
            d = c;
            c = b;
            b = a;
            a = temp1.wrapping_add(temp2);
        }

        // Update hash values
        self.h[0] = self.h[0].wrapping_add(a);
        self.h[1] = self.h[1].wrapping_add(b);
        self.h[2] = self.h[2].wrapping_add(c);
        self.h[3] = self.h[3].wrapping_add(d);
        self.h[4] = self.h[4].wrapping_add(e);
        self.h[5] = self.h[5].wrapping_add(f);
        self.h[6] = self.h[6].wrapping_add(g);
        self.h[7] = self.h[7].wrapping_add(h);
    }
}

/// Format helper for hex output
struct HexFormatter([u8; 32]);

impl std::fmt::LowerHex for HexFormatter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for byte in &self.0 {
            write!(f, "{:02x}", byte)?;
        }
        Ok(())
    }
}

impl Sha256 {
    fn finalize(self) -> HexFormatter {
        HexFormatter(self.finalize_inner())
    }

    fn finalize_inner(mut self) -> [u8; 32] {
        // Padding
        let bit_len = self.total_len * 8;
        self.block[self.block_len] = 0x80;
        self.block_len += 1;

        if self.block_len > 56 {
            self.block[self.block_len..64].fill(0);
            self.process_block();
            self.block_len = 0;
        }

        self.block[self.block_len..56].fill(0);
        self.block[56..64].copy_from_slice(&bit_len.to_be_bytes());
        self.process_block();

        // Output
        let mut result = [0u8; 32];
        for (i, &h) in self.h.iter().enumerate() {
            result[i * 4..(i + 1) * 4].copy_from_slice(&h.to_be_bytes());
        }
        result
    }
}

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

    #[test]
    fn test_model_info_creation() {
        let model = ModelInfo::new("test-model", "Test Model")
            .with_description("A test model")
            .with_input_config(32, 500.0, 500);

        assert_eq!(model.id, "test-model");
        assert_eq!(model.name, "Test Model");
        assert_eq!(model.n_channels, 32);
        assert_eq!(model.sfreq, 500.0);
    }

    #[test]
    fn test_model_source_url() {
        let hf = ModelSource::huggingface("user/repo", "model.onnx");
        assert_eq!(
            hf.to_url(),
            Some("https://huggingface.co/user/repo/resolve/main/model.onnx".to_string())
        );

        let gh = ModelSource::github("owner", "repo", "v1.0.0", "model.onnx");
        assert_eq!(
            gh.to_url(),
            Some("https://github.com/owner/repo/releases/download/v1.0.0/model.onnx".to_string())
        );

        let local = ModelSource::local("/path/to/model.onnx");
        assert!(local.to_url().is_none());
    }

    #[test]
    fn test_model_registry() {
        let registry = ModelRegistry::new();

        // Check built-in models are registered
        assert!(registry.get("eeg-artifact-v1").is_some());
        assert!(registry.get("meg-artifact-v1").is_some());

        // List all models
        let models = registry.list();
        assert!(!models.is_empty());
    }

    #[test]
    fn test_model_supports() {
        let model = ModelInfo::new("test", "Test")
            .with_artifact_types(vec![ArtifactType::EyeBlink, ArtifactType::Muscle]);

        assert!(model.supports(ArtifactType::EyeBlink));
        assert!(model.supports(ArtifactType::Muscle));
        assert!(!model.supports(ArtifactType::Heartbeat));
    }

    #[test]
    fn test_list_for_artifact() {
        let registry = ModelRegistry::new();

        let eye_models = registry.list_for_artifact(ArtifactType::EyeBlink);
        assert!(!eye_models.is_empty());

        // All eye models should support EyeBlink
        for model in eye_models {
            assert!(model.supports(ArtifactType::EyeBlink));
        }
    }

    #[test]
    fn test_cached_path() {
        let registry = ModelRegistry::new();
        let path = registry.cached_path("test-model");
        assert!(path.to_string_lossy().contains("test-model.onnx"));
    }

    #[test]
    fn test_sha256() {
        let mut hasher = Sha256::new();
        hasher.update(b"hello world");
        let result = format!("{:x}", hasher.finalize());
        assert_eq!(
            result,
            "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9"
        );
    }

    #[test]
    fn test_sha256_empty() {
        let hasher = Sha256::new();
        let result = format!("{:x}", hasher.finalize());
        assert_eq!(
            result,
            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
        );
    }

    #[test]
    fn test_sha256_large() {
        let mut hasher = Sha256::new();
        hasher.update(&[0u8; 1000]);
        let result = format!("{:x}", hasher.finalize());
        // Just verify it completes without error and produces valid hex
        assert_eq!(result.len(), 64);
    }
}