#!/usr/bin/env python3
"""
Model Gate — Icarus Staging Model Restriction
Enforces 3B parameter limit on staging environment
"""

import os
import sys
from typing import List, Optional

# Allowed models for Icarus staging (3B parameters or smaller)
ALLOWED_MODELS: List[str] = [
    "qwen2.5:3b",
    "llama3.2:3b",
    "phi3:mini",
    "phi3:3.8b",
]

# Blocked model patterns (anything larger than 4B)
BLOCKED_PATTERNS = [":7b", ":8b", ":13b", ":70b", ":34b", "-coder:", "-instruct:14b"]


def is_model_allowed(model_name: str) -> bool:
    """Check if model is allowed for Icarus staging."""
    model_lower = model_name.lower()
    
    # Exact match whitelist
    if model_lower in [m.lower() for m in ALLOWED_MODELS]:
        return True
    
    # Pattern blacklist
    for pattern in BLOCKED_PATTERNS:
        if pattern in model_lower:
            return False
    
    # Default allow for unknown small models (e.g., custom 3B models)
    # But require explicit approval for anything suspicious
    if any(size in model_lower for size in [":1b", ":2b", ":3b", ":4b"]):
        return True
    
    return False


def get_allowed_model(preferred: Optional[str] = None) -> str:
    """
    Get model for Icarus staging.
    
    If preferred is specified and allowed, use it.
    Otherwise fall back to first allowed model.
    """
    if preferred and is_model_allowed(preferred):
        return preferred
    
    if preferred and not is_model_allowed(preferred):
        print(f"⚠️  Model '{preferred}' blocked in staging. Using fallback.", file=sys.stderr)
    
    return ALLOWED_MODELS[0]  # Default to qwen2.5:3b


def validate_ollama_request(model: str) -> None:
    """
    Validate and raise if model is not allowed.
    Call this before any Ollama API call in Icarus staging.
    """
    if os.environ.get("ICARUS_ENV") != "staging":
        return  # Only enforce in staging
    
    if not is_model_allowed(model):
        raise ValueError(
            f"🚫 Model '{model}' is BLOCKED in Icarus staging.\n"
            f"Allowed models: {', '.join(ALLOWED_MODELS)}\n"
            f"Rationale: Staging simulates M4 Mac mini (3B target)"
        )


if __name__ == "__main__":
    # CLI test
    test_models = sys.argv[1:] if len(sys.argv) > 1 else [
        "qwen2.5:3b",
        "llama3.2:3b", 
        "qwen2.5-coder:7b",
        "llama3.1:8b",
        "unknown:5b"
    ]
    
    print("Icarus Staging Model Gate")
    print("=" * 40)
    for model in test_models:
        allowed = "✅ ALLOWED" if is_model_allowed(model) else "🚫 BLOCKED"
        print(f"{model:20} {allowed}")