"""
API Generator

Automatic REST API generation for trained models using FastAPI with
intelligent endpoint creation, input validation, and comprehensive documentation.
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Union, Type
from pydantic import BaseModel, Field, create_model
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import joblib
import json
import inspect
from datetime import datetime
import hashlib
from pathlib import Path
import warnings


class PredictionRequest(BaseModel):
    """Base model for prediction requests"""
    data: Dict[str, Any] = Field(..., description="Input data for prediction")
    return_proba: Optional[bool] = Field(False, description="Return prediction probabilities")


class BatchPredictionRequest(BaseModel):
    """Model for batch prediction requests"""
    data: List[Dict[str, Any]] = Field(..., description="List of input data for batch prediction")
    return_proba: Optional[bool] = Field(False, description="Return prediction probabilities")


class ModelInfo(BaseModel):
    """Model information response"""
    model_id: str
    model_type: str
    problem_type: str
    features: List[str]
    target: Optional[str]
    performance_metrics: Dict[str, float]
    created_at: str
    version: str


class PredictionResponse(BaseModel):
    """Prediction response"""
    predictions: List[Any]
    probabilities: Optional[List[List[float]]] = None
    model_id: str
    timestamp: str


class HealthResponse(BaseModel):
    """Health check response"""
    status: str
    timestamp: str
    model_loaded: bool
    version: str


class APIGenerator:
    """
    Automatic REST API generator for ML models using FastAPI.
    Creates comprehensive APIs with validation, documentation, and monitoring.
    """
    
    def __init__(self, model_path: str, model_id: Optional[str] = None):
        """
        Initialize API generator
        
        Args:
            model_path: Path to the saved model file
            model_id: Unique identifier for the model
        """
        self.model_path = Path(model_path)
        self.model_id = model_id or self.model_path.stem
        self.model = None
        self.feature_names = []
        self.target_name = None
        self.problem_type = None
        self.performance_metrics = {}
        self.app = None
        self.request_models = {}
        
        # Load model and metadata
        self._load_model()
    
    def _load_model(self):
        """Load model and extract metadata"""
        try:
            self.model = joblib.load(self.model_path)
        except Exception as e:
            raise ValueError(f"Failed to load model from {self.model_path}: {e}")
        
        # Try to load metadata
        metadata_path = self.model_path.parent / f"{self.model_path.stem}_metadata.json"
        if metadata_path.exists():
            with open(metadata_path, 'r') as f:
                metadata = json.load(f)
                self.feature_names = metadata.get('feature_names', [])
                self.target_name = metadata.get('target_name')
                self.problem_type = metadata.get('problem_type', 'unknown')
                self.performance_metrics = metadata.get('performance_metrics', {})
    
    def generate_api(self, title: str = "ML Model API", 
                    description: str = "Auto-generated API for ML model",
                    version: str = "1.0.0") -> FastAPI:
        """
        Generate FastAPI application for the model
        
        Args:
            title: API title
            description: API description
            version: API version
            
        Returns:
            FastAPI application instance
        """
        # Create FastAPI app
        self.app = FastAPI(
            title=title,
            description=description,
            version=version,
            docs_url="/docs",
            redoc_url="/redoc"
        )
        
        # Add CORS middleware
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )
        
        # Create dynamic request models based on features
        self._create_request_models()
        
        # Add endpoints
        self._add_endpoints()
        
        return self.app
    
    def _create_request_models(self):
        """Create dynamic Pydantic models for input validation"""
        if not self.feature_names:
            # Try to infer feature names from model
            if hasattr(self.model, 'feature_names_in_'):
                self.feature_names = list(self.model.feature_names_in_)
            elif hasattr(self.model, 'n_features_in_'):
                self.feature_names = [f"feature_{i}" for i in range(self.model.n_features_in_)]
        
        if not self.feature_names:
            # Default generic model
            self.request_models['single'] = PredictionRequest
            self.request_models['batch'] = BatchPredictionRequest
            return
        
        # Create single prediction model with specific fields
        fields = {}
        for feature in self.feature_names:
            fields[feature] = (Any, Field(..., description=f"Feature: {feature}"))
        
        # Add optional return_proba field
        fields['return_proba'] = (Optional[bool], Field(False, description="Return prediction probabilities"))
        
        SinglePredictionModel = create_model(
            'SinglePrediction',
            **fields,
            __base__=BaseModel
        )
        
        # Create batch prediction model
        BatchPredictionModel = create_model(
            'BatchPrediction',
            data=(List[Dict[str, Any]], Field(..., description="List of input data")),
            return_proba=(Optional[bool], Field(False, description="Return prediction probabilities")),
            __base__=BaseModel
        )
        
        self.request_models['single'] = SinglePredictionModel
        self.request_models['batch'] = BatchPredictionModel
    
    def _add_endpoints(self):
        """Add API endpoints"""
        
        @self.app.get("/", response_model=Dict[str, str])
        async def root():
            """Root endpoint"""
            return {
                "message": f"ML Model API for {self.model_id}",
                "docs": "/docs",
                "health": "/health"
            }
        
        @self.app.get("/health", response_model=HealthResponse)
        async def health_check():
            """Health check endpoint"""
            return HealthResponse(
                status="healthy",
                timestamp=datetime.now().isoformat(),
                model_loaded=self.model is not None,
                version="1.0.0"
            )
        
        @self.app.get("/info", response_model=ModelInfo)
        async def get_model_info():
            """Get model information"""
            return ModelInfo(
                model_id=self.model_id,
                model_type=type(self.model).__name__,
                problem_type=self.problem_type or "unknown",
                features=self.feature_names,
                target=self.target_name,
                performance_metrics=self.performance_metrics,
                created_at=datetime.now().isoformat(),
                version="1.0.0"
            )
        
        @self.app.post("/predict", response_model=PredictionResponse)
        async def predict(request: PredictionRequest):
            """Make single prediction"""
            try:
                # Convert input to DataFrame
                input_data = pd.DataFrame([request.data])
                
                # Ensure all required features are present
                if self.feature_names:
                    missing_features = set(self.feature_names) - set(input_data.columns)
                    if missing_features:
                        raise HTTPException(
                            status_code=400,
                            detail=f"Missing required features: {missing_features}"
                        )
                    
                    # Reorder columns to match training data
                    input_data = input_data[self.feature_names]
                
                # Make prediction
                prediction = self.model.predict(input_data)
                
                # Get probabilities if requested and available
                probabilities = None
                if request.return_proba and hasattr(self.model, 'predict_proba'):
                    probabilities = self.model.predict_proba(input_data)
                
                return PredictionResponse(
                    predictions=prediction.tolist(),
                    probabilities=probabilities.tolist() if probabilities is not None else None,
                    model_id=self.model_id,
                    timestamp=datetime.now().isoformat()
                )
                
            except Exception as e:
                raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")
        
        @self.app.post("/predict/batch", response_model=PredictionResponse)
        async def predict_batch(request: BatchPredictionRequest):
            """Make batch predictions"""
            try:
                # Convert input to DataFrame
                input_data = pd.DataFrame(request.data)
                
                # Ensure all required features are present
                if self.feature_names:
                    missing_features = set(self.feature_names) - set(input_data.columns)
                    if missing_features:
                        raise HTTPException(
                            status_code=400,
                            detail=f"Missing required features: {missing_features}"
                        )
                    
                    # Reorder columns to match training data
                    input_data = input_data[self.feature_names]
                
                # Make predictions
                predictions = self.model.predict(input_data)
                
                # Get probabilities if requested and available
                probabilities = None
                if request.return_proba and hasattr(self.model, 'predict_proba'):
                    probabilities = self.model.predict_proba(input_data)
                
                return PredictionResponse(
                    predictions=predictions.tolist(),
                    probabilities=probabilities.tolist() if probabilities is not None else None,
                    model_id=self.model_id,
                    timestamp=datetime.now().isoformat()
                )
                
            except Exception as e:
                raise HTTPException(status_code=500, detail=f"Batch prediction failed: {str(e)}")
        
        @self.app.get("/features")
        async def get_features():
            """Get model feature information"""
            return {
                "features": self.feature_names,
                "n_features": len(self.feature_names),
                "target": self.target_name,
                "problem_type": self.problem_type
            }
        
        @self.app.get("/metrics")
        async def get_metrics():
            """Get model performance metrics"""
            return {
                "model_id": self.model_id,
                "metrics": self.performance_metrics,
                "timestamp": datetime.now().isoformat()
            }
    
    def add_custom_endpoint(self, path: str, methods: List[str], 
                          handler_func: callable, **kwargs):
        """Add custom endpoint to the API"""
        if not self.app:
            raise ValueError("API not generated yet. Call generate_api() first.")
        
        # Add the endpoint
        self.app.add_api_route(
            path=path,
            endpoint=handler_func,
            methods=methods,
            **kwargs
        )
    
    def add_monitoring_middleware(self):
        """Add monitoring middleware for request tracking"""
        @self.app.middleware("http")
        async def log_requests(request, call_next):
            start_time = datetime.now()
            
            response = await call_next(request)
            
            process_time = (datetime.now() - start_time).total_seconds()
            
            # Add custom headers
            response.headers["X-Process-Time"] = str(process_time)
            response.headers["X-Model-ID"] = self.model_id
            
            return response
    
    def enable_rate_limiting(self, requests_per_minute: int = 60):
        """Add rate limiting middleware"""
        # This is a simplified version - in production, use Redis or similar
        request_counts = {}
        
        @self.app.middleware("http")
        async def rate_limit_middleware(request, call_next):
            client_ip = request.client.host
            current_time = datetime.now()
            
            # Clean old entries
            cutoff_time = current_time - pd.Timedelta(minutes=1)
            if client_ip in request_counts:
                request_counts[client_ip] = [
                    req_time for req_time in request_counts[client_ip]
                    if req_time > cutoff_time
                ]
            else:
                request_counts[client_ip] = []
            
            # Check rate limit
            if len(request_counts[client_ip]) >= requests_per_minute:
                return JSONResponse(
                    status_code=429,
                    content={"detail": "Rate limit exceeded"}
                )
            
            # Add current request
            request_counts[client_ip].append(current_time)
            
            response = await call_next(request)
            return response
    
    def add_authentication(self, api_key: str):
        """Add API key authentication"""
        @self.app.middleware("http")
        async def auth_middleware(request, call_next):
            # Skip auth for health and docs endpoints
            if request.url.path in ["/", "/health", "/docs", "/redoc", "/openapi.json"]:
                response = await call_next(request)
                return response
            
            # Check API key
            if "x-api-key" not in request.headers:
                return JSONResponse(
                    status_code=401,
                    content={"detail": "API key required"}
                )
            
            if request.headers["x-api-key"] != api_key:
                return JSONResponse(
                    status_code=401,
                    content={"detail": "Invalid API key"}
                )
            
            response = await call_next(request)
            return response
    
    def generate_openapi_spec(self) -> Dict[str, Any]:
        """Generate OpenAPI specification"""
        if not self.app:
            raise ValueError("API not generated yet. Call generate_api() first.")
        
        return self.app.openapi()
    
    def save_api_spec(self, output_path: str):
        """Save OpenAPI specification to file"""
        spec = self.generate_openapi_spec()
        
        with open(output_path, 'w') as f:
            json.dump(spec, f, indent=2)
    
    def generate_dockerfile(self, output_path: str = "Dockerfile"):
        """Generate Dockerfile for API deployment"""
        dockerfile_content = f"""FROM python:3.9-slim

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y \\
    gcc \\
    && rm -rf /var/lib/apt/lists/*

# Copy requirements first for better caching
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy application code
COPY . .

# Expose port
EXPOSE 8000

# Run the application
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
"""
        
        with open(output_path, 'w') as f:
            f.write(dockerfile_content)
    
    def generate_requirements(self, output_path: str = "requirements.txt"):
        """Generate requirements.txt for API deployment"""
        requirements = [
            "fastapi>=0.68.0",
            "uvicorn[standard]>=0.15.0",
            "pydantic>=1.8.0",
            "scikit-learn>=1.0.0",
            "pandas>=1.3.0",
            "numpy>=1.21.0",
            "joblib>=1.0.0",
            "python-multipart>=0.0.5"
        ]
        
        with open(output_path, 'w') as f:
            f.write('\n'.join(requirements))
    
    def generate_main_script(self, output_path: str = "main.py"):
        """Generate main.py script for running the API"""
        main_script_content = f'''from fastapi import FastAPI
import uvicorn
from api_generator import APIGenerator

# Initialize API generator
api_generator = APIGenerator(
    model_path="{self.model_path}",
    model_id="{self.model_id}"
)

# Generate FastAPI app
app = api_generator.generate_api(
    title="ML Model API",
    description="Auto-generated API for ML model",
    version="1.0.0"
)

# Add monitoring
api_generator.add_monitoring_middleware()

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
'''
        
        with open(output_path, 'w') as f:
            f.write(main_script_content)
    
    def test_api_locally(self, port: int = 8000):
        """Test the API locally"""
        import uvicorn
        
        if not self.app:
            raise ValueError("API not generated yet. Call generate_api() first.")
        
        print(f"Starting API server on http://localhost:{port}")
        print(f"API documentation available at http://localhost:{port}/docs")
        
        uvicorn.run(self.app, host="0.0.0.0", port=port)
    
    def create_test_client(self):
        """Create test client for the API"""
        from fastapi.testclient import TestClient
        
        if not self.app:
            raise ValueError("API not generated yet. Call generate_api() first.")
        
        return TestClient(self.app)
    
    def generate_test_suite(self, output_path: str = "test_api.py"):
        """Generate test suite for the API"""
        test_content = f'''import pytest
from fastapi.testclient import TestClient
from main import app

client = TestClient(app)

def test_root():
    response = client.get("/")
    assert response.status_code == 200
    assert "ML Model API" in response.json()["message"]

def test_health():
    response = client.get("/health")
    assert response.status_code == 200
    assert response.json()["status"] == "healthy"

def test_info():
    response = client.get("/info")
    assert response.status_code == 200
    assert "model_id" in response.json()

def test_predict():
    # Test prediction endpoint
    test_data = {{
        "data": {{
            {', '.join([f'"{feature}": 1.0' for feature in self.feature_names[:5]])}
        }},
        "return_proba": False
    }}
    
    response = client.post("/predict", json=test_data)
    assert response.status_code == 200
    assert "predictions" in response.json()

def test_predict_invalid_data():
    # Test with invalid data
    test_data = {{
        "data": {{}},
        "return_proba": False
    }}
    
    response = client.post("/predict", json=test_data)
    assert response.status_code == 400

def test_features():
    response = client.get("/features")
    assert response.status_code == 200
    assert "features" in response.json()

def test_metrics():
    response = client.get("/metrics")
    assert response.status_code == 200
    assert "metrics" in response.json()
'''
        
        with open(output_path, 'w') as f:
            f.write(test_content)
    
    def get_api_summary(self) -> Dict[str, Any]:
        """Get comprehensive API summary"""
        if not self.app:
            return {"error": "API not generated yet"}
        
        return {
            "model_id": self.model_id,
            "model_type": type(self.model).__name__,
            "problem_type": self.problem_type,
            "features": self.feature_names,
            "target": self.target_name,
            "performance_metrics": self.performance_metrics,
            "endpoints": [
                {"path": route.path, "methods": list(route.methods)}
                for route in self.app.routes
                if hasattr(route, 'path')
            ],
            "model_path": str(self.model_path),
            "generated_at": datetime.now().isoformat()
        }
