from fastapi import FastAPI, HTTPException, Depends, Header  # Import FastAPI framework, error handling, dependency injection, and header extraction utilities
import os  # Import the os module to read sensitive API keys from environment variables
import hmac  # Import hmac for cryptographically secure constant-time string comparison (prevents timing attacks)

app = FastAPI()  # Instantiate the FastAPI application; this object routes HTTP requests to handler functions

# In production, load from a secret store, not env vars
TENANT_KEYS = {  # Define a dictionary that maps each tenant ID to its corresponding secret API key
    "ministry-health": os.getenv("API_KEY_HEALTH"),  # Load the health ministry API key from the API_KEY_HEALTH environment variable
    "ministry-defense": os.getenv("API_KEY_DEFENSE"),  # Load the defense ministry API key from the API_KEY_DEFENSE environment variable
    "ministry-finance": os.getenv("API_KEY_FINANCE"),  # Load the finance ministry API key from the API_KEY_FINANCE environment variable
}

def verify_tenant(authorization: str = Header(None)):  # Define a FastAPI dependency that extracts the Authorization header; defaults to None if the header is absent
    if not authorization or not authorization.startswith("Bearer "):  # Check that the header exists AND follows the standard "Bearer <token>" format
        raise HTTPException(401, "Missing or invalid auth header")  # Return HTTP 401 Unauthorized immediately if the header is missing or malformed

    token = authorization.removeprefix("Bearer ")  # Strip the "Bearer " prefix (7 characters) to extract just the raw API key token

    for tenant_id, secret in TENANT_KEYS.items():  # Iterate through every known tenant ID and its corresponding secret key
        if hmac.compare_digest(token, secret):  # Use constant-time comparison to prevent timing side-channel attacks that could leak key character differences
            return tenant_id  # Return the tenant ID (e.g., "ministry-health") if the provided token matches a known secret

    raise HTTPException(401, "Invalid API key")  # Return HTTP 401 if no tenant's secret matched the provided token after exhausting all keys

@app.post("/v1/completions")  # Register a POST endpoint at /v1/completions; clients send prompts to this route
async def completions(  # Define an asynchronous handler function so the server can handle other requests while waiting for the LLM
    prompt: str,  # Accept the user's text prompt as a string (FastAPI extracts this from the JSON request body)
    tenant: str = Depends(verify_tenant),  # Use FastAPI's dependency injection system to run verify_tenant automatically; the returned tenant ID is injected here
):
    # tenant is always "ministry-health" or "ministry-defense" etc.
    # Never trust the client to provide the tenant ID
    return await route_to_llm(prompt, tenant)  # Forward the validated prompt and the authenticated tenant ID to the LLM routing function; await because it performs async I/O

import logging  # Import Python's built-in logging framework for structured, configurable log output
import json  # Import json to serialize log records as JSON strings for structured parsing
from datetime import datetime, timezone  # Import datetime for precise timestamps and timezone for UTC awareness

class TenantLogger:  # Define a custom logger that scopes all log output to a single tenant's dedicated file
    def __init__(self, tenant_id: str):  # Constructor accepts the unique tenant identifier
        self.tenant_id = tenant_id  # Store the tenant ID instance variable so all methods can reference it
        # Each tenant writes to their own file
        self.logger = logging.getLogger(f"llm-{tenant_id}")  # Create (or retrieve) a named logger unique to this tenant to avoid cross-tenant logger conflicts
        handler = logging.FileHandler(f"/var/log/llm/{tenant_id}/prompts.log")  # Create a file handler pointing to a tenant-specific path; the directory name is the tenant ID
        handler.setFormatter(logging.Formatter("%(message)s"))  # Set the formatter to output only the raw message string (our JSON), omitting default metadata like timestamps
        self.logger.addHandler(handler)  # Attach the file handler to the logger so messages are written to the tenant's file
        self.logger.setLevel(logging.INFO)  # Set the logging threshold to INFO; only messages at INFO level or above (WARNING, ERROR, CRITICAL) will be recorded

    def log_interaction(self, prompt: str, response: str, model: str):  # Method to log one LLM interaction for audit and monitoring purposes
        record = {  # Build a structured dictionary representing the log entry
            "timestamp": datetime.now(timezone.utc).isoformat(),  # Generate an ISO 8601 UTC timestamp for unambiguous chronological ordering
            "tenant": self.tenant_id,  # Include the tenant ID so logs can be filtered, attributed, or billed per tenant
            "model": model,  # Record which model name served this request (e.g., "llama3") for model performance tracking
            "prompt_hash": hash(prompt),  # Store a numeric hash of the prompt instead of raw plaintext to comply with data minimization principles
            "response_length": len(response),  # Record the character count of the response for usage metering and capacity planning
        }
        self.logger.info(json.dumps(record))  # Serialize the dictionary to a JSON string and write it as a single log line at INFO level

