#!/usr/bin/env python3
"""OpenAI-compatible RAG API proxy using llama.cpp for embeddings and Qdrant for retrieval.

Sits between clients and a llama.cpp LLM server, injecting relevant context
from a Qdrant vector database into chat completion requests.

Usage:
    rag_framework serve /rag/vector.db --embed-url URL --model-host HOST --model-port PORT
"""

import logging
import sys
import textwrap
import time
import uuid
from argparse import ArgumentParser, Namespace
from contextlib import asynccontextmanager
from functools import cache

import aiohttp
import openai
import qdrant_client
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

logger = logging.getLogger("rag_framework")

COLLECTION_NAME = "rag"


# ---------------------------------------------------------------------------
# OpenAI API Compatible Data Models
# ---------------------------------------------------------------------------
class ChatMessage(BaseModel):
    role: str = Field(description="The role of the message author")
    content: str = Field(description="The contents of the message")


class ChatCompletionRequest(BaseModel):
    messages: list[ChatMessage] = Field(description="A list of messages comprising the conversation")
    model: str = Field("", description="ID of the model to use")
    temperature: float | None = Field(1.0, ge=0, le=2, description="Sampling temperature")
    top_p: float | None = Field(1.0, ge=0, le=1, description="Nucleus sampling")
    n: int | None = Field(1, ge=1, le=128, description="Number of completions to generate")
    stream: bool | None = Field(False, description="Whether to stream back partial progress")
    stop: str | list[str] | None = Field(None, description="Sequences where the API will stop generating")
    max_completion_tokens: int | None = Field(None, ge=1, description="Maximum number of tokens to generate")
    presence_penalty: float | None = Field(0, ge=-2, le=2, description="Presence penalty")
    frequency_penalty: float | None = Field(0, ge=-2, le=2, description="Frequency penalty")
    user: str | None = Field(None, description="A unique identifier representing your end-user")


class Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


class Choice(BaseModel):
    index: int
    message: ChatMessage
    finish_reason: str | None = None


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


class Delta(BaseModel):
    role: str | None = None
    content: str | None = None
    reasoning_content: str | None = None


class StreamChoice(BaseModel):
    index: int
    delta: Delta
    finish_reason: str | None = None


class ChatCompletionStreamResponse(BaseModel):
    id: str
    object: str = "chat.completion.chunk"
    created: int
    model: str
    choices: list[StreamChoice]


# ---------------------------------------------------------------------------
# Embedding via llama.cpp
# ---------------------------------------------------------------------------
class LlamaCppEmbedder:
    """Generate embeddings via a llama.cpp server's /v1/embeddings endpoint."""

    def __init__(self, api_url):
        self.embeddings_url = f"{api_url.rstrip('/')}/v1/embeddings"

    async def embed_query(self, text):
        """Embed a single query and return its vector."""
        payload = {"input": [text]}
        async with aiohttp.ClientSession() as session:
            try:
                async with session.post(
                    self.embeddings_url, json=payload, timeout=aiohttp.ClientTimeout(total=30)
                ) as resp:
                    resp.raise_for_status()
                    result = await resp.json()
                    return result["data"][0]["embedding"]
            except aiohttp.ClientError as e:
                raise RuntimeError(f"Embedding server request failed: {e}") from e


