#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   server.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   OpenAI-compatible FastAPI server for local inference.
"""

import argparse
import time
import uuid
from contextlib import asynccontextmanager
from typing import Any

import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice
from openai.types.completion_usage import CompletionUsage

from vi.deployment.local.backends.base import InferenceBackend
from vi.deployment.local.backends.vi_backend import ViBackend
from vi.deployment.local.schemas import (
    ChatCompletionRequest,
    ChatCompletionStreamChoice,
    ChatCompletionStreamChunk,
    ErrorResponse,
    ModelListResponse,
    ModelObject,
)
from vi.deployment.local.utils import (
    build_generation_config_from_request,
    convert_response_format_to_pydantic,
)


class BackendRegistry:
    """Simple registry for the inference backend instance."""

    def __init__(self) -> None:
        """Initialize empty backend registry."""
        self._backend: InferenceBackend | None = None

    def set_backend(self, backend: InferenceBackend) -> None:
        """Set the backend instance.

        Args:
            backend: The inference backend to use

        """
        self._backend = backend

    def get_backend(self) -> InferenceBackend:
        """Get the backend instance.

        Returns:
            The current inference backend

        Raises:
            RuntimeError: If backend is not initialized

        """
        if self._backend is None:
            raise RuntimeError(
                "Backend not initialized. Use create_server() to initialize."
            )
        return self._backend


# Module-level registry instance
_registry = BackendRegistry()


@asynccontextmanager
async def lifespan(_app: FastAPI):
    """FastAPI lifespan manager for startup/shutdown."""
    # Startup
    try:
        _registry.get_backend()
    except RuntimeError as e:
        raise RuntimeError(
            "Backend not initialized. Use create_server() to create the app."
        ) from e
    yield
    # Shutdown (if needed)


app = FastAPI(
    title="Vi Inference Server",
    description="OpenAI-compatible inference server for Datature Vi models",
    version="0.1.0",
    lifespan=lifespan,
)


@app.get("/health")
async def health_check() -> dict[str, str]:
    """Health check endpoint.

    Returns:
        Status dictionary

    """
    try:
        backend = _registry.get_backend()
        # Verify backend has models available
        models = backend.available_models()
        if not models:
            return {"status": "unhealthy", "error": "No models available"}
        return {"status": "healthy", "models": models}
    except Exception as e:
        return {"status": "unhealthy", "error": str(e)}


@app.get("/v1/health")
async def health_check_v1() -> dict[str, str]:
    """Health check endpoint (v1).

    Returns:
        Status dictionary

    """
    return await health_check()


@app.get("/v1/models")
async def list_models() -> ModelListResponse:
    """List available models - matches OpenAI spec.

    Returns:
        ModelListResponse with available models

    """
    backend = _registry.get_backend()
    models = backend.available_models()

    return ModelListResponse(
        object="list",
        data=[
            ModelObject(id=model_id, object="model", owned_by="datature")
            for model_id in models
        ],
    )


@app.post("/v1/chat/completions", response_model=None)
async def chat_completions(
    request: ChatCompletionRequest,
) -> ChatCompletion | StreamingResponse:
    """Chat completion endpoint with streaming support.

    Supports both OpenAI format (messages) and Vi SDK format (source + user_prompt).

    Args:
        request: Chat completion request

    Returns:
        ChatCompletion for non-streaming requests
        StreamingResponse for streaming requests

    """
    try:
        backend = _registry.get_backend()

        # Validate that the requested model is available
        available_models = backend.available_models()
        if request.model not in available_models:
            raise HTTPException(
                status_code=404,
                detail=f"Model '{request.model}' not found. Available models: {', '.join(available_models)}",
            )

        # Determine request format and prepare parameters
        if request.source is not None or request.user_prompt is not None:
            # Vi SDK format: use source and user_prompt directly
            use_vi_format = True
            source = request.source
            user_prompt = request.user_prompt
        elif request.messages:
            # OpenAI format: extract from messages
            use_vi_format = False
        else:
            raise HTTPException(
                status_code=400,
                detail="Either 'messages' (OpenAI format) or 'source'/'user_prompt' (Vi SDK format) must be provided",
            )

        # Convert response_format to Pydantic model once for efficiency
        pydantic_schema = None
        if request.response_format:
            pydantic_schema = convert_response_format_to_pydantic(
                request.response_format
            )

        # Build generation config in the format expected by backends
        generation_config = build_generation_config_from_request(request)

        if request.stream:
            return StreamingResponse(
                stream_completion(
                    request, backend, use_vi_format, generation_config, pydantic_schema
                ),
                media_type="text/event-stream",
            )

        # Non-streaming response
        if use_vi_format:
            result = await backend.generate_vi_format(
                source=source,
                user_prompt=user_prompt,
                response_format=pydantic_schema,
                generation_config=generation_config,
            )
        else:
            result = await backend.generate(
                messages=[msg.model_dump() for msg in request.messages],
                response_format=pydantic_schema,
                generation_config=generation_config,
            )

        return ChatCompletion(
            id=f"chatcmpl-{uuid.uuid4().hex[:8]}",
            created=int(time.time()),
            model=request.model,
            object="chat.completion",
            choices=[
                Choice(
                    index=0,
                    message=ChatCompletionMessage(
                        role="assistant", content=result.text
                    ),
                    finish_reason=result.finish_reason,
                )
            ],
            usage=CompletionUsage(
                prompt_tokens=result.prompt_tokens,
                completion_tokens=result.completion_tokens,
                total_tokens=result.prompt_tokens + result.completion_tokens,
            ),
        )

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


async def stream_completion(
    request: ChatCompletionRequest,
    backend: InferenceBackend,
    use_vi_format: bool = False,
    generation_config: dict | None = None,
    pydantic_schema: Any | None = None,
) -> Any:
    """SSE streaming matching OpenAI's format.

    Args:
        request: Chat completion request
        backend: Inference backend
        use_vi_format: Whether to use Vi SDK format (source/user_prompt)
        generation_config: Dictionary of generation parameters (temperature, max_tokens, etc.)
        pydantic_schema: Parsed Pydantic model for structured output

    Yields:
        SSE formatted chunks

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

    try:
        if use_vi_format:
            stream_generator = backend.generate_stream_vi_format(
                source=request.source,
                user_prompt=request.user_prompt,
                response_format=pydantic_schema,
                generation_config=generation_config,
            )
        else:
            stream_generator = backend.generate_stream(
                messages=[msg.model_dump() for msg in request.messages],
                response_format=pydantic_schema,
                generation_config=generation_config,
            )

        async for token in stream_generator:
            # Ensure token is a string (safety check)
            if not isinstance(token, str):
                # Skip non-string tokens or convert to string
                if token is None:
                    continue
                token = str(token)

            chunk = ChatCompletionStreamChunk(
                id=completion_id,
                object="chat.completion.chunk",
                created=created,
                model=request.model,
                choices=[
                    ChatCompletionStreamChoice(
                        index=0,
                        delta={"content": token},
                        finish_reason=None,
                    )
                ],
            )
            yield f"data: {chunk.model_dump_json()}\n\n"

        # Final chunk with finish_reason
        final_chunk = ChatCompletionStreamChunk(
            id=completion_id,
            object="chat.completion.chunk",
            created=created,
            model=request.model,
            choices=[
                ChatCompletionStreamChoice(
                    index=0,
                    delta={},
                    finish_reason="stop",
                )
            ],
        )
        yield f"data: {final_chunk.model_dump_json()}\n\n"
        yield "data: [DONE]\n\n"

    except Exception as e:
        # Send error as SSE
        error_data = ErrorResponse(error={"message": str(e), "type": type(e).__name__})
        yield f"data: {error_data.model_dump_json()}\n\n"


