from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
import pandas as pd
import numpy as np
import joblib
import logging
import os
import sys
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.ensemble import IsolationForest

app = FastAPI(title="Isolation Forest Anomaly Detection")

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

# Log environment info for debugging
logger.info(f"Python version: {sys.version}")
logger.info(f"Current working directory: {os.getcwd()}")
logger.info(f"Directory contents: {os.listdir('.')}")

# Model paths - make them more flexible for different environments
MODEL_PATH = os.getenv("MODEL_PATH", "./isoforest_dos.pkl")
SCALER_PATH = os.getenv("SCALER_PATH", "./scaler_dos.pkl")
ENCODER_PATH = os.getenv("ENCODER_PATH", "./encoder_dos.pkl")

THRESHOLD = 0.4153
NUM_FEATS = ["inter_arrival_time", "packet_rate", "packet_length", "length_per_rate", 
             "packet_rate_mean", "packet_rate_var", "packet_rate_skew"]
CAT_FEATS = ["protocol"]
FEATURES = NUM_FEATS + CAT_FEATS

# Global variables to store loaded models
isoforest = None
scaler = None
encoder = None
models_loaded = False

def check_sklearn_versions():
    """Check and log sklearn version compatibility"""
    try:
        import sklearn
        logger.info(f"Scikit-learn version: {sklearn.__version__}")
        import numpy
        logger.info(f"NumPy version: {numpy.__version__}")
        return True
    except Exception as e:
        logger.error(f"Version check failed: {e}")
        return False

def load_models():
    """Load models with better error handling and fallback for encoder"""
    global isoforest, scaler, encoder, models_loaded
    
    # Check versions first
    if not check_sklearn_versions():
        return False
    
    try:
        # Check if files exist and log their details
        for path, name in [(MODEL_PATH, "model"), (SCALER_PATH, "scaler")]:
            if not os.path.exists(path):
                logger.error(f"{name} file not found at {path}")
                logger.info(f"Available files: {[f for f in os.listdir('.') if f.endswith('.pkl')]}")
                return False
            else:
                file_size = os.path.getsize(path)
                logger.info(f"{name} file found at {path} (size: {file_size} bytes)")
        
        # Try to load models with more specific error handling
        logger.info("Loading isolation forest model...")
        isoforest = joblib.load(MODEL_PATH)
        logger.info("✓ Isolation forest loaded")
        
        logger.info("Loading scaler...")
        scaler = joblib.load(SCALER_PATH)
        logger.info("✓ Scaler loaded")
        
        # Try to load encoder, but use fallback if it fails
        logger.info("Loading encoder...")
        try:
            if os.path.exists(ENCODER_PATH):
                encoder = joblib.load(ENCODER_PATH)
                logger.info("✓ Encoder loaded")
            else:
                logger.warning("Encoder file not found, will use fallback encoding")
                encoder = None
        except Exception as e:
            logger.warning(f"Failed to load encoder: {str(e)}. Will use fallback encoding")
            encoder = None
        
        models_loaded = True
        logger.info("Models loaded successfully")
        return True
        
    except ImportError as e:
        logger.error(f"Import error while loading models: {str(e)}")
        logger.error("This might be a version compatibility issue")
        return False
    except Exception as e:
        logger.error(f"Failed to load model or preprocessors: {str(e)}")
        logger.error(f"Error type: {type(e).__name__}")
        return False

# Add startup event
@app.on_event("startup")
async def startup_event():
    """Load models on startup"""
    global models_loaded
    logger.info("Starting model loading...")
    models_loaded = load_models()
    if models_loaded:
        logger.info("✓ Startup complete - models loaded successfully")
    else:
        logger.error("✗ Startup failed - models not loaded")

class NetworkData(BaseModel):
    inter_arrival_time: float
    packet_length: float
    protocol: str

class PredictionResponse(BaseModel):
    anomaly: int
    anomaly_score: float

@app.get("/")
async def root():
    """Root endpoint"""
    return {"message": "Isolation Forest Anomaly Detection API", "status": "running"}

@app.get("/health")
async def health_check():
    """Health check endpoint with more details"""
    if not models_loaded or isoforest is None or scaler is None:
        return {
            "status": "unhealthy", 
            "reason": "Critical models not loaded",
            "models_loaded": models_loaded,
            "isoforest_loaded": isoforest is not None,
            "scaler_loaded": scaler is not None,
            "encoder_loaded": encoder is not None
        }
    return {"status": "healthy", "models_loaded": True}

@app.get("/debug")
async def debug_info():
    """Debug endpoint to check environment"""
    import sklearn
    import numpy
    return {
        "sklearn_version": sklearn.__version__,
        "numpy_version": numpy.__version__,
        "working_directory": os.getcwd(),
        "files": os.listdir('.'),
        "pkl_files": [f for f in os.listdir('.') if f.endswith('.pkl')],
        "models_loaded": models_loaded,
        "isoforest_loaded": isoforest is not None,
        "scaler_loaded": scaler is not None,
        "encoder_loaded": encoder is not None
    }