import bentoml  # Import BentoML, an ML model serving framework that wraps models as scalable microservices
import numpy as np  # Import NumPy for numerical array operations, commonly used in embedding and inference pipelines
from bentoml.io import JSON  # Import the JSON input/output deserializer to define typed API endpoints

# Tenant-specific runners
health_runner = bentoml.transformers.get("llama3-health:latest").to_runner()  # Load the health ministry's fine-tuned Llama 3 model from the BentoML model store and convert it to an async runner
defense_runner = bentoml.transformers.get("llama3-defense:latest").to_runner()  # Load the defense ministry's fine-tuned Llama 3 model from the BentoML model store and convert it to an async runner

svc = bentoml.Service("multi-tenant-llm", runners=[health_runner, defense_runner])  # Create a BentoML service named "multi-tenant-llm" that manages both model runners for concurrent serving

TENANT_RUNNERS = {  # Build a lookup dictionary that maps each tenant ID to its dedicated model runner
    "ministry-health": health_runner,  # Associate the health ministry ID with the health-specific model runner
    "ministry-defense": defense_runner,  # Associate the defense ministry ID with the defense-specific model runner
}

@svc.api(input=JSON(), output=JSON())  # Decorate this method as a BentoML API endpoint; it accepts JSON input and returns JSON output
async def generate(input_data: dict) -> dict:  # Define an async handler that receives a dictionary and returns a dictionary
    tenant = input_data.get("tenant")  # Extract the tenant ID from the request JSON body; this was injected by the API gateway, never taken from client input
    prompt = input_data.get("prompt")  # Extract the user's prompt text from the request JSON body

    runner = TENANT_RUNNERS.get(tenant)  # Look up the correct model runner for this tenant from the dictionary; returns None if tenant is unknown
    if not runner:  # Check if the lookup returned None, meaning no dedicated runner exists for this tenant ID
        return {"error": "unknown tenant"}  # Return an error response immediately if the tenant is not recognized; this prevents accidental cross-tenant routing

    # System prompt enforces data isolation
    system = {  # Define a dictionary of tenant-specific system prompts that constrain model behavior
        "ministry-health": "You are a medical AI assistant. Never reference data from any other government department.",  # Health AI: restricts the model to medical domain only
        "ministry-defense": "You are a defense AI assistant. Never reference data from any other government department.",  # Defense AI: restricts the model to defense domain only
    }

    full_prompt = f"{system[tenant]}\n\n{user_prompt}"  # Prepend the tenant's system prompt to the user prompt, separated by a blank line for model instruction clarity
    result = await runner.async_run(full_prompt)  # Execute the model runner asynchronously with the composed full prompt and await the generated text response

    # Log with tenant context
    logger = TenantLogger(tenant)  # Instantiate a new TenantLogger scoped to this tenant for recording the interaction
    logger.log_interaction(prompt, result, "llama3")  # Log the original user prompt, the model's response, and the model name for compliance auditing

    return {"response": result}  # Return the model's generated text wrapped in a JSON response object with a "response" key

from qdrant_client import QdrantClient  # Import the Qdrant vector database client for similarity search and vector storage operations

qdrant = QdrantClient("http://qdrant:6333")  # Initialize a connection to the Qdrant server running in the cluster at the default gRPC/REST port 6333

TENANT_COLLECTIONS = {  # Map each tenant ID to its dedicated Qdrant collection name for strict data isolation at the database level
    "ministry-health": "documents_health",  # The health ministry stores its document vectors in the "documents_health" collection
    "ministry-defense": "documents_defense",  # The defense ministry stores its document vectors in the "documents_defense" collection
}

def retrieve_context(tenant: str, query: str, top_k: int = 5):  # Define a function to retrieve relevant documents for a given tenant and search query
    collection = TENANT_COLLECTIONS.get(tenant)  # Look up the collection name for this tenant; returns None if the tenant has no configured collection
    if not collection:  # Check whether the lookup returned None (unknown tenant or no collection assigned)
        return []  # Return an empty result list; this prevents any accidental cross-tenant data exposure

    results = qdrant.search(  # Execute a vector similarity search against the tenant's dedicated collection
        collection_name=collection,  # Specify the tenant-specific collection; Qdrant will only search within this collection's vectors
        query_vector=embed(query),  # Convert the query string into a dense vector embedding using a shared embedding model function
        limit=top_k,  # Limit the search results to the top-k most similar vectors (defaults to 5 if not overridden)
    )
    return [r.payload["text"] for r in results]  # Extract and return the original document text from each search result's metadata payload

