# monolith/app.py — heavily simplified
from fastapi import FastAPI, Request  # Import FastAPI framework and the Request type for HTTP handler signatures
import pickle  # Import pickle to deserialize the trained model and vectorizer from disk into usable Python objects

app = FastAPI()  # Create the FastAPI application instance — this is the entry point of the entire monolith process

# Loaded at startup — takes 8 seconds
model = pickle.load(open("/models/classifier.pkl", "rb"))  # Load the classifier model at import time, blocking server readiness for ~8s
vectorizer = pickle.load(open("/models/vectorizer.pkl", "rb"))  # Load the vectorizer at import time, adding 8+ seconds to startup

@app.post("/assistant")  # Register the /assistant POST endpoint that handles AI inference for 2 million users
async def assistant(request: Request):  # Define an async handler so FastAPI can interleave other requests during I/O waits
    data = await request.json()  # Await parsing of the incoming JSON body from the client
    user_text = data["text"]  # Extract the 'text' field from the parsed request body

    # Preprocessing
    tokens = tokenize(user_text)  # Tokenize the raw text into tokens — word splits, stemming, stopword removal, etc.
    features = vectorizer.transform([tokens])  # Transform the token list into the feature matrix the model was trained on

    # Inference — the slow part
    prediction = model.predict(features)[0]  # Run the model's predict method — this is where the 800ms average lives
    confidence = max(model.predict_proba(features)[0])  # Get the highest class probability as the confidence score

    # Post-processing
    response = format_response(prediction, confidence)  # Format the raw prediction and confidence into the API response shape
    return response  # Return the final response to the client — every other request is blocked during this call

@app.get("/health")  # Register a /health GET endpoint for Kubernetes liveness probes and load balancer health checks
async def health():  # Keep the health check trivially fast — it must never touch the model or database
    return {"status": "ok"}  # Return a simple status dict to indicate the process is alive and accepting connections

# Step 1: Add an inference client abstraction
import os  # Import os to read environment variables that control the routing mode without code changes
import httpx  # Import the async HTTP client library for making outbound requests to the inference service
import json  # Import json for serializing/deserializing request and response payloads over HTTP

class InferenceClient:
    """Routes to in-process model or remote service based on config."""  # Docstring: this class is the strangler fig facade
    def __init__(self):  # Constructor reads the routing mode from the environment at instantiation time
        self.mode = os.getenv("INFERENCE_MODE", "local")  # Read the routing mode; default "local" means zero change on first deploy
        self.remote_url = os.getenv("INFERENCE_URL", "http://localhost:8001/predict")  # Read the remote service URL with a sensible localhost default for dev
        self.model = None  # Initialize the model reference as None; only populated in local mode to avoid memory waste
        self.vectorizer = None  # Initialize the vectorizer reference as None; populated only when INFERENCE_MODE is "local"

        if self.mode == "local":  # Only load the model files when we are in local mode — saves startup time in remote mode
            import pickle  # Import pickle inside this conditional so it is never imported in remote/canary/shadow modes
            self.model = pickle.load(open("/models/classifier.pkl", "rb"))  # Deserialize the classifier model from disk
            self.vectorizer = pickle.load(open("/models/vectorizer.pkl", "rb"))  # Deserialize the vectorizer from disk

    async def predict(self, text: str):  # Async predict so the calling FastAPI handler does not block the event loop
        if self.mode == "local":  # Local mode: run inference in-process, identical to the original monolith behavior
            tokens = tokenize(text)  # Tokenize the input text using the existing preprocessing pipeline
            features = self.vectorizer.transform([tokens])  # Transform the token list into the feature matrix
            pred = self.model.predict(features)[0]  # Get the predicted class label from the loaded model
            conf = max(self.model.predict_proba(features)[0])  # Get the highest class probability as the confidence score
            return {"prediction": pred, "confidence": float(conf)}  # Return results in the exact shape the remote service will use
        else:  # Remote/shadow/canary mode: delegate the inference to the external inference service over HTTP
            async with httpx.AsyncClient() as client:  # Open an async HTTP client in a context manager for proper connection cleanup
                resp = await client.post(self.remote_url, json={"text": text}, timeout=5)  # POST the text to the inference service, timing out after 5 seconds
                resp.raise_for_status()  # Raise an exception on non-2xx responses so callers can trigger fallback logic
                return resp.json()  # Parse and return the JSON response from the remote service

