src/inference/api_server.py

"""
OpenMind API Server - OpenAI-Compatible Chat Completions API.

Serves the OpenMind model with:
- POST /v1/chat/completions (streaming + non-streaming)
- POST /v1/completions (legacy text completion)
- GET /v1/models (list available models)
- GET /health (health check)
- Static file serving for frontend at "/"

Fully compatible with OpenAI client libraries.

HF Spaces Deployment:
  - Listens on 0.0.0.0:7860 (required by HF Spaces)
  - MODEL_PATH env var controls weight loading
  - Rate limited: 5 requests/min per IP (via slowapi)
  - Input validation: max 500 characters per message
"""

import os
import sys
import json
import time
import uuid
import asyncio
import argparse
from pathlib import Path
from contextlib import asynccontextmanager
from typing import Optional, AsyncGenerator

from dotenv import load_dotenv

# Load .env if present (no-op in production / HF Spaces)
load_dotenv()

import torch
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field, field_validator
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded

sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))

from src.models.modeling_openmind import OpenMindModel
from src.models.config_openmind import OpenMindConfig
from src.data.tokenizer import BPETokenizer
from src.data.chat_templates import format_chat, SYSTEM_DEFAULT


# ─── Rate Limiter ──────────────────────────────────────────────────────────────

limiter = Limiter(key_func=get_remote_address)


# ─── Request/Response Models ──────────────────────────────────────────────────

class ChatMessage(BaseModel):
    role: str = "user"
    content: str = ""

    @field_validator("content")
    @classmethod
    def content_max_length(cls, v: str) -> str:
        if len(v) > 500:
            raise ValueError(
                f"Message content exceeds maximum length of 500 characters "
                f"(got {len(v)})."
            )
        return v


class ChatCompletionRequest(BaseModel):
    model: str = "openmind-125m"
    messages: list[ChatMessage]
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    top_p: float = Field(default=0.9, ge=0.0, le=1.0)
    top_k: int = Field(default=50, ge=0)
    max_tokens: int = Field(default=512, ge=1, le=4096)
    stream: bool = False
    stop: Optional[list[str]] = None
    presence_penalty: float = 0.0
    frequency_penalty: float = 0.0
    repetition_penalty: float = Field(default=1.15, ge=0.0)
    template: Optional[str] = "auto"


class CompletionRequest(BaseModel):
    model: str = "openmind-125m"
    prompt: str = Field(..., max_length=500)
    temperature: float = 0.7
    top_p: float = 0.9
    top_k: int = 50
    max_tokens: int = 256
    stream: bool = False
    stop: Optional[list[str]] = None
    repetition_penalty: float = 1.15


class ChatCompletionChoice(BaseModel):
    index: int = 0
    message: ChatMessage
    finish_reason: str = "stop"


class ChatCompletionResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: list[ChatCompletionChoice]
    usage: dict


class ModelInfo(BaseModel):
    id: str
    object: str = "model"
    created: int
    owned_by: str = "openmind"


# ─── Model Manager ────────────────────────────────────────────────────────────

class HFTokenizerWrapper:
    """Wrapper around HuggingFace tokenizer to match BPETokenizer API."""

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.eos_token_id = tokenizer.eos_token_id
        self.vocab_size = tokenizer.vocab_size

    def encode(self, text, allowed_special=None):
        return self.tokenizer.encode(text)

    def decode(self, ids):
        return self.tokenizer.decode(ids, skip_special_tokens=True)


