//! Model registry and management for Synor Compute.
//!
//! Provides:
//! - Pre-registered popular models (LLMs, vision, audio)
//! - Custom model uploads
//! - Model format conversion
//! - Model caching and warm-up

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use parking_lot::RwLock;

use crate::error::ComputeError;
use crate::processor::{Precision, ProcessorType};

/// Model identifier (CID or alias).
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ModelId(pub String);

impl ModelId {
    /// Creates from a CID.
    pub fn from_cid(cid: &str) -> Self {
        Self(cid.to_string())
    }

    /// Creates from an alias.
    pub fn from_alias(alias: &str) -> Self {
        Self(alias.to_string())
    }

    /// Checks if this is a CID (starts with Qm or bafy).
    pub fn is_cid(&self) -> bool {
        self.0.starts_with("Qm") || self.0.starts_with("bafy")
    }
}

impl std::fmt::Display for ModelId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl From<&str> for ModelId {
    fn from(s: &str) -> Self {
        Self(s.to_string())
    }
}

/// Model format specification.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ModelFormat {
    /// ONNX Runtime format.
    Onnx,
    /// PyTorch checkpoint (.pt, .pth).
    PyTorch,
    /// PyTorch TorchScript (.pt).
    TorchScript,
    /// TensorFlow SavedModel.
    TensorFlow,
    /// TensorFlow Lite.
    TFLite,
    /// SafeTensors format.
    SafeTensors,
    /// GGUF format (llama.cpp).
    Gguf,
    /// GGML format (legacy).
    Ggml,
    /// Custom binary format.
    Custom,
}

/// Model category.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ModelCategory {
    /// Large language models.
    Llm,
    /// Text embedding models.
    Embedding,
    /// Image classification.
    ImageClassification,
    /// Object detection.
    ObjectDetection,
    /// Image segmentation.
    Segmentation,
    /// Image generation (diffusion).
    ImageGeneration,
    /// Speech-to-text.
    SpeechToText,
    /// Text-to-speech.
    TextToSpeech,
    /// Video generation.
    VideoGeneration,
    /// Multi-modal (vision-language).
    MultiModal,
    /// Custom/other.
    Custom,
}

/// Model metadata.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ModelInfo {
    /// Model ID (alias or CID).
    pub id: ModelId,
    /// Human-readable name.
    pub name: String,
    /// Description.
    pub description: String,
    /// Model category.
    pub category: ModelCategory,
    /// Storage CID (actual location).
    pub cid: String,
    /// Model format.
    pub format: ModelFormat,
    /// Model size in bytes.
    pub size_bytes: u64,
    /// Parameter count.
    pub parameters: u64,
    /// Supported precisions.
    pub supported_precisions: Vec<Precision>,
    /// Recommended processor type.
    pub recommended_processor: ProcessorType,
    /// Context length (for LLMs).
    pub context_length: Option<u32>,
    /// Input schema (JSON Schema).
    pub input_schema: Option<String>,
    /// Output schema (JSON Schema).
    pub output_schema: Option<String>,
    /// License.
    pub license: String,
    /// Provider/author.
    pub provider: String,
    /// Version.
    pub version: String,
    /// Is public (anyone can use).
    pub is_public: bool,
    /// Owner address (for private models).
    pub owner: Option<[u8; 32]>,
}

impl ModelInfo {
    /// Creates a new LLM model info.
    pub fn llm(
        alias: &str,
        name: &str,
        cid: &str,
        parameters: u64,
        context_length: u32,
    ) -> Self {
        Self {
            id: ModelId::from_alias(alias),
            name: name.to_string(),
            description: format!("{} - {} parameter LLM", name, format_params(parameters)),
            category: ModelCategory::Llm,
            cid: cid.to_string(),
            format: ModelFormat::SafeTensors,
            size_bytes: parameters * 2, // ~2 bytes per param in fp16
            parameters,
            supported_precisions: vec![Precision::Fp16, Precision::Bf16, Precision::Int8, Precision::Int4],
            recommended_processor: ProcessorType::Lpu,
            context_length: Some(context_length),
            input_schema: None,
            output_schema: None,
            license: "Apache-2.0".to_string(),
            provider: "Synor".to_string(),
            version: "1.0".to_string(),
            is_public: true,
            owner: None,
        }
    }

