from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np
import redis
import json
from typing import List

app = FastAPI()
redis_client = redis.Redis(host='localhost', port=6379, db=0)

class PredictionRequest(BaseModel):
    user_id: str
    item_features: List[float]

class FeatureStore:
    def __init__(self, redis_client):
        self.redis = redis_client
    
    def get_features(self, entity_id: str, feature_names: List[str]) -> dict:
        features = {}
        for name in feature_names:
            value = self.redis.hget(f"features:{entity_id}", name)
            if value:
                features[name] = json.loads(value)
        return features
    
    def set_features(self, entity_id: str, features: dict):
        self.redis.hset(f"features:{entity_id}", mapping={
            k: json.dumps(v) for k, v in features.items()
        })

class ModelServer:
    def __init__(self, model_path):
        self.model = self.load_model(model_path)
    
    def predict(self, features: dict) -> float:
        X = self.preprocess(features)
        return float(self.model.predict(X)[0])
    
    def preprocess(self, features: dict) -> np.ndarray:
        return np.array([list(features.values())])

feature_store = FeatureStore(redis_client)
model_server = ModelServer("model.pkl")

@app.post("/predict")
async def predict(request: PredictionRequest):
    # 1. Get features from feature store
    features = feature_store.get_features(
        request.user_id,
        ["user_clicks", "user_purchases", "item_popularity"]
    )
    
    # 2. Get model prediction
    prediction = model_server.predict(features)
    
    # 3. Log prediction for monitoring
    log_prediction(request.user_id, prediction)
    
    return {"prediction": prediction, "features_used": list(features.keys())}

@app.post("/predict_batch")
async def predict_batch(user_ids: List[str]):
    # Batch prediction for efficiency
    results = []
    for user_id in user_ids:
        features = feature_store.get_features(user_id, [...])
        prediction = model_server.predict(features)
        results.append({"user_id": user_id, "prediction": prediction})
    return results

# Monitoring endpoint
@app.get("/monitoring/drift")
async def check_drift():
    recent_predictions = get_recent_predictions(last_hour=True)
    baseline_distribution = get_baseline_distribution()
    
    psi = calculate_psi(baseline_distribution, recent_predictions)
    
    return {
        "psi_score": psi,
        "drift_detected": psi > 0.1,
        "recommendation": "Retrain model" if psi > 0.25 else "Monitor"
    }