# ---------------------------------------------------------------------------
# RAG Service
# ---------------------------------------------------------------------------
class RagService:
    def __init__(self, vector_path, embed_url, model_host, model_port):
        self.qclient = qdrant_client.AsyncQdrantClient(path=vector_path)
        self.embedder = LlamaCppEmbedder(api_url=embed_url)
        self.llm = openai.AsyncOpenAI(
            api_key="ramalama",
            base_url=f"http://{model_host}:{model_port}",
            http_client=openai.DefaultAioHttpClient(),
        )

    async def search(self, query, limit=5, neighbor_window=1):
        """Embed query, search Qdrant, and expand with neighboring chunks.

        For each matched chunk, retrieves `neighbor_window` chunks before and
        after it within the same document. Each match produces its own context
        window. Adjacent or overlapping windows from the same document are merged.

        Returns a list of context strings, one per merged window.
        """
        vector = await self.embedder.embed_query(query)
        results = await self.qclient.query_points(
            collection_name=COLLECTION_NAME,
            query=vector,
            limit=limit,
            with_payload=True,
        )

        if not results.points:
            return []

        # Build expanded set of (doc_id, chunk_index) pairs
        expanded = set()
        has_metadata = False
        for point in results.points:
            if not point.payload or "chunk_index" not in point.payload:
                continue
            has_metadata = True
            doc_id = point.payload.get("doc_id", 0)
            chunk_idx = point.payload["chunk_index"]
            for offset in range(-neighbor_window, neighbor_window + 1):
                expanded.add((doc_id, chunk_idx + offset))

        # Fall back to matched chunks directly if no metadata
        if not has_metadata:
            return [point.payload["document"] for point in results.points if point.payload]

        # Group by doc_id and merge overlapping windows
        from collections import defaultdict

        by_doc = defaultdict(set)
        for doc_id, chunk_idx in expanded:
            by_doc[doc_id].add(chunk_idx)

        # Fetch all needed chunks in one query
        from qdrant_client import models

        all_indices = [idx for indices in by_doc.values() for idx in indices]
        neighbor_results = await self.qclient.scroll(
            collection_name=COLLECTION_NAME,
            scroll_filter=models.Filter(
                must=[
                    models.FieldCondition(
                        key="chunk_index",
                        match=models.MatchAny(any=all_indices),
                    )
                ]
            ),
            with_payload=True,
            limit=len(expanded),
        )

        # Index fetched chunks by (doc_id, chunk_index)
        chunk_map = {}
        for point in neighbor_results[0]:
            if point.payload:
                key = (point.payload.get("doc_id", 0), point.payload["chunk_index"])
                chunk_map[key] = point.payload["document"]

        # Build context windows: group contiguous indices per doc, sorted
        context_windows = []
        for doc_id in sorted(by_doc):
            indices = sorted(by_doc[doc_id])
            # Split into contiguous runs
            runs = []
            current_run = [indices[0]]
            for idx in indices[1:]:
                if idx == current_run[-1] + 1:
                    current_run.append(idx)
                else:
                    runs.append(current_run)
                    current_run = [idx]
            runs.append(current_run)

            for run in runs:
                parts = [chunk_map[(doc_id, i)] for i in run if (doc_id, i) in chunk_map]
                if parts:
                    context_windows.append("\n\n".join(parts))

        return context_windows

    async def create_chat_completion(self, request):
        """OpenAI-compatible chat completion with RAG context injection."""
        completion_id = f"chatcmpl-{uuid.uuid4().hex[:29]}"
        created = int(time.time())

        enhanced_messages = await self._build_rag_messages(request.messages)

        if request.stream:
            return await self._stream_completion(completion_id, created, request, enhanced_messages)
        return await self._complete_chat(completion_id, created, request, enhanced_messages)

    async def _build_rag_messages(self, messages):
        """Build messages with RAG context injected as a system message."""
        if not messages:
            return []

        latest = messages[-1]
        if latest.role != "user":
            return [{"role": m.role, "content": m.content} for m in messages]

        rag_chunks = await self.search(latest.content)
        if not rag_chunks:
            return [{"role": m.role, "content": m.content} for m in messages]

        context = "\n\n".join(rag_chunks)

        system_msg = {
            "role": "system",
            "content": textwrap.dedent(f"""\
                Use the provided context to answer questions accurately.
                If the answer is not in the context, say so - do not fabricate details.

                ### Retrieved Context:
                {context}"""),
        }

        result = [system_msg]
        result.extend({"role": m.role, "content": m.content} for m in messages)
        return result

    async def _complete_chat(self, completion_id, created, request, enhanced_messages):
        response = await self.llm.chat.completions.create(
            model=request.model,
            messages=enhanced_messages,
            temperature=request.temperature,
            max_completion_tokens=request.max_completion_tokens,
            stream=False,
        )
        return ChatCompletionResponse(
            id=completion_id,
            created=created,
            model=request.model,
            choices=[
                Choice(
                    index=0,
                    message=ChatMessage(role="assistant", content=response.choices[0].message.content),
                    finish_reason=response.choices[0].finish_reason,
                )
            ],
            usage=Usage(
                prompt_tokens=response.usage.prompt_tokens,
                completion_tokens=response.usage.completion_tokens,
                total_tokens=response.usage.total_tokens,
            ),
        )

    async def _stream_completion(self, completion_id, created, request, enhanced_messages):
        async def generate():
            first_chunk = ChatCompletionStreamResponse(
                id=completion_id,
                created=created,
                model=request.model,
                choices=[StreamChoice(index=0, delta=Delta(role="assistant"), finish_reason=None)],
            )
            yield f"data: {first_chunk.model_dump_json()}\n\n"

            response = await self.llm.chat.completions.create(
                model=request.model,
                messages=enhanced_messages,
                temperature=request.temperature,
                max_completion_tokens=request.max_completion_tokens,
                stream=True,
            )

            async for chunk in response:
                if chunk.choices and (delta := chunk.choices[0].delta):
                    content = getattr(delta, "content", None)
                    reasoning_content = getattr(delta, "reasoning_content", None)
                    if content is not None or reasoning_content is not None:
                        stream_chunk = ChatCompletionStreamResponse(
                            id=completion_id,
                            created=created,
                            model=request.model,
                            choices=[
                                StreamChoice(
                                    index=0,
                                    delta=Delta(content=content, reasoning_content=reasoning_content),
                                    finish_reason=None,
                                )
                            ],
                        )
                        yield f"data: {stream_chunk.model_dump_json()}\n\n"

                if chunk.choices and chunk.choices[0].finish_reason:
                    final_chunk = ChatCompletionStreamResponse(
                        id=completion_id,
                        created=created,
                        model=request.model,
                        choices=[
                            StreamChoice(index=0, delta=Delta(), finish_reason=chunk.choices[0].finish_reason)
                        ],
                    )
                    yield f"data: {final_chunk.model_dump_json()}\n\n"

            yield "data: [DONE]\n\n"

        return StreamingResponse(
            generate(),
            media_type="text/plain",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
        )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