def safe_clip_and_log(series, lower=None, upper=None):
    """Safely clip and apply log1p transformation"""
    if lower is not None:
        series = series.clip(lower=lower)
    if upper is not None:
        # Calculate quantile safely
        try:
            upper_val = series.quantile(0.98) if upper == "quantile_98" else upper
            series = series.clip(upper=upper_val)
        except:
            pass  # If quantile calculation fails, skip upper clipping
    return np.log1p(series)

def fallback_encode_protocol(protocols):
    """Fallback encoding for protocol column"""
    protocol_map = {
        "tcp": 0, 
        "udp": 1, 
        "icmp": 2, 
        "http": 3,
        "https": 4,
        "unknown": 5
    }
    return [protocol_map.get(p.lower(), 5) for p in protocols]

@app.post("/predict", response_model=List[PredictionResponse])
async def predict(data: List[NetworkData]):
    """Predict anomalies in network data"""
    
    # Check if critical models are loaded
    if not models_loaded or isoforest is None or scaler is None:
        raise HTTPException(
            status_code=503, 
            detail="Critical models not loaded. Service unavailable. Check /health for details."
        )
    
    try:
        # Convert input data to DataFrame
        df = pd.DataFrame([d.dict() for d in data])
        original_len = len(df)
        
        # Ensure minimum 5 rows for rolling calculations
        if len(df) < 5:
            padding_rows = 5 - len(df)
            padding_df = pd.DataFrame(
                [[0.001, 64, "tcp"]] * padding_rows, 
                columns=["inter_arrival_time", "packet_length", "protocol"]
            )
            df = pd.concat([padding_df, df], ignore_index=True)
        
        # Feature engineering with better error handling
        df["inter_arrival_time"] = safe_clip_and_log(df["inter_arrival_time"], lower=0.001)
        
        # Calculate packet_rate
        df["packet_rate"] = 1 / np.exp(df["inter_arrival_time"])
        df["packet_rate"] = safe_clip_and_log(df["packet_rate"])
        
        # Process packet_length
        df["packet_length"] = safe_clip_and_log(
            df["packet_length"].clip(lower=0), 
            upper="quantile_98"
        )
        
        # Calculate length_per_rate
        df["length_per_rate"] = np.exp(df["packet_length"]) / np.exp(df["packet_rate"])
        df["length_per_rate"] = safe_clip_and_log(df["length_per_rate"], upper="quantile_98")
        
        # Rolling statistics with better handling
        packet_rate_exp = np.exp(df["packet_rate"])
        
        # Rolling mean
        rolling_mean = packet_rate_exp.rolling(window=5, min_periods=1).mean()
        df["packet_rate_mean"] = safe_clip_and_log(rolling_mean.fillna(packet_rate_exp.median()))
        
        # Rolling variance
        rolling_std = packet_rate_exp.rolling(window=5, min_periods=1).std()
        df["packet_rate_var"] = safe_clip_and_log(
            rolling_std.fillna(packet_rate_exp.std() if packet_rate_exp.std() > 0 else 0.1)
        )
        
        # Rolling skewness
        rolling_skew = packet_rate_exp.rolling(window=5, min_periods=1).skew()
        df["packet_rate_skew"] = rolling_skew.fillna(0)
        # Handle negative skewness for log transformation
        df["packet_rate_skew"] = np.log1p(df["packet_rate_skew"] - df["packet_rate_skew"].min() + 0.001)
        
        # Process categorical features
        df["protocol"] = df["protocol"].astype(str).fillna("unknown")
        
        # Transform categorical features with fallback
        df_encoded = df.copy()
        try:
            if encoder is not None:
                df_encoded[CAT_FEATS] = encoder.transform(df[CAT_FEATS])
                logger.info("Used trained encoder")
            else:
                # Use fallback encoding
                df_encoded["protocol"] = fallback_encode_protocol(df["protocol"])
                logger.info("Used fallback encoding")
        except Exception as e:
            logger.warning(f"Encoding failed: {str(e)}. Using fallback encoding.")
            df_encoded["protocol"] = fallback_encode_protocol(df["protocol"])
        
        # Select features and scale
        X = df_encoded[FEATURES]
        
        # Handle any remaining NaN values
        X = X.fillna(0)
        
        try:
            X_scaled = scaler.transform(X)
        except Exception as e:
            logger.warning(f"Scaling failed: {str(e)}. Using unscaled features.")
            X_scaled = X.values
        
        # Predict anomalies
        try:
            anomaly_scores = -isoforest.score_samples(X_scaled)
            anomalies = (anomaly_scores > THRESHOLD).astype(int)
        except Exception as e:
            logger.error(f"Prediction failed: {str(e)}")
            # Fallback: return all as normal
            anomaly_scores = np.zeros(len(X_scaled))
            anomalies = np.zeros(len(X_scaled), dtype=int)
        
        # Return only the original data predictions (skip padding)
        start_idx = len(df) - original_len
        response = [
            PredictionResponse(anomaly=int(anomaly), anomaly_score=float(score))
            for anomaly, score in zip(anomalies[start_idx:], anomaly_scores[start_idx:])
        ]
        
        logger.info(f"Processed {original_len} records successfully")
        return response
        
    except Exception as e:
        logger.error(f"Prediction error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    port = int(os.getenv("PORT", 8000))
    uvicorn.run(app, host="0.0.0.0", port=port)