scripts/run_server.py

#!/usr/bin/env python3
"""
Server startup script with pre-flight checks.

Usage:
    python scripts/run_server.py
    python scripts/run_server.py --port 8080 --reload
    python scripts/run_server.py --host 0.0.0.0 --port 8000
"""

import argparse
import sys
import time
from pathlib import Path

import requests
from rich.console import Console
from rich.panel import Panel
from rich.table import Table

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from config.settings import LLMProvider, Settings


console = Console()


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Start EyeWiki RAG API server with pre-flight checks"
    )

    parser.add_argument(
        "--host",
        type=str,
        default="0.0.0.0",
        help="Host to bind (default: 0.0.0.0)",
    )

    parser.add_argument(
        "--port",
        type=int,
        default=8000,
        help="Port number (default: 8000)",
    )

    parser.add_argument(
        "--reload",
        action="store_true",
        help="Enable hot reload for development",
    )

    parser.add_argument(
        "--skip-checks",
        action="store_true",
        help="Skip pre-flight checks (not recommended)",
    )

    return parser.parse_args()


def print_header():
    """Print welcome header."""
    console.print()
    console.print(
        Panel.fit(
            "[bold blue]EyeWiki RAG API Server[/bold blue]\n"
            "[dim]Retrieval-Augmented Generation for Medical Knowledge[/dim]",
            border_style="blue",
        )
    )
    console.print()


def check_ollama(settings: Settings) -> bool:
    """
    Check if Ollama is running and has required models.

    Args:
        settings: Application settings

    Returns:
        True if check passed, False otherwise
    """
    console.print("[bold cyan]1. Checking Ollama service...[/bold cyan]")

    try:
        # Check if Ollama is running
        response = requests.get(f"{settings.ollama_base_url}/api/tags", timeout=5)
        response.raise_for_status()

        models_data = response.json()
        available_models = [model["name"] for model in models_data.get("models", [])]

        # Check for required LLM model (embedding model is sentence-transformers, not Ollama)
        required_models = {
            "LLM": settings.llm_model,
        }

        table = Table(show_header=True, header_style="bold magenta")
        table.add_column("Model Type", style="cyan")
        table.add_column("Required Model", style="yellow")
        table.add_column("Status", style="green")

        all_found = True
        for model_type, model_name in required_models.items():
            # Check if model name (with or without tag) is in available models
            found = any(
                model_name in model or model in model_name for model in available_models
            )

            status = "[green]✓ Found[/green]" if found else "[red]✗ Missing[/red]"
            table.add_row(model_type, model_name, status)

            if not found:
                all_found = False

        console.print(table)

        if not all_found:
            console.print(
                "\n[red]Error:[/red] Some required models are missing. "
                "Pull them with:"
            )
            for model_type, model_name in required_models.items():
                if not any(
                    model_name in model or model in model_name
                    for model in available_models
                ):
                    console.print(f"  [yellow]ollama pull {model_name}[/yellow]")
            console.print()
            return False

        console.print("[green]✓ Ollama is running with all required models[/green]\n")
        return True

    except requests.RequestException as e:
        console.print(f"[red]✗ Failed to connect to Ollama:[/red] {e}")
        console.print(
            f"\nMake sure Ollama is running at [yellow]{settings.ollama_base_url}[/yellow]"
        )
        console.print("Start it with: [yellow]ollama serve[/yellow]\n")
        return False


def check_openai_config(settings: Settings) -> bool:
    """
    Check if OpenAI-compatible API is configured with required API key.

    Args:
        settings: Application settings

    Returns:
        True if check passed, False otherwise
    """
    console.print("[bold cyan]1. Checking OpenAI-compatible API configuration...[/bold cyan]")

    table = Table(show_header=True, header_style="bold magenta")
    table.add_column("Property", style="cyan")
    table.add_column("Value", style="yellow")
    table.add_column("Status", style="green")

    # Check API key
    has_key = bool(settings.openai_api_key)
    key_display = f"{settings.openai_api_key[:8]}..." if has_key else "(not set)"
    key_status = "[green]✓ Set[/green]" if has_key else "[red]✗ Missing[/red]"
    table.add_row("API Key", key_display, key_status)

    # Show base URL
    base_url = settings.openai_base_url or "(OpenAI default)"
    table.add_row("Base URL", base_url, "[green]✓[/green]")

    # Show model
    table.add_row("Model", settings.openai_model, "[green]✓[/green]")

    console.print(table)

    if not has_key:
        console.print(
            "\n[red]Error:[/red] API key is required for OpenAI-compatible provider."
        )
        console.print(
            "Set the [yellow]OPENAI_API_KEY[/yellow] environment variable or add it to your [yellow].env[/yellow] file.\n"
        )
        return False

    console.print("[green]✓ OpenAI-compatible API configuration looks good[/green]\n")
    return True