# Replace the direct model call
client = InferenceClient()  # Instantiate the client once at module level — routing mode is locked in at construction time

@app.post("/assistant")  # Register the same /assistant endpoint the monolith always had — the interface does not change
async def assistant(request: Request):  # The handler signature remains identical, so no client or test code needs updating
    data = await request.json()  # Parse the incoming JSON body from the request
    result = await client.predict(data["text"])  # Delegate all inference work to InferenceClient regardless of the mode
    response = format_response(result["prediction"], result["confidence"])  # Format the prediction into the same response shape as before
    return response  # Return the formatted response — callers cannot tell whether inference ran locally or remotely

# inference_service/app.py  # This is a brand-new standalone service, deployed separately from the monolith
from fastapi import FastAPI, Request  # Use FastAPI just like the monolith for consistency and developer familiarity
import pickle  # Import pickle to load the model and vectorizer in this dedicated service's memory
import time  # Import time to measure and report per-request inference latency

app = FastAPI()  # Create the standalone inference service instance — runs on its own port (e.g., 8001)

model = pickle.load(open("/models/classifier.pkl", "rb"))  # Load the classifier model at startup — inference now lives here
vectorizer = pickle.load(open("/models/vectorizer.pkl", "rb"))  # Load the vectorizer at startup — model loading is isolated from the monolith

@app.post("/predict")  # Register the /predict endpoint that the monolith's InferenceClient will call via HTTP
async def predict(request: Request):  # Async handler so this service can handle many concurrent inference requests
    data = await request.json()  # Parse the incoming JSON payload sent by the InferenceClient
    start = time.time()  # Record the start time for computing per-request latency

    tokens = tokenize(data["text"])  # Tokenize the input text using the same preprocessing as the original monolith
    features = vectorizer.transform([tokens])  # Transform the tokens into the feature matrix the model expects
    pred = model.predict(features)[0]  # Run the model to get the predicted class label
    conf = max(model.predict_proba(features)[0])  # Extract the highest class probability as the confidence score

    return {  # Return the response in the exact shape the InferenceClient expects
        "prediction": pred,  # The predicted class label
        "confidence": float(conf),  # The confidence score converted to a JSON-safe float
        "latency_ms": round((time.time() - start) * 1000),  # Report per-request latency so the monolith can track performance metrics
    }

@app.get("/health")  # Register a separate /health endpoint for the inference service's own Kubernetes probe
async def health():  # Dedicated health check — each service has independent lifecycle management
    return {"status": "ok"}  # Return a simple status to indicate the inference service is ready and the model is loaded

class InferenceClient:
    async def predict(self, text: str):  # Extend the predict method with shadow-mode support — local result is still the source of truth
        local_result = await self._predict_local(text)  # Always compute the local prediction first — this is what the user actually sees

        if self.mode == "shadow":  # Shadow mode: fire a copy of every request to the remote service without using its response
            # Fire-and-forget to the remote service
            try:  # Wrap the remote call in try/except so a slow or failing remote never blocks the user response
                async with httpx.AsyncClient() as client:  # Open an async HTTP client for the shadow request
                    await client.post(  # Fire the request to the remote service but discard the response — fire-and-forget
                        self.remote_url,  # The URL of the new inference service
                        json={"text": text, "shadow": True},  # Include a "shadow": True flag so the remote can log it differently
                        timeout=2,  # Use a shorter timeout for shadow traffic — the user is not waiting on this call
                    )
            except Exception as e:  # Catch any network error, timeout, or server error from the shadow request
                logger.warning(f"Shadow inference failed: {e}")  # Log the shadow failure but do NOT alert — shadow is best-effort

        return local_result  # Always return the local result — the shadow response is purely for validation and comparison

import structlog  # Import structlog for structured, searchable log output that is easy to query in production

shadow_logger = structlog.get_logger("shadow-inference")  # Create a dedicated logger with a distinct name so logs can be routed to a separate stream

def log_shadow_comparison(text, local_result, remote_result, local_latency, remote_latency):  # Log a side-by-side comparison of local vs remote results
    if local_result["prediction"] != remote_result["prediction"]:  # Only log when predictions diverge — agreement is the happy path and too noisy to log
        shadow_logger.warning("prediction_mismatch",  # Emit a warning-level structured log for every prediction mismatch
            local=local_result,  # Include the full local result for debugging
            remote=remote_result,  # Include the full remote result for comparison
            text_hash=hash(text),  # Include a hash of the input text (not the raw text) for lookup without storing PII
        )