    /// Creates a custom model info.
    pub fn custom(
        cid: &str,
        name: &str,
        category: ModelCategory,
        format: ModelFormat,
        size_bytes: u64,
        owner: [u8; 32],
    ) -> Self {
        Self {
            id: ModelId::from_cid(cid),
            name: name.to_string(),
            description: format!("Custom model: {}", name),
            category,
            cid: cid.to_string(),
            format,
            size_bytes,
            parameters: 0,
            supported_precisions: vec![Precision::Fp32, Precision::Fp16],
            recommended_processor: ProcessorType::Cpu(crate::processor::CpuVariant::default()),
            context_length: None,
            input_schema: None,
            output_schema: None,
            license: "Custom".to_string(),
            provider: "User".to_string(),
            version: "1.0".to_string(),
            is_public: false,
            owner: Some(owner),
        }
    }
}

/// Format parameter count for display.
fn format_params(params: u64) -> String {
    if params >= 1_000_000_000_000 {
        format!("{:.1}T", params as f64 / 1e12)
    } else if params >= 1_000_000_000 {
        format!("{:.1}B", params as f64 / 1e9)
    } else if params >= 1_000_000 {
        format!("{:.1}M", params as f64 / 1e6)
    } else {
        format!("{}", params)
    }
}

/// Model registry.
pub struct ModelRegistry {
    /// Registered models by ID/alias.
    models: RwLock<HashMap<String, ModelInfo>>,
    /// Alias to CID mapping.
    aliases: RwLock<HashMap<String, String>>,
}

impl ModelRegistry {
    /// Creates a new model registry with default models.
    pub fn new() -> Self {
        let registry = Self {
            models: RwLock::new(HashMap::new()),
            aliases: RwLock::new(HashMap::new()),
        };
        registry.register_default_models();
        registry
    }

