from fastapi import APIRouter, HTTPException, UploadFile, File
from pydantic import BaseModel
from typing import Dict, Any, List, Optional
import numpy as np
import onnxruntime as ort
import json
import logging

logger = logging.getLogger(__name__)

router = APIRouter()

# In-memory model storage (in production, use proper model registry)
loaded_models = {}

class PredictionRequest(BaseModel):
    model_id: str
    features: Dict[str, float]

class PredictionResponse(BaseModel):
    model_id: str
    prediction: float
    probability: Optional[Dict[str, float]] = None
    metadata: Optional[Dict[str, Any]] = None

class ModelInfo(BaseModel):
    model_id: str
    name: str
    version: str
    type: str
    input_features: List[str]
    output_shape: List[int]
    metadata: Dict[str, Any]

@router.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    """
    Run inference on a loaded model
    """
    try:
        if request.model_id not in loaded_models:
            raise HTTPException(status_code=404, detail=f"Model {request.model_id} not found")
        
        model_info = loaded_models[request.model_id]
        session = model_info['session']
        
        # Prepare input
        input_features = model_info['input_features']
        input_array = np.array([[request.features.get(f, 0.0) for f in input_features]], dtype=np.float32)
        
        # Run inference
        input_name = session.get_inputs()[0].name
        output = session.run(None, {input_name: input_array})
        
        # Process output
        prediction = float(output[0][0])
        
        # For classification models, get probabilities
        probability = None
        if model_info['type'] == 'classification' and len(output[0][0]) > 1:
            probability = {
                f"class_{i}": float(p) 
                for i, p in enumerate(output[0][0])
            }
        
        return PredictionResponse(
            model_id=request.model_id,
            prediction=prediction,
            probability=probability,
            metadata={
                "model_version": model_info['version'],
                "timestamp": np.datetime64('now').tolist()
            }
        )
        
    except Exception as e:
        logger.error(f"Prediction failed: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

@router.post("/load")
async def load_model(
    model_id: str,
    model_file: UploadFile = File(...),
    metadata: str = None
):
    """
    Load an ONNX model for inference
    """
    try:
        # Read model file
        content = await model_file.read()
        
        # Create ONNX session
        session = ort.InferenceSession(content)
        
        # Parse metadata
        model_metadata = json.loads(metadata) if metadata else {}
        
        # Extract model info
        input_features = [inp.name for inp in session.get_inputs()]
        output_shape = [out.shape for out in session.get_outputs()]
        
        # Store model
        loaded_models[model_id] = {
            'session': session,
            'input_features': model_metadata.get('feature_names', input_features),
            'type': model_metadata.get('model_type', 'regression'),
            'version': model_metadata.get('version', '1.0'),
            'metadata': model_metadata
        }
        
        return {
            "message": f"Model {model_id} loaded successfully",
            "input_features": input_features,
            "output_shape": output_shape
        }
        
    except Exception as e:
        logger.error(f"Failed to load model: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Failed to load model: {str(e)}")

@router.get("/list", response_model=List[ModelInfo])
async def list_models():
    """
    List all loaded models
    """
    models = []
    
    for model_id, info in loaded_models.items():
        session = info['session']
        models.append(ModelInfo(
            model_id=model_id,
            name=info['metadata'].get('name', model_id),
            version=info['version'],
            type=info['type'],
            input_features=info['input_features'],
            output_shape=[out.shape for out in session.get_outputs()],
            metadata=info['metadata']
        ))
    
    return models

@router.delete("/{model_id}")
async def unload_model(model_id: str):
    """
    Unload a model from memory
    """
    if model_id not in loaded_models:
        raise HTTPException(status_code=404, detail=f"Model {model_id} not found")
    
    del loaded_models[model_id]
    
    return {"message": f"Model {model_id} unloaded successfully"}

@router.post("/batch_predict")
async def batch_predict(
    model_id: str,
    features: List[Dict[str, float]]
):
    """
    Run batch predictions
    """
    try:
        if model_id not in loaded_models:
            raise HTTPException(status_code=404, detail=f"Model {model_id} not found")
        
        predictions = []
        
        for feature_set in features:
            request = PredictionRequest(model_id=model_id, features=feature_set)
            result = await predict(request)
            predictions.append(result.dict())
        
        return {
            "model_id": model_id,
            "predictions": predictions,
            "count": len(predictions)
        }
        
    except Exception as e:
        logger.error(f"Batch prediction failed: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Batch prediction failed: {str(e)}")