async def _request(host, port, path, timeout=10):
    async with aiohttp.ClientSession(
        f"http://{host}:{port}",
        timeout=aiohttp.ClientTimeout(total=timeout),
    ) as session:
        async with session.get(path) as resp:
            data = await resp.json()
            return resp, data


async def _wait_for_llm(host, port, total_timeout=120):
    import asyncio

    end_time = time.monotonic() + total_timeout
    while time.monotonic() < end_time:
        try:
            resp, _ = await _request(host, port, "/health", timeout=2)
            if resp.status == 200:
                return True
            if resp.status == 404:
                try:
                    models_resp, _ = await _request(host, port, "/models", timeout=2)
                    if models_resp.status == 200:
                        return True
                except aiohttp.ClientError:
                    pass
        except aiohttp.ClientError:
            pass
        await asyncio.sleep(1)
    raise TimeoutError(f"LLM server at {host}:{port} did not become ready after {total_timeout}s")


# ---------------------------------------------------------------------------
# FastAPI Application
# ---------------------------------------------------------------------------
@cache
def _get_service():
    args = _get_args()
    return RagService(args.vector_path, args.embed_url, args.model_host, args.model_port)


@asynccontextmanager
async def lifespan(app: FastAPI):
    args = _get_args()
    _get_service()
    await _wait_for_llm(args.model_host, args.model_port)
    yield


app = FastAPI(title="RAG-Enhanced OpenAI Compatible API", lifespan=lifespan)


@app.post("/v1/chat/completions", response_model=None)
async def create_chat_completion(request: ChatCompletionRequest):
    try:
        return await _get_service().create_chat_completion(request)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/models")
async def llama_models():
    args = _get_args()
    try:
        resp, data = await _request(args.model_host, args.model_port, "/models")
        if resp.status == 200:
            for model in data.get("models", []):
                model["name"] += "+rag"
            return data
        raise HTTPException(status_code=resp.status, detail=resp.reason)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=503, detail=f"LLM service unavailable: {e}")


@app.get("/v1/models")
async def list_models():
    try:
        service = _get_service()
        llm_models = await service.llm.models.list()
        return {
            "object": "list",
            "data": [
                {
                    "id": f"{m.id}+rag",
                    "object": "model",
                    "created": getattr(m, "created", int(time.time())),
                    "owned_by": getattr(m, "owned_by", "ramalama"),
                }
                for m in llm_models.data
            ],
        }
    except Exception as e:
        raise HTTPException(status_code=503, detail=f"LLM service unavailable: {e}")


@app.get("/health")
async def health_check():
    try:
        args = _get_args()
        resp, _ = await _request(args.model_host, args.model_port, "/health")
        if resp.status != 200:
            raise HTTPException(status_code=resp.status, detail=resp.reason)
        service = _get_service()
        if service.qclient is None or service.llm is None:
            raise HTTPException(status_code=503, detail="RAG has not been initialized")
        return {"status": "ok", "rag": "ok"}
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=503, detail=f"LLM service unavailable: {e}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
parser = ArgumentParser(description="OpenAI-compatible RAG API server")
parser.add_argument("--debug", action="store_true", help="Enable debug logging")
subparsers = parser.add_subparsers(dest="command")

serve_parser = subparsers.add_parser("serve", help="Run RAG as OpenAI-compatible HTTP API server")
serve_parser.add_argument("vector_path", type=str, help="Path to the Qdrant vector database")
serve_parser.add_argument("--embed-url", dest="embed_url", required=True, help="URL of the llama.cpp embedding server")
serve_parser.add_argument("--model-host", default="localhost", help="Hostname of the LLM server")
serve_parser.add_argument("--model-port", default=8080, type=int, help="Port of the LLM server")
serve_parser.add_argument("--host", default="0.0.0.0", help="Host to bind server")
serve_parser.add_argument("--port", default=8081, type=int, help="Port for RAG API")


@cache
def _get_args():
    return parser.parse_args()


if __name__ == "__main__":
    try:
        args = _get_args()
        logging.basicConfig(
            level=logging.DEBUG if args.debug else logging.WARNING,
            style="{",
            format="{asctime} {name} {levelname}: {message}",
            force=True,
        )
        if args.command == "serve":
            uvicorn.run(app, host=args.host, port=args.port)
        else:
            parser.print_help()
    except ValueError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
