#!/usr/bin/env python3
"""
LockStock FastAPI Inference Server Example
-------------------------------------------
Complete working example of a FastAPI inference server with LockStock chain auth.

Usage:
    1. Provision agent: https://lockstock-api.onrender.com
    2. Bind to hardware: lockstock-guard bind --agent <agent_id> --token <token>
    3. Run server: python examples/fastapi_inference_server.py
    4. Test: curl -X POST http://localhost:8000/v1/chat/completions \
             -H "Content-Type: application/json" \
             -d '{"messages": [{"role": "user", "content": "Hello"}]}'
"""

import os
import time
import uvicorn
from typing import List, Optional
from pydantic import BaseModel
from fastapi import FastAPI, HTTPException

# Import LockStock middleware
from lockstock_fastapi import LockStockMiddleware


# Request/Response Models
class Message(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    messages: List[Message]
    temperature: float = 0.7
    max_tokens: int = 512


class ChatResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: List[dict]


# Create FastAPI app
app = FastAPI(
    title="LockStock Inference Server",
    description="FastAPI inference server with LockStock chain authentication",
    version="1.0.0",
)


# Add LockStock Middleware
# Option 1: From environment variables
if os.getenv("LOCKSTOCK_AGENT_ID"):
    app.add_middleware(LockStockMiddleware)
    print("✓ LockStock middleware loaded from environment")
else:
    # Option 2: Explicit configuration (for testing)
    print("⚠ LOCKSTOCK_AGENT_ID not set - running WITHOUT chain auth")
    print("  Set LOCKSTOCK_AGENT_ID to enable LockStock")


# Mock inference function (replace with your model)
async def generate_response(messages: List[Message]) -> str:
    """
    Replace this with your actual model inference.

    Examples:
    - OpenAI API call
    - Local model (transformers, vllm, etc.)
    - Another inference service
    """
    # Mock response for demonstration
    last_message = messages[-1].content
    return f"Mock response to: {last_message}"


@app.post("/v1/chat/completions", response_model=ChatResponse)
async def chat_completions(request: ChatRequest):
    """
    OpenAI-compatible chat completions endpoint.

    LockStock middleware automatically:
    1. Advances chain before inference
    2. Adds X-LockStock-* headers to response
    3. Logs to audit trail
    """
    try:
        # Generate response (replace with your model)
        response_text = await generate_response(request.messages)

        # Build response
        return ChatResponse(
            id=f"chatcmpl-{int(time.time())}",
            created=int(time.time()),
            model="lockstock-demo",
            choices=[
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": response_text},
                    "finish_reason": "stop",
                }
            ],
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
async def health():
    """
    Health check endpoint (not protected by LockStock).

    Use paths parameter in middleware to exclude:
    app.add_middleware(LockStockMiddleware, paths=["/v1/chat/completions"])
    """
    return {"status": "healthy"}


@app.get("/")
async def root():
    """Root endpoint with usage instructions."""
    return {
        "service": "LockStock Inference Server",
        "version": "1.0.0",
        "endpoints": {
            "/v1/chat/completions": "OpenAI-compatible chat endpoint (protected)",
            "/health": "Health check (not protected)",
        },
        "lockstock": {
            "enabled": bool(os.getenv("LOCKSTOCK_AGENT_ID")),
            "agent_id": os.getenv("LOCKSTOCK_AGENT_ID", "not set"),
        },
        "usage": """
curl -X POST http://localhost:8000/v1/chat/completions \\
  -H "Content-Type: application/json" \\
  -d '{"messages": [{"role": "user", "content": "Hello"}]}'
        """.strip(),
    }


if __name__ == "__main__":
    print("=" * 70)
    print("LockStock FastAPI Inference Server")
    print("=" * 70)
    print()

    agent_id = os.getenv("LOCKSTOCK_AGENT_ID")
    if agent_id:
        print(f"✓ LockStock enabled")
        print(f"  Agent ID: {agent_id}")
        print(f"  Socket: {os.getenv('LOCKSTOCK_SOCKET', '/var/run/lockstock-guard/guard.sock')}")
        print(f"  Task: {os.getenv('LOCKSTOCK_TASK', 'network')}")
        print(f"  Block on failure: {os.getenv('LOCKSTOCK_BLOCK_ON_FAILURE', 'false')}")
    else:
        print("⚠ LockStock DISABLED")
        print("  Set LOCKSTOCK_AGENT_ID to enable chain authentication")

    print()
    print("Server starting on http://0.0.0.0:8000")
    print("Docs: http://localhost:8000/docs")
    print("=" * 70)
    print()

    uvicorn.run(app, host="0.0.0.0", port=8000)