    /// Registers default/popular models.
    fn register_default_models(&self) {
        let default_models = vec![
            // ===== LLMs =====
            // Llama 3 family
            ModelInfo::llm("llama-3-8b", "Llama 3 8B", "QmLlama3_8B_placeholder", 8_000_000_000, 8192),
            ModelInfo::llm("llama-3-70b", "Llama 3 70B", "QmLlama3_70B_placeholder", 70_000_000_000, 8192),
            ModelInfo::llm("llama-3.1-8b", "Llama 3.1 8B", "QmLlama31_8B_placeholder", 8_000_000_000, 128000),
            ModelInfo::llm("llama-3.1-70b", "Llama 3.1 70B", "QmLlama31_70B_placeholder", 70_000_000_000, 128000),
            ModelInfo::llm("llama-3.1-405b", "Llama 3.1 405B", "QmLlama31_405B_placeholder", 405_000_000_000, 128000),

            // Mistral family
            ModelInfo::llm("mistral-7b", "Mistral 7B", "QmMistral7B_placeholder", 7_000_000_000, 32768),
            ModelInfo::llm("mixtral-8x7b", "Mixtral 8x7B", "QmMixtral8x7B_placeholder", 46_000_000_000, 32768),
            ModelInfo::llm("mixtral-8x22b", "Mixtral 8x22B", "QmMixtral8x22B_placeholder", 176_000_000_000, 65536),

            // Qwen family
            ModelInfo::llm("qwen-2.5-7b", "Qwen 2.5 7B", "QmQwen25_7B_placeholder", 7_000_000_000, 128000),
            ModelInfo::llm("qwen-2.5-72b", "Qwen 2.5 72B", "QmQwen25_72B_placeholder", 72_000_000_000, 128000),

            // DeepSeek family
            ModelInfo::llm("deepseek-v2", "DeepSeek V2", "QmDeepSeekV2_placeholder", 236_000_000_000, 128000),
            ModelInfo::llm("deepseek-coder-33b", "DeepSeek Coder 33B", "QmDeepSeekCoder33B_placeholder", 33_000_000_000, 16384),

            // Phi family (small/efficient)
            ModelInfo::llm("phi-3-mini", "Phi 3 Mini", "QmPhi3Mini_placeholder", 3_800_000_000, 128000),
            ModelInfo::llm("phi-3-medium", "Phi 3 Medium", "QmPhi3Medium_placeholder", 14_000_000_000, 128000),

            // Code models
            ModelInfo::llm("codellama-34b", "Code Llama 34B", "QmCodeLlama34B_placeholder", 34_000_000_000, 16384),
            ModelInfo::llm("starcoder2-15b", "StarCoder2 15B", "QmStarCoder2_15B_placeholder", 15_000_000_000, 16384),

            // ===== Embedding Models =====
            ModelInfo {
                id: ModelId::from_alias("bge-large"),
                name: "BGE Large".to_string(),
                description: "BAAI General Embedding - Large".to_string(),
                category: ModelCategory::Embedding,
                cid: "QmBGELarge_placeholder".to_string(),
                format: ModelFormat::SafeTensors,
                size_bytes: 1_300_000_000,
                parameters: 335_000_000,
                supported_precisions: vec![Precision::Fp32, Precision::Fp16],
                recommended_processor: ProcessorType::Gpu(crate::processor::GpuVariant::default()),
                context_length: Some(512),
                input_schema: None,
                output_schema: None,
                license: "MIT".to_string(),
                provider: "BAAI".to_string(),
                version: "1.5".to_string(),
                is_public: true,
                owner: None,
            },
            ModelInfo {
                id: ModelId::from_alias("e5-large-v2"),
                name: "E5 Large v2".to_string(),
                description: "Microsoft E5 Embedding - Large".to_string(),
                category: ModelCategory::Embedding,
                cid: "QmE5LargeV2_placeholder".to_string(),
                format: ModelFormat::SafeTensors,
                size_bytes: 1_300_000_000,
                parameters: 335_000_000,
                supported_precisions: vec![Precision::Fp32, Precision::Fp16],
                recommended_processor: ProcessorType::Gpu(crate::processor::GpuVariant::default()),
                context_length: Some(512),
                input_schema: None,
                output_schema: None,
                license: "MIT".to_string(),
                provider: "Microsoft".to_string(),
                version: "2.0".to_string(),
                is_public: true,
                owner: None,
            },

            // ===== Vision Models =====
            ModelInfo {
                id: ModelId::from_alias("stable-diffusion-xl"),
                name: "Stable Diffusion XL".to_string(),
                description: "SDXL 1.0 - High quality image generation".to_string(),
                category: ModelCategory::ImageGeneration,
                cid: "QmSDXL_placeholder".to_string(),
                format: ModelFormat::SafeTensors,
                size_bytes: 6_900_000_000,
                parameters: 2_600_000_000,
                supported_precisions: vec![Precision::Fp16, Precision::Bf16],
                recommended_processor: ProcessorType::Gpu(crate::processor::GpuVariant::default()),
                context_length: None,
                input_schema: None,
                output_schema: None,
                license: "CreativeML Open RAIL++-M".to_string(),
                provider: "Stability AI".to_string(),
                version: "1.0".to_string(),
                is_public: true,
                owner: None,
            },
            ModelInfo {
                id: ModelId::from_alias("flux-1-dev"),
                name: "FLUX.1 Dev".to_string(),
                description: "FLUX.1 Development - State of the art image generation".to_string(),
                category: ModelCategory::ImageGeneration,
                cid: "QmFLUX1Dev_placeholder".to_string(),
                format: ModelFormat::SafeTensors,
                size_bytes: 23_000_000_000,
                parameters: 12_000_000_000,
                supported_precisions: vec![Precision::Fp16, Precision::Bf16],
                recommended_processor: ProcessorType::Gpu(crate::processor::GpuVariant::default()),
                context_length: None,
                input_schema: None,
                output_schema: None,
                license: "FLUX.1-dev Non-Commercial".to_string(),
                provider: "Black Forest Labs".to_string(),
                version: "1.0".to_string(),
                is_public: true,
                owner: None,
            },

            // ===== Speech Models =====
            ModelInfo {
                id: ModelId::from_alias("whisper-large-v3"),
                name: "Whisper Large v3".to_string(),
                description: "OpenAI Whisper - Speech recognition".to_string(),
                category: ModelCategory::SpeechToText,
                cid: "QmWhisperLargeV3_placeholder".to_string(),
                format: ModelFormat::SafeTensors,
                size_bytes: 3_100_000_000,
                parameters: 1_550_000_000,
                supported_precisions: vec![Precision::Fp32, Precision::Fp16],
                recommended_processor: ProcessorType::Gpu(crate::processor::GpuVariant::default()),
                context_length: None,
                input_schema: None,
                output_schema: None,
                license: "MIT".to_string(),
                provider: "OpenAI".to_string(),
                version: "3.0".to_string(),
                is_public: true,
                owner: None,
            },

            // ===== Multi-Modal Models =====
            ModelInfo {
                id: ModelId::from_alias("llava-1.5-13b"),
                name: "LLaVA 1.5 13B".to_string(),
                description: "Large Language and Vision Assistant".to_string(),
                category: ModelCategory::MultiModal,
                cid: "QmLLaVA15_13B_placeholder".to_string(),
                format: ModelFormat::SafeTensors,
                size_bytes: 26_000_000_000,
                parameters: 13_000_000_000,
                supported_precisions: vec![Precision::Fp16, Precision::Int8],
                recommended_processor: ProcessorType::Gpu(crate::processor::GpuVariant::default()),
                context_length: Some(4096),
                input_schema: None,
                output_schema: None,
                license: "Llama 2".to_string(),
                provider: "LLaVA".to_string(),
                version: "1.5".to_string(),
                is_public: true,
                owner: None,
            },
        ];

        let mut models = self.models.write();
        let mut aliases = self.aliases.write();

        for model in default_models {
            let alias = model.id.0.clone();
            let cid = model.cid.clone();

            models.insert(alias.clone(), model.clone());
            models.insert(cid.clone(), model);
            aliases.insert(alias, cid);
        }
    }