class InferenceClient:
    async def predict(self, user_id: str, text: str):  # Extended predict now takes a user_id for deterministic canary routing
        if self.mode == "remote" or (self.mode == "canary" and hash(user_id) % 100 < 1):  # Route to remote if full-remote mode, or if canary mode AND user_id hashes to the first percentile
            try:  # Wrap the remote call so any failure triggers a graceful fallback to the local model
                result = await self._predict_remote(text)  # Call the remote inference service
                return result  # Return the remote result to the user — this is the faster path we want to validate
            except Exception as e:  # Catch any remote failure — network error, timeout, server error, etc.
                logger.error(f"Remote inference failed, falling back to local: {e}")  # Log the failure for operational awareness
                return await self._predict_local(text)  # Fall back to the local model — the user sees slower response instead of an error
        else:  # User is not in the canary cohort or mode is still "local"
            return await self._predict_local(text)  # Use the local model for this request — business as usual

# metrics.py
from prometheus_client import Histogram, Counter, Gauge  # Import Prometheus metric types for observability in production

inference_latency = Histogram(  # Create a Histogram to track inference latency distribution across modes
    "inference_latency_seconds",  # Metric name exposed on the /metrics endpoint
    "Inference latency",  # Help text displayed in Prometheus for this metric
    ["mode"],  # Label the histogram by mode (local, remote, shadow) so we can compare latency distributions
)

prediction_mismatch = Counter(  # Create a Counter to track how often local and remote predictions disagree
    "prediction_mismatches_total",  # Metric name — a monotonically increasing counter
    "Count of prediction mismatches between local and remote",  # Help text explaining what this counter represents
)

def record_inference(mode, latency, prediction, expected=None):  # Utility function to record inference metrics consistently
    inference_latency.labels(mode=mode).observe(latency)  # Record the latency in the histogram under the given mode label
    if expected is not None and prediction != expected:  # Only increment the mismatch counter when we have an expected value AND the prediction differs
        prediction_mismatch.inc()  # Increment the mismatch counter — a spike here triggers an alert to halt the canary rollout

# Clean up the monolith
class InferenceClient:
    def __init__(self):  # After full rollout, the constructor no longer has a local mode branch — it only knows remote
        self.remote_url = os.getenv("INFERENCE_URL", "http://inference:8001/predict")  # Read the remote service URL; the default now assumes a Kubernetes service DNS name

    async def predict(self, text: str):  # The predict method is now pure HTTP — no local inference code path remains
        async with httpx.AsyncClient() as client:  # Open an async HTTP client for each prediction call
            resp = await client.post(self.remote_url, json={"text": text}, timeout=5)  # POST the text to the inference service with a 5-second timeout
            resp.raise_for_status()  # Raise on non-2xx responses — there is no local fallback anymore, so this is a real error
            return resp.json()  # Parse and return the JSON response from the dedicated inference service

# inference_service/cache.py
import redis  # Import the Redis client for caching repeated inference results
import json  # Import json for serializing and deserializing cached predictions

cache = redis.Redis(host="redis", port=6379, decode_responses=True)  # Connect to Redis — decode_responses=True returns strings instead of bytes for convenience

@app.post("/predict")  # The same /predict endpoint, now augmented with a cache check before running inference
async def predict(request: Request):  # The handler signature stays the same so the InferenceClient does not need changes
    data = await request.json()  # Parse the incoming request JSON body
    text_hash = hashlib.sha256(data["text"].encode()).hexdigest()  # Compute a SHA-256 hash of the input text to use as the cache key — this avoids storing PII in Redis

    # Check cache
    cached = cache.get(f"pred:{text_hash}")  # Look up the cache key in Redis — returns the cached JSON string or None
    if cached:  # Cache hit: the exact same query was processed recently
        return json.loads(cached)  # Deserialize and return the cached result immediately — this takes under 5ms instead of 800ms

    # Run inference
    result = run_inference(data["text"])  # Cache miss: run the full inference pipeline (tokenize, vectorize, predict, format)

    # Cache for 1 hour
    cache.setex(f"pred:{text_hash}", 3600, json.dumps(result))  # Store the result in Redis with a 1-hour TTL (3600 seconds) as a JSON string
    return result  # Return the freshly computed result — subsequent identical queries will hit the cache