from flask import Flask, jsonify, request

import numpy as np

import joblib

import traceback

from pathlib import Path

from src.preprocessing.preprocessing import create_ordered_medical_pipeline

from src.utils.logger import get_logger

from typing import Dict, Any

# Initialize Flask app

app = Flask(__name__)

class PredictionService:

    """Service for handling model predictions"""

    def __init__(self, model_path: str = 'src/models/model.joblib'):

        self.model_path = Path(model_path)

        self.pipeline = None

        self.preprocessor = None

        self.logger = get_logger(__name__)

        try:

            self.initialize()

        except Exception as e:

            error_traceback = traceback.format_exc()

            self.logger.error(f"Initialization failed with error: {str(e)}")

            self.logger.error(f"Traceback: {error_traceback}")

            raise

    def initialize(self):

        """Initialize model pipeline"""

        try:

            self.logger.info(f"Looking for model at: {self.model_path.absolute()}")

            # Load model pipeline

            if not self.model_path.exists():

                raise FileNotFoundError(f"Model file not found at {self.model_path.absolute()}")

            # Load the pipeline

            self.pipeline = joblib.load(self.model_path)

            self.logger.info("Model pipeline loaded successfully")

            # Log pipeline contents

            self.logger.info("Pipeline contents:")

            for key in self.pipeline.keys():

                self.logger.info(f"- Found component: {key}")

            # Initialize preprocessor

            self.logger.info("Initializing preprocessor...")

            self.preprocessor = create_ordered_medical_pipeline()

            self.logger.info("Preprocessor initialized successfully")

            # Log feature dimensions

            self.logger.info(f"Expected feature dimension: {self.pipeline['feature_dim']}")

            self.logger.info(f"Vectorizer vocabulary size: {len(self.pipeline['vectorizer'].vocabulary_)}")

        except Exception as e:

            error_traceback = traceback.format_exc()

            self.logger.error(f"Initialization failed with error: {str(e)}")

            self.logger.error(f"Traceback: {error_traceback}")

            raise

    def predict(self, text: str) -> Dict[str, Any]:

        """Make prediction on input text"""

        try:

            # Validate input

            if not isinstance(text, str):

                raise ValueError("Input must be a string")

            if not text.strip():

                raise ValueError("Input text cannot be empty")

            # Preprocess text

            processed_text = self.preprocessor.process(text)

            if isinstance(processed_text, tuple):

                processed_text = processed_text[0]

            # Extract features using vectorizer

            features = self.pipeline['vectorizer'].transform([processed_text])

            features = features.toarray()

            # Verify feature dimension

            if features.shape[1] != self.pipeline['feature_dim']:

                raise ValueError(

                    f"Feature dimension mismatch: got {features.shape[1]}, "

                    f"expected {self.pipeline['feature_dim']}"

                )

            # Scale features

            features_scaled = self.pipeline['scaler'].transform(features)

            # Get the model - the model itself is a VotingClassifier

            model = self.pipeline['model']

            # Make prediction

            prediction = model.predict(features_scaled)[0]

            # Get prediction probability if available

            confidence = None

            if hasattr(model, 'predict_proba'):

                probabilities = model.predict_proba(features_scaled)[0]

                confidence = float(np.max(probabilities))

            # Convert prediction using label encoder if available

            if 'metadata' in self.pipeline and 'label_encoder' in self.pipeline['metadata']:

                prediction = self.pipeline['metadata']['label_encoder'].inverse_transform([prediction])[0]

            return {

                'status': 'success',

                'prediction': prediction,

                'confidence': confidence

            }

        except Exception as e:

            error_traceback = traceback.format_exc()

            self.logger.error(f"Prediction failed: {str(e)}")

            self.logger.error(f"Traceback: {error_traceback}")

            raise

# Initialize service with error handling

try:

    prediction_service = PredictionService()

    app.logger.info("PredictionService initialized successfully")

except Exception as e:

    app.logger.error(f"Failed to initialize PredictionService: {str(e)}")

    traceback.print_exc()

    prediction_service = None

@app.route("/health")

def health_check():

    """Health check endpoint"""

    if prediction_service is None:

        return jsonify({

            'status': 'unhealthy',

            'error': 'Prediction service failed to initialize'

        }), 500

    health_info = {

        'status': 'healthy',

        'components': {

            'model': prediction_service.pipeline is not None and 'model' in prediction_service.pipeline,

            'vectorizer': prediction_service.pipeline is not None and 'vectorizer' in prediction_service.pipeline,

            'scaler': prediction_service.pipeline is not None and 'scaler' in prediction_service.pipeline,

            'preprocessor': prediction_service.preprocessor is not None

        }

    }

    return jsonify(health_info)

@app.route("/predict", methods=["POST"])

def predict():

    """Prediction endpoint"""

    if prediction_service is None:

        return jsonify({

            'status': 'error',

            'message': 'Prediction service is not available'

        }), 503

    try:

        # Get request data

        data = request.get_json()

        # Validate request data

        if not data:

            return jsonify({

                'status': 'error',

                'message': 'No data provided'

            }), 400

        if 'description' not in data:

            return jsonify({

                'status': 'error',

                'message': 'No description provided'

            }), 400

        # Get prediction

        result = prediction_service.predict(data['description'])

        return jsonify(result)

    except ValueError as e:

        return jsonify({

            'status': 'error',

            'message': str(e)

        }), 400

    except Exception as e:

        error_traceback = traceback.format_exc()

        app.logger.error(f"Prediction failed: {str(e)}")

        app.logger.error(f"Traceback: {error_traceback}")

        return jsonify({

            'status': 'error',

            'message': str(e),

            'traceback': error_traceback

        }), 500

if __name__ == "__main__":

    app.run(debug=True, host='0.0.0.0', port=5000)