import hashlib  # Import hashlib for SHA-256 cryptographic hashing to build tamper-evident hash chains
import json  # Import json for serializing audit records into a portable text format
from datetime import datetime, timezone  # Import datetime for precise event timestamps and timezone for UTC normalization

class AuditLogger:  # Define an audit logger that produces hash-chained, tamper-evident log entries
    def __init__(self, tenant_id: str):  # Constructor accepts the tenant identifier to scope audit files
        self.tenant_id = tenant_id  # Store the tenant ID so all log writes go to the correct tenant directory
        self.prev_hash = self._load_last_hash()  # Load the hash from the most recent audit entry to continue the cryptographic chain

    def _load_last_hash(self):  # Private helper that reads the previous entry's hash from disk to maintain chain continuity
        # Read the last line's hash from the audit file
        try:  # Wrap in try/except because the audit file may not exist yet (first entry) or could be empty
            with open(f"/audit/{self.tenant_id}/audit.log", "r") as f:  # Open the tenant's audit file in read mode
                last_line = f.readlines()[-1]  # Read all lines and select the last line, which is the most recent audit entry
                return json.loads(last_line)["hash"]  # Parse the JSON and extract the "hash" field from the most recent record
        except (FileNotFoundError, IndexError, KeyError):  # Catch three scenarios: file missing, file empty (no lines), or JSON missing "hash" key
            return "0" * 64  # Return a 64-character hex string of zeros as the genesis hash for the very first audit entry

    def log(self, action: str, details: dict):  # Public method to append a new tamper-evident audit entry to the chain
        record = {  # Build the core audit record with metadata about the event being logged
            "timestamp": datetime.now(timezone.utc).isoformat(),  # Generate an ISO 8601 UTC timestamp for precise, timezone-agnostic event ordering
            "tenant": self.tenant_id,  # Tag the record with the tenant ID so audit trails are attributable per tenant
            "action": action,  # Describe the action that occurred (e.g., "prompt_submitted", "model_deployed", "api_key_rotated")
            "details": details,  # Include arbitrary structured data with event-specific context for forensic investigation
        }

        # Chain hashes: each record includes the hash of the previous record
        record_str = json.dumps(record, sort_keys=True)  # Serialize the record with sorted keys to ensure deterministic, reproducible hashing across platforms
        record["prev_hash"] = self.prev_hash  # Insert the previous entry's hash into this record, forming a cryptographic linked list
        record["hash"] = hashlib.sha256(  # Compute a SHA-256 hash that binds this record to the previous one
            (record_str + self.prev_hash).encode()  # Concatenate the canonical JSON string with the previous hash and encode to UTF-8 bytes for the hash function
        ).hexdigest()  # Convert the raw binary digest to a 64-character hexadecimal string for human-readable storage

        with open(f"/audit/{self.tenant_id}/audit.log", "a") as f:  # Open the tenant's audit file in append mode (never overwrites)
            f.write(json.dumps(record) + "\n")  # Append the complete record (with prev_hash and hash) as a newline-terminated JSON line

        self.prev_hash = record["hash"]  # Update the stored previous hash to the newly computed hash so the next call chains correctly

# Pseudocode — this runs as a Kubernetes MutatingWebhookConfiguration
def mutate_pod(pod):  # Define the webhook handler function; Kubernetes calls this before persisting every pod creation
    namespace = pod.metadata.namespace  # Extract the namespace from the incoming pod's object metadata

    # Only apply to ministry namespaces
    if not namespace.startswith("ministry-"):  # Check if the namespace name begins with the "ministry-" prefix
        return pod  # Return the pod unchanged; non-ministry namespaces are outside this webhook's scope

    tenant_id = namespace.replace("ministry-", "")  # Derive the tenant ID by removing the "ministry-" prefix from the namespace (e.g., "ministry-health" -> "health")

    # Inject tenant label
    pod.metadata.labels["tenant-id"] = tenant_id  # Add a "tenant-id" label with the tenant value so NetworkPolicies and other controllers can identify tenant affiliation

    # Inject sidecar that enforces egress rules
    sidecar = Container(  # Create a sidecar container definition that will be injected alongside the user's application containers
        name="tenant-proxy",  # Name the sidecar container "tenant-proxy" for clear identification in logs and monitoring
        image="tenant-proxy:1.0",  # Use a purpose-built proxy image that enforces tenant-scoped egress rules at the network layer
        env=[EnvVar(name="TENANT_ID", value=tenant_id)],  # Pass the derived tenant ID as an environment variable so the proxy knows which tenant rules to enforce
    )
    pod.spec.containers.append(sidecar)  # Append the sidecar container to the pod's existing container list; it will share the pod's network namespace

    return pod  # Return the modified pod object so Kubernetes creates the pod with the injected label and sidecar