server.py

#!/usr/bin/env python3
"""
BlitzKode backend server.

Exposes a FastAPI backend for local GGUF inference through llama.cpp.
Model is loaded lazily so the module stays importable in tests and
environments where the model artifact is not present yet.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import re
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from collections import deque
from collections.abc import Callable
from contextlib import asynccontextmanager, suppress
from dataclasses import dataclass
from dataclasses import field as dataclass_field
from html.parser import HTMLParser
from pathlib import Path
from typing import Any, Literal, cast

import llama_cpp
import uvicorn
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field
from starlette.middleware.base import BaseHTTPMiddleware

APP_NAME = "BlitzKode"
APP_VERSION = "2.0"
CREATOR = "Sajad"
ROOT_DIR = Path(__file__).resolve().parent
DEFAULT_MODEL_PATH = ROOT_DIR / "blitzkode.gguf"
DEFAULT_CONTEXT = 2048
DEFAULT_MAX_PROMPT_LENGTH = 4000
DEFAULT_MAX_TOKENS = 512
DEFAULT_RATE_LIMIT_MAX = 30
DEFAULT_MAX_SEARCH_RESULTS = 5
DEFAULT_SEARCH_TIMEOUT_SECONDS = 8
DEFAULT_SEARCH_CACHE_TTL_SECONDS = 300
DEFAULT_MAX_MESSAGES = 20
DEFAULT_BATCH = 256
DEFAULT_PROMPT_CACHE_BYTES = 64 * 1024 * 1024
STOP_TOKENS = ["<|im_end|>", "<|im_start|>user"]

SYSTEM_PROMPT = (
    "<|im_start|>system\n"
    "You are BlitzKode, an AI coding assistant created by Sajad. "
    "You are an expert in Python, JavaScript, Java, C++, and other programming languages. "
    "For coding work, first identify the user's goal, constraints, and any unknowns. "
    "If asked about a library, API, file, function, citation, execution result, or repository detail that is not provided, "
    "do not fabricate it. Say you do not know and explain how to verify it. "
    "Prefer safe minimal fixes over speculative code. "
    "Write clean, efficient, and well-documented code. Keep responses concise and practical.<|im_end|>"
)

logger = logging.getLogger("blitzkode")


def _bool_from_env(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def _int_from_env(name: str, default: int) -> int:
    value = os.getenv(name)
    if not value:
        return default
    try:
        return int(value)
    except ValueError:
        return default


def _path_from_env(name: str, default: Path) -> Path:
    value = os.getenv(name)
    return Path(value) if value else default


def _validate_prompt(prompt: str, max_length: int) -> tuple[str, JSONResponse | None]:
    prompt = prompt.strip()
    if not prompt:
        return prompt, JSONResponse({"error": "Prompt is required"}, status_code=400)
    if len(prompt) > max_length:
        return prompt, JSONResponse(
            {"error": f"Prompt too long. Max {max_length} chars."},
            status_code=400,
        )
    return prompt, None


_SIGNATURE_QUERY_RE = re.compile(
    r"(?:signature|how\s+(?:do|can)\s+i\s+use|usage\s+of|docs?\s+for).{0,120}?"
    r"(?P<symbol>[A-Za-z_][\w.]*\s*\([^)]*\)|[A-Za-z_][\w.]*\s+function)",
    re.IGNORECASE | re.DOTALL,
)
_CODE_CONTEXT_RE = re.compile(r"```|\b(def|class|function|interface|type|import|from)\b|\{\s*\"|\bsource\s+code\b", re.IGNORECASE)


def _grounding_guard_response(prompt: str, has_external_context: bool = False) -> str | None:
    """Prevents confident fabricated API/function signatures when no source/docs are provided.

    Small local models can ignore system instructions around unknown symbols. This guardrail only triggers for direct API/signature
    lookup questions that contain no pasted source/docs context. It does not block normal code-generation tasks.
    """
    if has_external_context or _CODE_CONTEXT_RE.search(prompt):
        return None

    match = _SIGNATURE_QUERY_RE.search(prompt)
    if not match:
        return None

    symbol = " ".join(match.group("symbol").split())
    return (
        f"I don't have enough verified context to know the signature or usage of `{symbol}`. "
        "Please provide the source code or official documentation, or enable research mode so I can ground the answer."
    )


@dataclass(slots=True)
class Settings:
    root_dir: Path = ROOT_DIR
    model_path: Path = dataclass_field(default_factory=lambda: _path_from_env("BLITZKODE_MODEL_PATH", DEFAULT_MODEL_PATH))
    host: str = os.getenv("BLITZKODE_HOST", "0.0.0.0")
    port: int = _int_from_env("BLITZKODE_PORT", 7860)
    n_gpu_layers: int = _int_from_env("BLITZKODE_GPU_LAYERS", 0)
    n_ctx: int = _int_from_env("BLITZKODE_N_CTX", DEFAULT_CONTEXT)
    n_threads: int = _int_from_env("BLITZKODE_THREADS", max(1, min(8, os.cpu_count() or 1)))
    n_threads_batch: int = _int_from_env("BLITZKODE_THREADS_BATCH", max(1, min(8, os.cpu_count() or 1)))
    n_batch: int = _int_from_env("BLITZKODE_BATCH", DEFAULT_BATCH)
    n_ubatch: int = _int_from_env("BLITZKODE_UBATCH", min(DEFAULT_BATCH, 128))
    prompt_cache_enabled: bool = _bool_from_env("BLITZKODE_PROMPT_CACHE", default=True)
    prompt_cache_bytes: int = _int_from_env("BLITZKODE_PROMPT_CACHE_BYTES", DEFAULT_PROMPT_CACHE_BYTES)
    use_mmap: bool = _bool_from_env("BLITZKODE_USE_MMAP", default=True)
    use_mlock: bool = _bool_from_env("BLITZKODE_USE_MLOCK", default=False)
    offload_kqv: bool = _bool_from_env("BLITZKODE_OFFLOAD_KQV", default=True)
    max_prompt_length: int = _int_from_env("BLITZKODE_MAX_PROMPT_LENGTH", DEFAULT_MAX_PROMPT_LENGTH)
    preload_model: bool = _bool_from_env("BLITZKODE_PRELOAD_MODEL", default=False)
    cors_origins: str = os.getenv("BLITZKODE_CORS_ORIGINS", "http://localhost:7860")
    api_key: str = os.getenv("BLITZKODE_API_KEY", "")
    web_search_enabled: bool = _bool_from_env("BLITZKODE_WEB_SEARCH", default=True)
    search_timeout_seconds: int = _int_from_env("BLITZKODE_SEARCH_TIMEOUT", DEFAULT_SEARCH_TIMEOUT_SECONDS)
    max_search_results: int = _int_from_env("BLITZKODE_MAX_SEARCH_RESULTS", DEFAULT_MAX_SEARCH_RESULTS)
    search_cache_ttl_seconds: int = _int_from_env("BLITZKODE_SEARCH_CACHE_TTL", DEFAULT_SEARCH_CACHE_TTL_SECONDS)


class MessageItem(BaseModel):
    role: Literal["user", "assistant"]
    content: str = Field(min_length=1, max_length=DEFAULT_MAX_PROMPT_LENGTH)


class GenerateRequest(BaseModel):
    prompt: str
    messages: list[MessageItem] = Field(default_factory=list, max_length=DEFAULT_MAX_MESSAGES)
    temperature: float = Field(default=0.5, ge=0.0, le=2.0)
    max_tokens: int = Field(default=256, ge=1, le=DEFAULT_MAX_TOKENS)
    top_p: float = Field(default=0.95, gt=0.0, le=1.0)
    top_k: int = Field(default=20, ge=1, le=200)
    repeat_penalty: float = Field(default=1.05, ge=0.8, le=2.0)


class SearchRequest(BaseModel):
    query: str = Field(min_length=1, max_length=500)
    max_results: int = Field(default=DEFAULT_MAX_SEARCH_RESULTS, ge=1, le=10)
    deep: bool = False


class ResearchGenerateRequest(GenerateRequest):
    search_query: str | None = Field(default=None, max_length=500)
    search_results: int = Field(default=DEFAULT_MAX_SEARCH_RESULTS, ge=1, le=10)
    deep_search: bool = False


@dataclass(slots=True)
class SearchResult:
    title: str
    url: str
    snippet: str
    source: str = "DuckDuckGo"

    def as_dict(self) -> dict[str, str]:
        return {
            "title": self.title,
            "url": self.url,
            "snippet": self.snippet,
            "source": self.source,
        }


class DuckDuckGoHTMLParser(HTMLParser):
    def __init__(self, max_results: int):
        super().__init__(convert_charrefs=True)
        self.max_results = max_results
        self.results: list[dict[str, str]] = []
        self._active_field: Literal["title", "snippet"] | None = None
        self._active_href = ""
        self._text_parts: list[str] = []

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        if tag != "a":
            return

        attr_map = {name: value or "" for name, value in attrs}
        classes = set(attr_map.get("class", "").split())
        if "result__a" in classes:
            if len(self.results) >= self.max_results:
                return
            self._active_field = "title"
        elif "result__snippet" in classes:
            self._active_field = "snippet"
        else:
            return

        self._active_href = attr_map.get("href", "")
        self._text_parts = []

    def handle_data(self, data: str) -> None:
        if self._active_field:
            self._text_parts.append(data)

    def handle_endtag(self, tag: str) -> None:
        if tag != "a" or not self._active_field:
            return

        text = " ".join("".join(self._text_parts).split())
        url = self._unwrap_result_url(self._active_href)

        if self._active_field == "title" and text and url and len(self.results) < self.max_results:
            self.results.append({"title": text, "url": url, "snippet": ""})
        elif self._active_field == "snippet" and text and self.results:
            target = next((item for item in reversed(self.results) if not item["snippet"]), None)
            if target:
                target["snippet"] = text

        self._active_field = None
        self._active_href = ""
        self._text_parts = []

    @staticmethod
    def _unwrap_result_url(href: str) -> str:
        href = href.strip()
        if href.startswith("//"):
            href = f"https:{href}"

        parsed = urllib.parse.urlparse(href)
        if "duckduckgo.com" in parsed.netloc and parsed.path.startswith("/l/"):
            target = urllib.parse.parse_qs(parsed.query).get("uddg", [""])[0]
            if target:
                return target
        return href


class WebSearchService:
    def __init__(self, settings: Settings):
        self.settings = settings
        self._cache: dict[tuple[str, int, bool], tuple[float, list[dict[str, str]]]] = {}
        self._cache_lock = threading.Lock()

    @property
    def enabled(self) -> bool:
        return self.settings.web_search_enabled

    def _query_variants(self, query: str, deep: bool) -> list[str]:
        query = " ".join(query.split())
        if not deep:
            return [query]
        return [
            query,
            f"{query} official documentation",
            f"{query} best practices",
        ]

    def _append_result(
        self, results: list[SearchResult], seen_urls: set[str], title: str, url: str, snippet: str, max_results: int
    ) -> None:
        title = " ".join((title or "Untitled").split())[:200]
        url = (url or "").strip()
        snippet = " ".join((snippet or "").split())[:500]
        if not url or url in seen_urls or len(results) >= max_results:
            return
        seen_urls.add(url)
        results.append(SearchResult(title=title, url=url, snippet=snippet))

    def _collect_related_topics(self, topics: list[dict], results: list[SearchResult], seen_urls: set[str], max_results: int) -> None:
        for topic in topics:
            if len(results) >= max_results:
                return
            if "Topics" in topic:
                self._collect_related_topics(topic.get("Topics", []), results, seen_urls, max_results)
                continue
            text = topic.get("Text", "")
            url = topic.get("FirstURL", "")
            if text and url:
                title = text.split(" - ", 1)[0]
                self._append_result(results, seen_urls, title, url, text, max_results)

    def _read_search_payload(self, request: urllib.request.Request) -> dict[str, Any]:
        with urllib.request.urlopen(request, timeout=self.settings.search_timeout_seconds) as response:
            raw = response.read().decode("utf-8")
        try:
            payload = json.loads(raw)
        except json.JSONDecodeError as exc:
            raise RuntimeError("Search provider returned an invalid JSON response") from exc
        if not isinstance(payload, dict):
            raise RuntimeError("Search provider returned an unexpected response shape")
        return payload

    def _search_html(
        self,
        query: str,
        results: list[SearchResult],
        seen_urls: set[str],
        max_results: int,
    ) -> None:
        params = urllib.parse.urlencode({"q": query})
        request = urllib.request.Request(
            f"https://html.duckduckgo.com/html/?{params}",
            headers={"User-Agent": f"Mozilla/5.0 {APP_NAME}/{APP_VERSION}"},
        )
        with urllib.request.urlopen(request, timeout=self.settings.search_timeout_seconds) as response:
            raw = response.read().decode("utf-8", errors="replace")

        parser = DuckDuckGoHTMLParser(max_results)
        parser.feed(raw)
        for item in parser.results:
            self._append_result(results, seen_urls, item["title"], item["url"], item["snippet"], max_results)

    def search(self, query: str, max_results: int = DEFAULT_MAX_SEARCH_RESULTS, deep: bool = False) -> list[dict[str, str]]:
        if not self.enabled:
            raise RuntimeError("Web search is disabled. Set BLITZKODE_WEB_SEARCH=true to enable it.")

        query = " ".join(query.split())
        if not query:
            raise ValueError("Search query is required")

        limit = min(max_results, max(1, self.settings.max_search_results), 10)
        cache_key = (query.lower(), limit, deep)
        now = time.monotonic()
        with self._cache_lock:
            cached = self._cache.get(cache_key)
            if cached and now - cached[0] < self.settings.search_cache_ttl_seconds:
                return [dict(item) for item in cached[1]]

        results: list[SearchResult] = []
        seen_urls: set[str] = set()

        for variant in self._query_variants(query, deep):
            if len(results) >= limit:
                break
            params = urllib.parse.urlencode(
                {
                    "q": variant,
                    "format": "json",
                    "no_html": "1",
                    "skip_disambig": "1",
                }
            )
            request = urllib.request.Request(
                f"https://api.duckduckgo.com/?{params}",
                headers={"User-Agent": f"{APP_NAME}/{APP_VERSION}"},
            )
            try:
                payload = self._read_search_payload(request)
            except RuntimeError as exc:
                result_count = len(results)
                self._search_html(variant, results, seen_urls, limit)
                if len(results) == result_count:
                    raise exc
                continue

            self._append_result(
                results,
                seen_urls,
                payload.get("Heading") or variant,
                payload.get("AbstractURL", ""),
                payload.get("AbstractText", ""),
                limit,
            )
            self._collect_related_topics(payload.get("RelatedTopics", []), results, seen_urls, limit)

        search_payload = [result.as_dict() for result in results]
        with self._cache_lock:
            self._cache[cache_key] = (time.monotonic(), search_payload)
        return [dict(item) for item in search_payload]


class ModelService:
    def __init__(self, settings: Settings):
        self.settings = settings
        self._llm: llama_cpp.Llama | None = None
        self._init_lock = threading.Lock()
        self._load_time_seconds: float | None = None
        self._last_error: str | None = None
        self._busy: bool = False

    @property
    def model_loaded(self) -> bool:
        return self._llm is not None

    @property
    def model_exists(self) -> bool:
        return self.settings.model_path.exists()

    @property
    def last_error(self) -> str | None:
        return self._last_error

    @property
    def load_time_seconds(self) -> float | None:
        return self._load_time_seconds

    @property
    def busy(self) -> bool:
        return self._busy

    def load_model(self):
        if self._llm is not None:
            return self._llm

        with self._init_lock:
            if self._llm is not None:
                return self._llm

            if not self.model_exists:
                self._last_error = f"Model not found at {self.settings.model_path}"
                raise FileNotFoundError(self._last_error)

            start_time = time.perf_counter()
            try:
                self._llm = self._create_llama()
                self._load_time_seconds = time.perf_counter() - start_time
                self._last_error = None
                self._configure_prompt_cache(self._llm)
                logger.info(
                    "Model loaded in %.2fs (gpu_layers=%d, ctx=%d, threads=%d, batch=%d)",
                    self._load_time_seconds,
                    self.settings.n_gpu_layers,
                    self.settings.n_ctx,
                    self.settings.n_threads,
                    self.settings.n_batch,
                )
            except Exception as exc:
                self._last_error = str(exc)
                logger.error("Model load failed: %s", exc)
                raise

        return self._llm

    def _create_llama(self) -> llama_cpp.Llama:
        kwargs: dict[str, Any] = {
            "model_path": str(self.settings.model_path),
            "n_gpu_layers": self.settings.n_gpu_layers,
            "n_ctx": self.settings.n_ctx,
            "n_threads": self.settings.n_threads,
            "n_threads_batch": self.settings.n_threads_batch,
            "n_batch": self.settings.n_batch,
            "n_ubatch": self.settings.n_ubatch,
            "offload_kqv": self.settings.offload_kqv,
            "verbose": False,
            "use_mmap": self.settings.use_mmap,
            "use_mlock": self.settings.use_mlock,
            "seed": -1,
        }
        try:
            return llama_cpp.Llama(**kwargs)
        except TypeError as exc:
            message = str(exc)
            unsupported = [key for key in ("n_threads_batch", "n_ubatch", "offload_kqv") if key in message]
            if not unsupported:
                raise
            for key in unsupported:
                kwargs.pop(key, None)
            logger.warning("Retrying model load without unsupported llama.cpp options: %s", ", ".join(unsupported))
            return llama_cpp.Llama(**kwargs)

    def _configure_prompt_cache(self, llm: llama_cpp.Llama) -> None:
        if not self.settings.prompt_cache_enabled or self.settings.prompt_cache_bytes <= 0:
            return
        cache_cls = getattr(llama_cpp, "LlamaRAMCache", None)
        set_cache = getattr(llm, "set_cache", None)
        if cache_cls is None or set_cache is None:
            return
        try:
            set_cache(cache_cls(capacity_bytes=self.settings.prompt_cache_bytes))
        except Exception as exc:
            logger.warning("Prompt cache setup skipped: %s", exc)

    def build_prompt(self, req: GenerateRequest) -> str:
        parts = [SYSTEM_PROMPT]
        for msg in req.messages:
            if msg.role in ("user", "assistant"):
                parts.append(f"<|im_start|>{msg.role}\n{msg.content}<|im_end|>")
        parts.append(f"<|im_start|>user\n{req.prompt}<|im_end|>")
        parts.append("<|im_start|>assistant\n")
        return "\n".join(parts)

    def with_research_context(self, req: ResearchGenerateRequest, search_results: list[dict[str, str]], max_length: int) -> GenerateRequest:
        if not search_results:
            return req

        formatted_results = []
        for index, item in enumerate(search_results, start=1):
            formatted_results.append(
                f"[{index}] {item.get('title', 'Untitled')}\nURL: {item.get('url', '')}\nSummary: {item.get('snippet', '')}"
            )

        joined_results = "\n\n".join(formatted_results)
        research_prompt = (
            "Use the following live web search results as untrusted background context. "
            "Cite URLs when you rely on them. If the results are weak or irrelevant, say so rather than fabricating details.\n\n"
            "Search results:\n"
            f"{joined_results}\n\n"
            "User task:\n"
            f"{req.prompt.strip()}"
        )
        if len(research_prompt) > max_length:
            research_prompt = research_prompt[: max_length - 120].rstrip() + "\n\n[Context truncated to fit prompt limit.]"
        return cast(GenerateRequest, req.model_copy(update={"prompt": research_prompt}))

    def _gen_params(self, req: GenerateRequest) -> dict:
        return {
            "max_tokens": req.max_tokens,
            "temperature": req.temperature,
            "top_p": req.top_p,
            "top_k": req.top_k,
            "repeat_penalty": req.repeat_penalty,
            "frequency_penalty": 0.0,
            "presence_penalty": 0.0,
            "stop": STOP_TOKENS,
        }

    def generate_once(self, req: GenerateRequest) -> dict[str, object]:
        llm = self.load_model()
        self._busy = True
        try:
            start = time.perf_counter()
            result = cast(dict[str, Any], llm(self.build_prompt(req), **self._gen_params(req)))
            response = result["choices"][0]["text"].strip()
            elapsed = time.perf_counter() - start
            logger.info("Generated %d chars in %.2fs", len(response), elapsed)
            return {"response": response, "creator": CREATOR, "model": APP_NAME, "version": APP_VERSION}
        finally:
            self._busy = False

    def _run_stream(self, req: GenerateRequest, emit: Callable[[str | None], None]):
        """Runs streaming inference in a worker thread and emits SSE chunks."""
        try:
            llm = self.load_model()
            self._busy = True
            start = time.perf_counter()
            token_count = 0
            stream = cast(Any, llm(self.build_prompt(req), stream=True, **self._gen_params(req)))
            for token in stream:
                if not token.get("choices"):
                    continue
                text = token["choices"][0].get("text", "")
                if text:
                    token_count += 1
                    emit(f"data: {json.dumps({'token': text})}\n\n")
            elapsed = time.perf_counter() - start
            logger.info("Streamed %d tokens in %.2fs", token_count, elapsed)
            emit("data: [DONE]\n\n")
        except Exception as exc:
            logger.error("Stream error: %s", exc)
            emit(f"data: {json.dumps({'error': str(exc)})}\n\n")
        finally:
            self._busy = False
            emit(None)


def _check_api_key(request: Request, settings: Settings) -> JSONResponse | None:
    if not settings.api_key:
        return None
    auth = request.headers.get("Authorization", "")
    token = auth[7:] if auth.startswith("Bearer ") else auth

    # Timing-safe comparison (prevent timing attacks)
    import hmac

    if not hmac.compare_digest(token, settings.api_key):
        return JSONResponse({"error": "Unauthorized"}, status_code=401)
    return None


class RateLimitMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, max_requests: int = DEFAULT_RATE_LIMIT_MAX, window_seconds: int = 60):
        super().__init__(app)
        self._max = max_requests
        self._window = window_seconds
        self._clients: dict[str, deque[float]] = {}
        self._lock = threading.Lock()
        self._cleanup_done = 0

    async def dispatch(self, request: Request, call_next):
        client_ip = request.client.host if request.client else "unknown"
        now = time.monotonic()

        # Cleanup old entries periodically (every 1000 requests)
        self._cleanup_done += 1
        if self._cleanup_done > 1000:
            self._cleanup_done = 0
            with self._lock:
                cutoff = now - self._window
                self._clients = {ip: deque(t for t in ts if t >= cutoff) for ip, ts in self._clients.items() if ts}

        with self._lock:
            timestamps = self._clients.setdefault(client_ip, deque())
            cutoff = now - self._window
            while timestamps and timestamps[0] < cutoff:
                timestamps.popleft()
            if len(timestamps) >= self._max:
                return JSONResponse(
                    {"error": "Rate limit exceeded. Try again later."},
                    status_code=429,
                    headers={"Retry-After": str(self._window)},
                )
            timestamps.append(now)
        return await call_next(request)


class RequestSizeLimitMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, max_bytes: int = 50_000):
        super().__init__(app)
        self._max = max_bytes

    async def dispatch(self, request: Request, call_next):
        content_length = request.headers.get("content-length")
        if content_length:
            try:
                if int(content_length) > self._max:
                    return JSONResponse({"error": "Request body too large"}, status_code=413)
            except ValueError:
                return JSONResponse({"error": "Invalid Content-Length header"}, status_code=400)
        return await call_next(request)


def create_app(settings: Settings | None = None) -> FastAPI:
    settings = settings or Settings()
    model_service = ModelService(settings)
    search_service = WebSearchService(settings)
    model_lock = asyncio.Lock()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(message)s", datefmt="%H:%M:%S")

    @asynccontextmanager
    async def lifespan(_: FastAPI):
        if settings.preload_model:
            with suppress(Exception):
                await asyncio.to_thread(model_service.load_model)
        yield

    app = FastAPI(title=f"{APP_NAME} API", version=APP_VERSION, lifespan=lifespan)
    app.state.settings = settings
    app.state.model_service = model_service
    app.state.search_service = search_service

    cors_origins = [o.strip() for o in settings.cors_origins.split(",") if o.strip()]
    app.add_middleware(
        CORSMiddleware,
        allow_origins=cors_origins,
        allow_methods=["POST", "GET", "OPTIONS"],
        allow_headers=["Content-Type", "Authorization"],
    )

    if _bool_from_env("BLITZKODE_RATE_LIMIT", default=True):
        app.add_middleware(RateLimitMiddleware, max_requests=_int_from_env("BLITZKODE_RATE_LIMIT_MAX", DEFAULT_RATE_LIMIT_MAX))
    app.add_middleware(RequestSizeLimitMiddleware, max_bytes=_int_from_env("BLITZKODE_MAX_REQUEST_BYTES", 50_000))

    @app.get("/")
    async def root():
        return JSONResponse(
            {
                "name": APP_NAME,
                "version": APP_VERSION,
                "message": "BlitzKode backend API is running. Use /info for endpoint details.",
            }
        )

    @app.get("/health")
    async def health():
        status = "healthy" if model_service.model_exists and not model_service.last_error else "degraded"
        return JSONResponse(
            {
                "status": status,
                "model_loaded": model_service.model_loaded,
                "model_path": str(settings.model_path),
                "model_exists": model_service.model_exists,
                "version": APP_VERSION,
                "gpu_layers": settings.n_gpu_layers,
                "last_error": model_service.last_error,
                "busy": model_service.busy,
            }
        )

    @app.post("/generate")
    async def generate(req: GenerateRequest, request: Request):
        auth_err = _check_api_key(request, settings)
        if auth_err:
            return auth_err

        prompt, err = _validate_prompt(req.prompt, settings.max_prompt_length)
        if err:
            return err

        guard_response = _grounding_guard_response(prompt)
        if guard_response:
            return JSONResponse(
                {"response": guard_response, "creator": CREATOR, "model": APP_NAME, "version": APP_VERSION, "guarded": True}
            )

        async with model_lock:
            try:
                sanitized = req.model_copy(update={"prompt": prompt})
                payload = await asyncio.to_thread(model_service.generate_once, sanitized)
                return JSONResponse(payload)
            except FileNotFoundError as exc:
                return JSONResponse({"error": str(exc)}, status_code=503)
            except Exception as exc:
                return JSONResponse({"error": str(exc)}, status_code=500)

    @app.post("/generate/research")
    async def generate_research(req: ResearchGenerateRequest, request: Request):
        auth_err = _check_api_key(request, settings)
        if auth_err:
            return auth_err

        prompt, err = _validate_prompt(req.prompt, settings.max_prompt_length)
        if err:
            return err

        if not search_service.enabled:
            return JSONResponse({"error": "Web search is disabled"}, status_code=503)

        search_query = (req.search_query or prompt).strip()
        try:
            results = await asyncio.to_thread(search_service.search, search_query, req.search_results, req.deep_search)
            sanitized = req.model_copy(update={"prompt": prompt})
            enriched = model_service.with_research_context(sanitized, results, settings.max_prompt_length)
            async with model_lock:
                payload = await asyncio.to_thread(model_service.generate_once, enriched)
            payload["search_results"] = results
            return JSONResponse(payload)
        except FileNotFoundError as exc:
            return JSONResponse({"error": str(exc)}, status_code=503)
        except (RuntimeError, urllib.error.URLError, TimeoutError) as exc:
            return JSONResponse({"error": f"Search failed: {exc}"}, status_code=502)
        except Exception as exc:
            return JSONResponse({"error": str(exc)}, status_code=500)

    @app.post("/generate/stream")
    async def generate_stream(req: GenerateRequest, request: Request):
        auth_err = _check_api_key(request, settings)
        if auth_err:
            return auth_err

        prompt, err = _validate_prompt(req.prompt, settings.max_prompt_length)
        if err:
            return err

        if not model_service.model_exists:
            return JSONResponse({"error": f"Model not found at {settings.model_path}"}, status_code=503)

        guard_response = _grounding_guard_response(prompt)
        if guard_response:

            async def _guarded_stream():
                yield f"data: {json.dumps({'token': guard_response})}\n\n"
                yield "data: [DONE]\n\n"

            return StreamingResponse(_guarded_stream(), media_type="text/event-stream")

        sanitized = req.model_copy(update={"prompt": prompt})

        async def _locked_stream():
            async with model_lock:
                loop = asyncio.get_running_loop()
                token_q: asyncio.Queue[str | None] = asyncio.Queue()

                def emit(chunk: str | None) -> None:
                    loop.call_soon_threadsafe(token_q.put_nowait, chunk)

                thread = threading.Thread(
                    target=model_service._run_stream,
                    args=(sanitized, emit),
                    daemon=True,
                )
                thread.start()
                while True:
                    chunk = await token_q.get()
                    if chunk is None:
                        break
                    yield chunk

        return StreamingResponse(
            _locked_stream(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
        )

    @app.post("/search/web")
    async def search_web(req: SearchRequest, request: Request):
        auth_err = _check_api_key(request, settings)
        if auth_err:
            return auth_err

        if not search_service.enabled:
            return JSONResponse({"error": "Web search is disabled"}, status_code=503)

        try:
            results = await asyncio.to_thread(search_service.search, req.query, req.max_results, req.deep)
            return JSONResponse({"query": req.query.strip(), "deep": req.deep, "results": results})
        except (RuntimeError, urllib.error.URLError, TimeoutError) as exc:
            return JSONResponse({"error": f"Search failed: {exc}"}, status_code=502)
        except Exception as exc:
            return JSONResponse({"error": str(exc)}, status_code=500)

    @app.get("/info")
    async def info():
        return JSONResponse(
            {
                "name": APP_NAME,
                "creator": CREATOR,
                "version": APP_VERSION,
                "status": "ready" if model_service.model_exists else "model-missing",
                "mode": f"{'GPU' if settings.n_gpu_layers > 0 else 'CPU'} (llama.cpp)",
                "gpu_layers": settings.n_gpu_layers,
                "context_window": settings.n_ctx,
                "threads": settings.n_threads,
                "threads_batch": settings.n_threads_batch,
                "batch": settings.n_batch,
                "ubatch": settings.n_ubatch,
                "prompt_cache_enabled": settings.prompt_cache_enabled,
                "model_loaded": model_service.model_loaded,
                "load_time_seconds": model_service.load_time_seconds,
                "busy": model_service.busy,
                "web_search_enabled": search_service.enabled,
                "endpoints": {
                    "generate": "POST /generate",
                    "research_generate": "POST /generate/research",
                    "stream": "POST /generate/stream",
                    "search": "POST /search/web",
                    "health": "GET /health",
                    "info": "GET /info",
                },
            }
        )

    return app


app = create_app()


def main() -> None:
    s = Settings()
    print(f"\n{'=' * 50}")
    print(f"{APP_NAME.upper()} v{APP_VERSION}")
    print(f"Creator: {CREATOR}")
    print(f"{'=' * 50}")
    print(f"Model:  {s.model_path}")
    print(f"GPU:    {s.n_gpu_layers} layers")
    print(f"Ctx:    {s.n_ctx} | Threads: {s.n_threads}")
    print(f"URL:    http://localhost:{s.port}\n")

    uvicorn.run(app, host=s.host, port=s.port, log_level="warning")


if __name__ == "__main__":
    main()