"""
Model configuration for the unified model manager
Defines default models, fallback options, and model type mappings
"""

# Default models for different tasks
DEFAULT_MODELS = {
    "text-generation": {
        "primary": "facebook/bart-base",
        "fallback": "facebook/bart-large-cnn",
        "description": "Text generation models for QA and medical data extraction"
    },
    "summarization": {
        "primary": "Falconsai/medical_summarization",
        "fallback": "facebook/bart-large-cnn",
        "description": "Text summarization models for medical reports"
    },
    "ner": {
        "primary": "dslim/bert-base-NER",
        "fallback": "dslim/bert-base-NER",
        "description": "Named Entity Recognition for medical entities"
    },
    "gguf": {
        "primary": "microsoft/Phi-3-mini-4k-instruct-gguf",
        "fallback": "microsoft/Phi-3-mini-4k-instruct-gguf",
        "description": "GGUF models for patient summaries and medical tasks"
    },
    "openvino": {
        "primary": "microsoft/Phi-3-mini-4k-instruct",
        "fallback": "microsoft/Phi-3-mini-4k-instruct",
        "description": "OpenVINO optimized models"
    }
}

# Model type mappings for automatic detection
MODEL_TYPE_MAPPINGS = {
    # GGUF models
    ".gguf": "gguf",
    "gguf": "gguf",
    
    # OpenVINO models
    "openvino": "openvino",
    "ov": "openvino",
    
    # Transformers models
    "text-generation": "text-generation",
    "summarization": "summarization",
    "ner": "ner",
    "question-answering": "text-generation",
    "translation": "text-generation"
}

# Memory-optimized models for Hugging Face Spaces
SPACES_OPTIMIZED_MODELS = {
    "text-generation": "facebook/bart-base",
    "summarization": "Falconsai/medical_summarization",
    "ner": "dslim/bert-base-NER",
    "gguf": "microsoft/Phi-3-mini-4k-instruct-gguf"
}

# Model validation rules
MODEL_VALIDATION_RULES = {
    "text-generation": {
        "min_tokens": 100,
        "max_tokens": 4000,
        "supported_formats": ["huggingface", "local"]
    },
    "summarization": {
        "min_tokens": 50,
        "max_tokens": 4096,
        "supported_formats": ["huggingface", "local"]
    },
    "ner": {
        "min_tokens": 50,
        "max_tokens": 512,
        "supported_formats": ["huggingface", "local"]
    },
    "gguf": {
        "min_tokens": 100,
        "max_tokens": 4096,
        "supported_formats": ["huggingface", "local", "remote"]
    },
    "openvino": {
        "min_tokens": 100,
        "max_tokens": 4000,
        "supported_formats": ["huggingface", "local"]
    }
}

def get_default_model(model_type: str, use_spaces_optimized: bool = False) -> str:
    """Get the default model for a given type"""
    if use_spaces_optimized and model_type in SPACES_OPTIMIZED_MODELS:
        return SPACES_OPTIMIZED_MODELS[model_type]
    
    if model_type in DEFAULT_MODELS:
        return DEFAULT_MODELS[model_type]["primary"]
    
    # Fallback to text-generation if type not found
    return DEFAULT_MODELS["text-generation"]["primary"]

def get_fallback_model(model_type: str) -> str:
    """Get the fallback model for a given type"""
    if model_type in DEFAULT_MODELS:
        return DEFAULT_MODELS[model_type]["fallback"]
    
    return DEFAULT_MODELS["text-generation"]["fallback"]

def detect_model_type(model_name: str) -> str:
    """Automatically detect model type from model name"""
    model_name_lower = model_name.lower()
    
    # Check for explicit type indicators
    for indicator, model_type in MODEL_TYPE_MAPPINGS.items():
        if indicator in model_name_lower:
            return model_type
    
    # Check file extensions
    if model_name.endswith('.gguf'):
        return "gguf"
    
    # Default to text-generation for unknown types
    return "text-generation"

def validate_model_config(model_name: str, model_type: str) -> dict:
    """Validate model configuration and return validation result"""
    result = {
        "valid": True,
        "warnings": [],
        "errors": [],
        "recommendations": []
    }
    
    # Check if model type is supported
    if model_type not in MODEL_VALIDATION_RULES:
        result["valid"] = False
        result["errors"].append(f"Unsupported model type: {model_type}")
        return result
    
    # Check model name format
    if model_type == "gguf":
        if not (model_name.endswith('.gguf') or '/' in model_name):
            result["warnings"].append("GGUF model should have .gguf extension or be in repo/filename format")
    
    # Check for memory optimization recommendations
    if model_type in ["text-generation", "summarization"]:
        if "large" in model_name.lower() or "xl" in model_name.lower():
            result["warnings"].append("Large models may cause memory issues on limited resources")
            result["recommendations"].append("Consider using a smaller model for better performance")
    
    return result

def get_model_info(model_name: str, model_type: str) -> dict:
    """Get comprehensive information about a model configuration"""
    validation = validate_model_config(model_name, model_type)
    
    return {
        "model_name": model_name,
        "model_type": model_type,
        "detected_type": detect_model_type(model_name),
        "default_model": get_default_model(model_type),
        "fallback_model": get_fallback_model(model_type),
        "validation": validation,
        "supported_formats": MODEL_VALIDATION_RULES.get(model_type, {}).get("supported_formats", []),
        "description": DEFAULT_MODELS.get(model_type, {}).get("description", "Unknown model type")
    }