import joblib, numpy as np, time
from typing import Dict, Any
from prometheus_client import Counter, Histogram, Gauge, start_http_server
import logging, traceback
from kserve import Model, ModelServer

# Custom metrics
PREDICTION_COUNTER = Counter('model_predictions_total', 'Total predictions')
PREDICTION_ERRORS  = Counter('model_prediction_errors_total', 'Total errors')
PREDICTION_LATENCY = Histogram('model_prediction_latency_seconds', 'Prediction latency')
MODEL_ACCURACY     = Gauge('model_accuracy', 'Model accuracy')

class TaxiModel(Model):
    def __init__(self, name: str):
        super().__init__(name)
        self.name = name
        self.model = joblib.load('/mnt/models/model')
        start_http_server(8000)  # Expose metrics endpoint
        MODEL_ACCURACY.set(0.93) # Current model accuracy baseline
        self.ready = True

    def predict(self, request, headers: dict = None) -> dict:
        logging.info(f"Received request: {request}")
        start = time.time()
        try:
            instance = request["instances"][0]
            start_location = instance["start_location"]
            end_location = instance["end_location"]
            X = np.array([[start_location, end_location]])
            price = self.model.predict(X)[0]
            PREDICTION_COUNTER.inc()
            PREDICTION_LATENCY.observe(time.time() - start)
            return {"predictions": [{"predicted_price": float(price)}]}
        except Exception:
            PREDICTION_ERRORS.inc()
            raise

if __name__ == "__main__":
    logging.error("In main")
    model = TaxiModel("taxi-model")

    #print("Attempting to load model manually in the main thread...")
    #model.load()
    #print("Model loaded successfully in main thread. Starting server...")

    logging.error("Starting model server")
    try:
        ModelServer().start([model])
    except Exception as e:
        # This block catches any exception raised during startup or readiness check
        logging.error("!!!!!!!!!!!!!!!! FATAL ERROR CAUGHT DURING STARTUP !!!!!!!!!!!!!!!!")
        logging.error(f"Error type: {type(e).__name__}")
        logging.error(f"Error message: {e}")

        # This function prints the full stack trace to stderr
        traceback.print_exc()

        logging.error("!!!!!!!!!!!!!!!! END OF ERROR TRACEBACK !!!!!!!!!!!!!!!!")

        # Ensure the process exits with an error status code
        exit(1)