def check_vector_store(settings: Settings) -> bool:
    """
    Check if vector store exists and has documents.

    Args:
        settings: Application settings

    Returns:
        True if check passed, False otherwise
    """
    console.print("[bold cyan]2. Checking vector store...[/bold cyan]")

    qdrant_path = Path(settings.qdrant_path)
    collection_name = settings.qdrant_collection_name

    # Check if Qdrant directory exists
    if not qdrant_path.exists():
        console.print(f"[red]✗ Qdrant directory not found:[/red] {qdrant_path}")
        console.print(
            "\nRun the indexing pipeline first:\n"
            "  [yellow]python scripts/build_index.py --index-vectors[/yellow]\n"
        )
        return False

    # Try to connect to Qdrant and check collection
    try:
        from qdrant_client import QdrantClient

        client = QdrantClient(path=str(qdrant_path))

        # Check if collection exists
        collections = client.get_collections().collections
        collection_names = [col.name for col in collections]

        if collection_name not in collection_names:
            console.print(
                f"[red]✗ Collection '{collection_name}' not found[/red]\n"
                f"Available collections: {collection_names}"
            )
            console.print(
                "\nRun the indexing pipeline first:\n"
                "  [yellow]python scripts/build_index.py --index-vectors[/yellow]\n"
            )
            return False

        # Get collection info
        collection_info = client.get_collection(collection_name)
        points_count = collection_info.points_count

        if points_count == 0:
            console.print(
                f"[yellow]⚠ Collection '{collection_name}' exists but is empty[/yellow]"
            )
            console.print(
                "\nRun the indexing pipeline:\n"
                "  [yellow]python scripts/build_index.py --index-vectors[/yellow]\n"
            )
            return False

        # Print stats
        table = Table(show_header=True, header_style="bold magenta")
        table.add_column("Property", style="cyan")
        table.add_column("Value", style="yellow")

        table.add_row("Collection", collection_name)
        table.add_row("Location", str(qdrant_path))
        table.add_row("Documents", f"{points_count:,}")

        console.print(table)
        console.print("[green]✓ Vector store is ready[/green]\n")
        return True

    except Exception as e:
        console.print(f"[red]✗ Failed to access vector store:[/red] {e}\n")
        return False


def check_required_files() -> bool:
    """
    Check if all required files exist.

    Returns:
        True if all files exist, False otherwise
    """
    console.print("[bold cyan]3. Checking required files...[/bold cyan]")

    required_files = {
        "System Prompt": project_root / "prompts" / "system_prompt.txt",
        "Query Prompt": project_root / "prompts" / "query_prompt.txt",
        "Medical Disclaimer": project_root / "prompts" / "medical_disclaimer.txt",
    }

    table = Table(show_header=True, header_style="bold magenta")
    table.add_column("File", style="cyan")
    table.add_column("Path", style="yellow")
    table.add_column("Status", style="green")

    all_exist = True
    for name, path in required_files.items():
        exists = path.exists()
        status = "[green]✓ Found[/green]" if exists else "[red]✗ Missing[/red]"
        table.add_row(name, str(path.relative_to(project_root)), status)

        if not exists:
            all_exist = False

    console.print(table)

    if not all_exist:
        console.print(
            "\n[red]Error:[/red] Some required files are missing.\n"
            "Make sure all prompt files are in the [yellow]prompts/[/yellow] directory.\n"
        )
        return False

    console.print("[green]✓ All required files found[/green]\n")
    return True