    /// Resolves a model ID (alias or CID) to model info.
    pub fn resolve(&self, id: &str) -> Result<ModelInfo, ComputeError> {
        let models = self.models.read();

        // Try direct lookup
        if let Some(model) = models.get(id) {
            return Ok(model.clone());
        }

        // Try alias lookup
        let aliases = self.aliases.read();
        if let Some(cid) = aliases.get(id) {
            if let Some(model) = models.get(cid) {
                return Ok(model.clone());
            }
        }

        Err(ComputeError::ModelNotFound(id.to_string()))
    }

    /// Registers a custom model.
    pub fn register(&self, model: ModelInfo) -> Result<(), ComputeError> {
        let mut models = self.models.write();
        let mut aliases = self.aliases.write();

        let id = model.id.0.clone();
        let cid = model.cid.clone();

        // Register by both ID and CID
        models.insert(id.clone(), model.clone());
        models.insert(cid.clone(), model);

        if !id.starts_with("Qm") && !id.starts_with("bafy") {
            aliases.insert(id, cid);
        }

        Ok(())
    }

    /// Lists all available models.
    pub fn list(&self) -> Vec<ModelInfo> {
        let models = self.models.read();
        let aliases = self.aliases.read();

        // Return only alias entries to avoid duplicates
        aliases
            .keys()
            .filter_map(|alias| models.get(alias).cloned())
            .collect()
    }

    /// Lists models by category.
    pub fn list_by_category(&self, category: ModelCategory) -> Vec<ModelInfo> {
        self.list()
            .into_iter()
            .filter(|m| m.category == category)
            .collect()
    }

    /// Searches models by name/description.
    pub fn search(&self, query: &str) -> Vec<ModelInfo> {
        let query_lower = query.to_lowercase();
        self.list()
            .into_iter()
            .filter(|m| {
                m.name.to_lowercase().contains(&query_lower)
                    || m.description.to_lowercase().contains(&query_lower)
                    || m.id.0.to_lowercase().contains(&query_lower)
            })
            .collect()
    }

    /// Gets model by alias.
    pub fn get_by_alias(&self, alias: &str) -> Option<ModelInfo> {
        self.resolve(alias).ok()
    }

    /// Checks if a model exists.
    pub fn exists(&self, id: &str) -> bool {
        self.resolve(id).is_ok()
    }
}

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

/// Model upload request.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ModelUploadRequest {
    /// Model name.
    pub name: String,
    /// Model description.
    pub description: Option<String>,
    /// Model category.
    pub category: ModelCategory,
    /// Model format.
    pub format: ModelFormat,
    /// Model file data (bytes).
    #[serde(skip)]
    pub data: Vec<u8>,
    /// Optional alias (must be unique).
    pub alias: Option<String>,
    /// Is public.
    pub is_public: bool,
    /// License.
    pub license: Option<String>,
}

/// Model upload response.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ModelUploadResponse {
    /// Assigned model ID.
    pub model_id: ModelId,
    /// Storage CID.
    pub cid: String,
    /// Size in bytes.
    pub size_bytes: u64,
}

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

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

        // Test default models exist
        let llama = registry.resolve("llama-3-70b").unwrap();
        assert_eq!(llama.parameters, 70_000_000_000);
        assert_eq!(llama.category, ModelCategory::Llm);
    }

    #[test]
    fn test_model_list() {
        let registry = ModelRegistry::new();
        let models = registry.list();
        assert!(!models.is_empty());
    }

    #[test]
    fn test_model_search() {
        let registry = ModelRegistry::new();
        let results = registry.search("llama");
        assert!(!results.is_empty());
        assert!(results.iter().all(|m| m.name.to_lowercase().contains("llama")));
    }

    #[test]
    fn test_model_by_category() {
        let registry = ModelRegistry::new();
        let llms = registry.list_by_category(ModelCategory::Llm);
        assert!(!llms.is_empty());
        assert!(llms.iter().all(|m| m.category == ModelCategory::Llm));
    }

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

        let custom = ModelInfo::custom(
            "QmCustomModel123",
            "My Custom Model",
            ModelCategory::Custom,
            ModelFormat::Onnx,
            1_000_000,
            [1u8; 32],
        );

        registry.register(custom).unwrap();

        let resolved = registry.resolve("QmCustomModel123").unwrap();
        assert_eq!(resolved.name, "My Custom Model");
    }
}