class ModelManager:
    """Manages model loading and inference."""

    def __init__(self):
        self.model: Optional[OpenMindModel] = None
        self.tokenizer: Optional[BPETokenizer] = None
        self.model_name: str = ""
        self.device: str = "cpu"
        self.chat_template: str = "chat"

    def load(self, model_path: str, device: str = None):
        """Load model and tokenizer.

        model_path must be the path to the .pt weights FILE
        (e.g. "./weights/model.pt"). The parent directory is derived
        from it and used for config.json / tokenizer lookup.

          model_file = Path(model_path)          # ./weights/model.pt
          model_dir  = Path(model_path).parent   # ./weights/
        """
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = device

        # Split into the weights FILE and its parent DIRECTORY
        model_file = Path(model_path)
        model_dir  = model_file.parent  # e.g. ./weights/

        print(f"Loading model weights from {model_file} (dir: {model_dir}) on {device}...")

        # Load weights from the .pt file; config.json is read from model_dir
        self.model = OpenMindModel.from_pretrained(str(model_dir), device=device)
        self.model.eval()
        # Use the directory name as the model label (e.g. "weights" or "openmind-125m")
        self.model_name = model_dir.name

        # Auto-detect prompt template based on the weights path string
        path_str = str(model_path).lower()
        if "sft" in path_str or "aligned" in path_str:
            self.chat_template = "chat"
            print("Auto-detected SFT/Aligned model: default chat template set to 'chat'")
        else:
            self.chat_template = "alpaca"
            print("Auto-detected Base model: default chat template set to 'alpaca' (Instruction-Tuning)")

        # Load tokenizer — search inside model_dir (not the .pt file path)
        if self.model.config.vocab_size == 50257:
            from transformers import AutoTokenizer
            print("Loading HuggingFace GPT-2 tokenizer...")
            hf_tokenizer = AutoTokenizer.from_pretrained("gpt2")
            self.tokenizer = HFTokenizerWrapper(hf_tokenizer)
        else:
            tokenizer_dir = model_dir / "tokenizer"
            if tokenizer_dir.exists():
                self.tokenizer = BPETokenizer.load(str(tokenizer_dir))
            else:
                # Search model_dir for a custom vocab file
                for f in os.listdir(str(model_dir)):
                    if f.endswith("_vocab.json"):
                        name = f.replace("_vocab.json", "")
                        self.tokenizer = BPETokenizer.load(str(model_dir), name)
                        break

            if self.tokenizer is None:
                print("Warning: No tokenizer found, creating default")
                self.tokenizer = BPETokenizer(vocab_size=32000)

        print(f"Model '{self.model_name}' loaded successfully!")

    def generate_text(
        self,
        prompt: str,
        max_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 50,
        repetition_penalty: float = 1.0,
        stop: Optional[list[str]] = None,
    ) -> str:
        """Generate text from a prompt."""
        input_ids = self.tokenizer.encode(prompt, allowed_special={"all"})
        input_tensor = torch.tensor([input_ids], dtype=torch.long).to(self.device)

        output_ids = self.model.generate(
            input_tensor,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            eos_token_id=self.tokenizer.eos_token_id,
            repetition_penalty=repetition_penalty,
        )

        # Decode only the generated tokens
        generated_ids = output_ids[0, len(input_ids):].tolist()
        response_text = self.tokenizer.decode(generated_ids)

        if stop:
            for stop_seq in stop:
                if stop_seq in response_text:
                    response_text = response_text.split(stop_seq)[0]

        return response_text

    async def stream_generate(
        self,
        prompt: str,
        max_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 50,
        repetition_penalty: float = 1.0,
        stop: Optional[list[str]] = None,
    ) -> AsyncGenerator[str, None]:
        """Stream-generate tokens one at a time."""
        input_ids = self.tokenizer.encode(prompt, allowed_special={"all"})
        input_tensor = torch.tensor([input_ids], dtype=torch.long).to(self.device)

        past_key_values = [None] * self.model.config.n_layers
        generated = input_tensor
        generated_text = ""

        for _ in range(max_tokens):
            if past_key_values[0] is not None:
                curr_input = generated[:, -1:]
            else:
                curr_input = generated

            with torch.no_grad():
                outputs = self.model(
                    curr_input,
                    past_key_values=past_key_values,
                    use_cache=True,
                )

            logits = outputs["logits"][:, -1, :]
            past_key_values = outputs["past_key_values"]

            # Apply repetition penalty
            if repetition_penalty != 1.0:
                for i in range(logits.shape[0]):
                    for token_id in set(generated[i].tolist()):
                        logit = logits[i, token_id].item()
                        if logit < 0:
                            logits[i, token_id] = logit * repetition_penalty
                        else:
                            logits[i, token_id] = logit / repetition_penalty

            # Apply temperature
            logits = logits / max(temperature, 1e-8)

            # Top-k filtering
            if top_k > 0:
                top_k_vals = torch.topk(logits, min(top_k, logits.size(-1)))
                mask = logits < top_k_vals.values[..., -1, None]
                logits[mask] = float("-inf")

            # Top-p filtering
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(
                    torch.softmax(sorted_logits, dim=-1), dim=-1
                )
                sorted_remove = cumulative_probs > top_p
                sorted_remove[..., 1:] = sorted_remove[..., :-1].clone()
                sorted_remove[..., 0] = 0
                remove = sorted_remove.scatter(1, sorted_indices, sorted_remove)
                logits[remove] = float("-inf")

            probs = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)

            generated = torch.cat([generated, next_token], dim=-1)

            token_id = next_token[0, 0].item()
            if token_id == self.tokenizer.eos_token_id:
                break

            token_text = self.tokenizer.decode([token_id])

            potential_text = generated_text + token_text
            stopped = False
            if stop:
                for stop_seq in stop:
                    if stop_seq in potential_text:
                        stop_idx = potential_text.find(stop_seq)
                        yield potential_text[len(generated_text):stop_idx]
                        stopped = True
                        break
            if stopped:
                break

            generated_text = potential_text
            yield token_text

            # Small delay for streaming effect
            await asyncio.sleep(0)