def create_server(
    backend: InferenceBackend | None = None,
    pretrained_model_name_or_path: str | None = None,
    run_id: str | None = None,
    secret_key: str | None = None,
    organization_id: str | None = None,
    model_name: str | None = None,
    **kwargs: Any,
) -> FastAPI:
    """Create and configure the FastAPI server.

    Args:
        backend: Pre-configured backend instance (takes precedence)
        pretrained_model_name_or_path: Path to local model or HuggingFace model ID
            (for auto-creating backend)
        run_id: Datature run ID (for auto-creating backend)
        secret_key: Datature API secret key
        organization_id: Datature organization ID
        model_name: Custom model identifier to use in API requests.
            If not provided, will be auto-generated from the model path/run_id
        **kwargs: Additional backend initialization arguments

    Returns:
        Configured FastAPI application

    Example:
        ```python
        # Using Vi backend with local model
        app = create_server(pretrained_model_name_or_path="./my-model")

        # Using Vi backend with HuggingFace model
        app = create_server(pretrained_model_name_or_path="nvidia/Cosmos-Reason2-2B")

        # Using Vi backend with custom model name
        app = create_server(
            pretrained_model_name_or_path="nvidia/Cosmos-Reason2-2B",
            model_name="my-vision-model",
        )

        # Using Vi backend with Datature model
        app = create_server(
            run_id="abc123",
            secret_key="your-key",
            organization_id="your-org",
        )

        # Using custom backend
        from vi.deployment.local.backends import ViBackend

        backend = ViBackend(pretrained_model_name_or_path="./my-model")
        app = create_server(backend=backend)

        # Run with uvicorn
        import uvicorn

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

    """
    if backend is None:
        # Auto-create Vi backend
        backend = ViBackend(
            pretrained_model_name_or_path=pretrained_model_name_or_path,
            run_id=run_id,
            secret_key=secret_key,
            organization_id=organization_id,
            model_name=model_name,
            **kwargs,
        )

    _registry.set_backend(backend)
    return app


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run Vi Inference Server")
    parser.add_argument(
        "--pretrained-model-name-or-path",
        type=str,
        help="Path to local model or HuggingFace model ID (e.g., nvidia/Cosmos-Reason2-2B)",
    )
    parser.add_argument("--run-id", type=str, help="Datature run ID")
    parser.add_argument("--secret-key", type=str, help="Datature API secret key")
    parser.add_argument("--organization-id", type=str, help="Datature organization ID")
    parser.add_argument(
        "--model-name",
        type=str,
        help="Custom model identifier to use in API requests (defaults to auto-generated from model path/run_id)",
    )
    parser.add_argument("--host", type=str, default="0.0.0.0", help="Server host")
    parser.add_argument("--port", type=int, default=8000, help="Server port")

    args = parser.parse_args()

    # Create server
    server = create_server(
        pretrained_model_name_or_path=args.pretrained_model_name_or_path,
        run_id=args.run_id,
        secret_key=args.secret_key,
        organization_id=args.organization_id,
        model_name=args.model_name,
    )

    # Run server
    uvicorn.run(server, host=args.host, port=args.port)