def run_preflight_checks(skip_checks: bool = False) -> bool:
    """
    Run all pre-flight checks.

    Args:
        skip_checks: Skip all checks if True

    Returns:
        True if all checks passed, False otherwise
    """
    if skip_checks:
        console.print("[yellow]⚠ Skipping pre-flight checks[/yellow]\n")
        return True

    console.print("[bold yellow]Running Pre-flight Checks...[/bold yellow]\n")

    # Load settings
    try:
        settings = Settings()
    except Exception as e:
        console.print(f"[red]✗ Failed to load settings:[/red] {e}\n")
        return False

    console.print(f"[dim]LLM Provider: {settings.llm_provider.value}[/dim]\n")

    # Check LLM provider (Ollama or OpenAI-compatible)
    if settings.llm_provider == LLMProvider.OLLAMA:
        llm_check = check_ollama(settings)
    else:
        llm_check = check_openai_config(settings)

    # Run checks
    checks = [
        llm_check,
        check_vector_store(settings),
        check_required_files(),
    ]

    if not all(checks):
        console.print("[bold red]✗ Pre-flight checks failed[/bold red]")
        console.print("Fix the issues above and try again.\n")
        return False

    console.print("[bold green]✓ All pre-flight checks passed![/bold green]\n")
    return True


def print_access_urls(host: str, port: int):
    """
    Print access URLs for the server.

    Args:
        host: Server host
        port: Server port
    """
    # Determine display host
    display_host = "localhost" if host in ["0.0.0.0", "127.0.0.1"] else host

    table = Table(
        show_header=True,
        header_style="bold magenta",
        title="[bold green]Server Access URLs[/bold green]",
        title_style="bold green",
    )
    table.add_column("Service", style="cyan", width=20)
    table.add_column("URL", style="yellow")
    table.add_column("Description", style="dim")

    urls = [
        ("API Root", f"http://{display_host}:{port}", "API information"),
        ("Health Check", f"http://{display_host}:{port}/health", "Service health status"),
        (
            "Interactive Docs",
            f"http://{display_host}:{port}/docs",
            "Swagger UI documentation",
        ),
        ("ReDoc", f"http://{display_host}:{port}/redoc", "Alternative API docs"),
        (
            "Gradio UI",
            f"http://{display_host}:{port}/ui",
            "Web chat interface",
        ),
    ]

    for service, url, description in urls:
        table.add_row(service, url, description)

    console.print()
    console.print(table)
    console.print()

    # Print quick start commands
    console.print("[bold cyan]Quick Test Commands:[/bold cyan]")
    console.print(
        f"  [dim]# Test health endpoint[/dim]\n"
        f"  [yellow]curl http://{display_host}:{port}/health[/yellow]\n"
    )
    console.print(
        f"  [dim]# Query the API[/dim]\n"
        f"  [yellow]curl -X POST http://{display_host}:{port}/query \\[/yellow]\n"
        f'  [yellow]  -H "Content-Type: application/json" \\[/yellow]\n'
        f'  [yellow]  -d \'{{"question": "What is glaucoma?"}}\' [/yellow]\n'
    )


def start_server(host: str, port: int, reload: bool):
    """
    Start the uvicorn server.

    Args:
        host: Server host
        port: Server port
        reload: Enable hot reload
    """
    console.print("[bold green]Starting server...[/bold green]\n")

    # Print URLs before starting
    print_access_urls(host, port)

    # Import uvicorn here to avoid import errors if not installed
    try:
        import uvicorn
    except ImportError:
        console.print("[red]Error:[/red] uvicorn is not installed")
        console.print("Install it with: [yellow]pip install uvicorn[/yellow]\n")
        sys.exit(1)

    # Start server
    try:
        console.print(
            f"[dim]Server listening on {host}:{port}[/dim]",
            f"[dim](Press CTRL+C to stop)[/dim]\n",
        )

        uvicorn.run(
            "src.api.main:app",
            host=host,
            port=port,
            reload=reload,
            log_level="info",
        )

    except KeyboardInterrupt:
        console.print("\n\n[yellow]Server stopped by user[/yellow]")
    except Exception as e:
        console.print(f"\n[red]Error starting server:[/red] {e}")
        sys.exit(1)


def main():
    """Main entry point."""
    args = parse_args()

    print_header()

    # Run pre-flight checks
    if not run_preflight_checks(skip_checks=args.skip_checks):
        console.print("[red]Startup aborted due to failed checks[/red]\n")
        sys.exit(1)

    # Start server
    start_server(host=args.host, port=args.port, reload=args.reload)


if __name__ == "__main__":
    main()