# ─── Global model manager ─────────────────────────────────────────────────────

manager = ModelManager()


# ─── App lifespan (load model at startup) ─────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load model on startup using MODEL_PATH env var.

    MODEL_PATH must point to the weights FILE (e.g. ./weights/model.pt).
    The parent directory (./weights/) is where config.json lives.
    """
    model_path = os.getenv("MODEL_PATH", "./weights/model.pt")
    model_file = Path(model_path)

    # Check that the weights FILE exists (not just the directory)
    if model_file.exists() and model_file.is_file():
        try:
            manager.load(str(model_file))
        except Exception as exc:
            print(f"[WARNING] Could not load model from {model_file}: {exc}")
    else:
        print(
            f"[INFO] MODEL_PATH={model_path!r} — weights file not found at startup. "
            "Model will not be loaded until weights are present."
        )
    yield
    # Cleanup (if needed)


# ─── FastAPI Application ──────────────────────────────────────────────────────

app = FastAPI(
    title="OpenMind API",
    description="OpenAI-compatible API for the OpenMind language model",
    version="0.1.0",
    lifespan=lifespan,
)

# Attach rate limiter
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ─── Health & Model Endpoints ─────────────────────────────────────────────────

@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {
        "status": "ok",
        "model_loaded": manager.model is not None,
        "model_name": manager.model_name,
        "device": manager.device,
    }


@app.get("/v1/models")
async def list_models():
    """List available models."""
    models = []
    if manager.model is not None:
        models.append(ModelInfo(
            id=manager.model_name,
            created=int(time.time()),
        ))
    return {"object": "list", "data": [m.dict() for m in models]}


# ─── Chat Completions ─────────────────────────────────────────────────────────

@app.post("/v1/chat/completions")
@limiter.limit("5/minute")
async def chat_completions(request: Request, body: ChatCompletionRequest):
    """OpenAI-compatible chat completions endpoint (max 5 req/min per IP)."""
    if manager.model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    # Determine template
    template = body.template
    if not template or template == "auto":
        template = manager.chat_template

    # Format messages into prompt
    messages = [{"role": m.role, "content": m.content} for m in body.messages]
    prompt = format_chat(messages, add_generation_prompt=True, template=template)

    # Determine stop sequences
    stop_sequences = body.stop
    if not stop_sequences:
        if template == "chat":
            stop_sequences = ["<|user|>", "<|system|>", "<|endoftext|>"]
        elif template == "alpaca":
            stop_sequences = ["###", "Instruction:", "Response:", "<|endoftext|>"]
        else:
            stop_sequences = ["<|endoftext|>"]

    completion_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"

    if body.stream:
        return StreamingResponse(
            _stream_chat_response(
                completion_id, prompt, body, stop_sequences
            ),
            media_type="text/event-stream",
        )

    # Non-streaming response
    response_text = manager.generate_text(
        prompt,
        max_tokens=body.max_tokens,
        temperature=body.temperature,
        top_p=body.top_p,
        top_k=body.top_k,
        repetition_penalty=body.repetition_penalty,
        stop=stop_sequences,
    )

    # Count tokens (approximate)
    prompt_tokens = len(manager.tokenizer.encode(prompt, allowed_special={"all"}))
    completion_tokens = len(manager.tokenizer.encode(response_text))

    return ChatCompletionResponse(
        id=completion_id,
        created=int(time.time()),
        model=manager.model_name,
        choices=[
            ChatCompletionChoice(
                message=ChatMessage(role="assistant", content=response_text),
                finish_reason="stop",
            )
        ],
        usage={
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        },
    )


async def _stream_chat_response(
    completion_id: str,
    prompt: str,
    request: ChatCompletionRequest,
    stop_sequences: Optional[list[str]],
) -> AsyncGenerator[str, None]:
    """Generate streaming SSE response."""
    # Initial chunk with role
    chunk = {
        "id": completion_id,
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": manager.model_name,
        "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
    }
    yield f"data: {json.dumps(chunk)}\n\n"

    # Stream tokens
    async for token in manager.stream_generate(
        prompt,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_p=request.top_p,
        top_k=request.top_k,
        repetition_penalty=request.repetition_penalty,
        stop=stop_sequences,
    ):
        chunk = {
            "id": completion_id,
            "object": "chat.completion.chunk",
            "created": int(time.time()),
            "model": manager.model_name,
            "choices": [{"index": 0, "delta": {"content": token}, "finish_reason": None}],
        }
        yield f"data: {json.dumps(chunk)}\n\n"

    # Final chunk
    chunk = {
        "id": completion_id,
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": manager.model_name,
        "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
    }
    yield f"data: {json.dumps(chunk)}\n\n"
    yield "data: [DONE]\n\n"


# ─── Legacy Text Completions ──────────────────────────────────────────────────

@app.post("/v1/completions")
@limiter.limit("5/minute")
async def text_completions(request: Request, body: CompletionRequest):
    """Legacy text completion endpoint (max 5 req/min per IP)."""
    if manager.model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    # For text completions, default stop sequence is <|endoftext|> if not provided
    stop_sequences = body.stop or ["<|endoftext|>"]

    response_text = manager.generate_text(
        body.prompt,
        max_tokens=body.max_tokens,
        temperature=body.temperature,
        top_p=body.top_p,
        top_k=body.top_k,
        repetition_penalty=body.repetition_penalty,
        stop=stop_sequences,
    )

    return {
        "id": f"cmpl-{uuid.uuid4().hex[:8]}",
        "object": "text_completion",
        "created": int(time.time()),
        "model": manager.model_name,
        "choices": [{"text": response_text, "index": 0, "finish_reason": "stop"}],
        "usage": {
            "prompt_tokens": len(manager.tokenizer.encode(body.prompt)),
            "completion_tokens": len(manager.tokenizer.encode(response_text)),
        },
    }


# ─── Frontend Static Files ────────────────────────────────────────────────────

def setup_static_files():
    """Mount frontend/ as static files at '/' (called after route registration)."""
    # Resolve frontend directory relative to the project root
    project_root = Path(__file__).resolve().parent.parent.parent
    frontend_dir = project_root / "frontend"

    if frontend_dir.exists():
        # Serve CSS and JS subdirectories
        for subdir in ["css", "js", "assets"]:
            asset_dir = frontend_dir / subdir
            if asset_dir.exists():
                app.mount(f"/{subdir}", StaticFiles(directory=str(asset_dir)), name=subdir)

        # Serve root index.html at "/"
        @app.get("/")
        async def serve_frontend():
            index_path = frontend_dir / "index.html"
            if index_path.exists():
                return FileResponse(str(index_path))
            return {"message": "OpenMind API is running"}

        print(f"Serving frontend from {frontend_dir}")
    else:
        print(f"[INFO] Frontend directory not found at {frontend_dir}. Skipping static serving.")


# Register static files
setup_static_files()


# ─── Entry Point ──────────────────────────────────────────────────────────────

def start_server(
    model_path: str = None,
    host: str = None,
    port: int = None,
    device: str = None,
    chat_template: str = "auto",
):
    """Start the API server.

    All parameters fall back to environment variables:
      MODEL_PATH  — path to weights (default: ./weights/model.pt)
      SERVER_HOST — bind host (default: 0.0.0.0)
      SERVER_PORT — bind port (default: 7860)
    """
    if model_path is None:
        model_path = os.getenv("MODEL_PATH", "./weights/model.pt")
    if host is None:
        host = os.getenv("SERVER_HOST", "0.0.0.0")
    if port is None:
        port = int(os.getenv("SERVER_PORT", "7860"))

    if chat_template != "auto" and manager.model is not None:
        manager.chat_template = chat_template

    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="OpenMind API Server")
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="Path to model directory/file (overrides MODEL_PATH env var)",
    )
    parser.add_argument("--host", type=str, default=None, help="Bind host (default: 0.0.0.0)")
    parser.add_argument("--port", type=int, default=None, help="Bind port (default: 7860)")
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument(
        "--chat-template",
        type=str,
        default="auto",
        choices=["auto", "chat", "alpaca", "raw"],
        help="Chat template override",
    )
    args = parser.parse_args()

    # Allow --model to override env var
    if args.model:
        os.environ["MODEL_PATH"] = args.model

    start_server(
        model_path=args.model,
        host=args.host,
        port=args.port,
        device=args.device,
        chat_template=args.chat_template,
    )