# handler.py - Hugging Face Inference Endpoints Compatible Version
import time
import threading
import logging
import importlib
import gc
import os
from typing import Dict, Any, Optional
import json
from model_manager import discover_models_in_src, ModelManager
from download_models import download_models

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# Global model manager
_model_manager = None


class EndpointHandler:
    """
    Hugging Face Inference Endpoints compatible handler for multiple models

    This class implements the required interface for HF Inference Endpoints:
    - __init__(self, path=".")
    - __call__(self, data)
    """

    def __init__(self, path="."):
        """Initialize the endpoint handler"""
        global _model_manager

        logger.info(f"Initializing EndpointHandler with path: {path}")
        download_models()

        # Initialize model manager
        _model_manager = ModelManager()

        # Change to the model directory
        original_cwd = os.getcwd()
        if path != ".":
            os.chdir(path)

        try:
            # Discover and register models
            discovered_models = discover_models_in_src()

            for name, module in discovered_models.items():
                _model_manager.register_model_module(name, module)

            # Load models at startup (optional - you can also load on-demand)
            self.preload_models = os.getenv("PRELOAD_MODELS", "true").lower() == "true"
            if self.preload_models and discovered_models:
                logger.info("Preloading all models...")
                for model_name in discovered_models.keys():
                    _model_manager.load_model(model_name)

            logger.info(f"Handler initialized with {len(discovered_models)} models")

        finally:
            # Restore original working directory
            os.chdir(original_cwd)

    def __call__(self, data):
        """
        Handle inference requests

        Expected input formats:
        1. Single model: {"model": "resnet", "inputs": {...}}
        2. Multiple models: {"models": [{"name": "resnet", "inputs": {...}}, {"name": "yolo", "inputs": {...}}]}
        3. Direct inputs (uses default model): {"input_tensor": [...]}
        """
        try:
            # Handle different input formats
            if isinstance(data, dict):

                # Format 1: Single model specified
                if "model" in data and "inputs" in data:
                    model_name = data["model"]
                    inputs = data["inputs"]
                    result = _model_manager.predict(model_name, **inputs)
                    return {"model": model_name, "prediction": result}

                # Format 2: Multiple models
                elif "models" in data:
                    results = []
                    for model_request in data["models"]:
                        model_name = model_request["name"]
                        inputs = model_request.get("inputs", {})
                        result = _model_manager.predict(model_name, **inputs)
                        results.append({"model": model_name, "prediction": result})
                    return {"results": results}

                # Format 3: Direct inputs - try to infer model or use default
                else:
                    # Try to use the first available model as default
                    available_models = list(_model_manager.model_modules.keys())
                    if not available_models:
                        raise ValueError("No models available")

                    default_model = available_models[0]
                    logger.info(f"Using default model: {default_model}")

                    result = _model_manager.predict(default_model, **data)
                    return {"model": default_model, "prediction": result}

            else:
                raise ValueError("Invalid input format. Expected dictionary.")

        except Exception as e:
            logger.error(f"Inference error: {e}")
            return {"error": str(e)}


# Standalone functions for compatibility (if needed)
def load_models(model_names: Optional[list] = None) -> Dict[str, bool]:
    """Load specified models or all models if none specified"""
    if _model_manager is None:
        raise RuntimeError("Handler not initialized")

    if model_names is None:
        results = {}
        for model_name in _model_manager.model_modules.keys():
            results[model_name] = _model_manager.load_model(model_name)
        return results

    results = {}
    for name in model_names:
        results[name] = _model_manager.load_model(name)
    return results


def predict(model_name: str, **kwargs) -> Any:
    """Make prediction using specified model"""
    if _model_manager is None:
        raise RuntimeError("Handler not initialized")
    return _model_manager.predict(model_name, **kwargs)