"""
Inference service for email threat classification:
Loads serialized ML pipelines, runs feature transformations, generates probabilities,
and assembles security intelligence reports.
"""

import os
import joblib
import pandas as pd
import numpy as np
from typing import Dict, Any

from src.data_loader import CLASS_NAMES, ID_TO_CLASS
from src.feature_engineering import extract_threat_indicators
from src.security_analyzer import assess_threat_level


class EmailThreatPredictor:
    """
    High-performance production predictor for email threat classification.
    """

    def __init__(self, model_path: str = None):
        if model_path is None:
            # Default location
            default_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "models", "artifacts"))
            model_path = os.path.join(default_dir, "best_threat_model.joblib")

        self.model_path = model_path
        self.pipeline = None
        self._load_model()

    def _load_model(self):
        """Loads serialized pipeline with safety checks."""
        if not os.path.exists(self.model_path):
            raise FileNotFoundError(
                f"Model artifact not found at: {self.model_path}. "
                "Please run train.py to build and save the model pipeline."
            )
        self.pipeline = joblib.load(self.model_path)

    def predict(self, subject: str, body: str) -> Dict[str, Any]:
        """
        Executes end-to-end inference on email subject and body.
        Handles edge cases such as empty strings, whitespace, or missing inputs.
        """
        clean_subj = str(subject or "").strip()
        clean_body = str(body or "").strip()

        # Handle purely empty input edge case
        if not clean_subj and not clean_body:
            return {
                "predicted_class": "Legitimate",
                "predicted_label_id": 0,
                "probabilities": {"Legitimate": 0.90, "Spam": 0.05, "Suspicious": 0.05},
                "confidence": 0.90,
                "severity": "LOW RISK",
                "badge_color": "#38a169",
                "risk_score_out_of_10": 0.5,
                "summary": "Empty message contains no malicious tokens or phishing indicators.",
                "indicators": extract_threat_indicators("", ""),
                "recommendations": ["No content detected. Enter email text to analyze."],
                "warning": "Disclaimer: Research prototype designed for defense-in-depth assistance. Machine learning models cannot guarantee 100% email safety."
            }

        input_df = pd.DataFrame([{"subject": clean_subj, "body": clean_body}])

        # Generate predictions
        try:
            probas = self.pipeline.predict_proba(input_df)[0]
            pred_id = int(np.argmax(probas))
        except Exception:
            # Fallback if predict_proba is unavailable
            pred_id = int(self.pipeline.predict(input_df)[0])
            probas = [0.0, 0.0, 0.0]
            probas[pred_id] = 1.0

        predicted_class = ID_TO_CLASS[pred_id]
        
        prob_dict = {
            CLASS_NAMES[i]: round(float(probas[i]), 4)
            for i in range(len(CLASS_NAMES))
        }

        # Extract indicators & assess security posture
        indicators = extract_threat_indicators(clean_subj, clean_body)
        assessment = assess_threat_level(predicted_class, prob_dict, indicators)

        return {
            "predicted_class": predicted_class,
            "predicted_label_id": pred_id,
            "probabilities": prob_dict,
            "confidence": round(float(np.max(probas)), 4),
            "severity": assessment["severity"],
            "badge_color": assessment["badge_color"],
            "risk_score_out_of_10": assessment["risk_score_out_of_10"],
            "summary": assessment["summary"],
            "indicators": indicators,
            "recommendations": assessment["recommendations"],
            "warning": "Disclaimer: Research prototype designed for defense-in-depth assistance. Machine learning models cannot guarantee 100% email safety. Always adhere to organizational security protocols."
        }