headquarter/apps/api/src/main.py

import logging
import os

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from sqlalchemy import select, text

from src.api.auth import router as auth_router
from src.api.dashboard import router as dashboard_router
from src.api.git_repositories import router as git_repositories_router
from src.api.projects import router as projects_router
from src.api.ssh_keys import router as ssh_keys_router
from src.api.tool_types import router as tool_types_router
from src.api.user_config import router as user_config_router
from src.api.users import router as users_router
from src.config import Settings
from src.database import SessionLocal, init_database
from src.logging_config import (
    ExceptionLoggingMiddleware,
    RequestLoggingMiddleware,
    configure_logging,
)
from src.models.tool_type import ToolType

# Configure logging early
log_level = os.getenv("LOG_LEVEL", "INFO").upper()
configure_logging(level=getattr(logging, log_level, logging.INFO))

logger = logging.getLogger(__name__)
settings = Settings()
app = FastAPI(title="Headquarter API")

# Configure CORS - must be before other middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=[settings.web_base_url],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.add_middleware(RequestLoggingMiddleware)
app.add_middleware(ExceptionLoggingMiddleware)


async def _table_exists(session, table_name: str) -> bool:
    """Check if a table exists in the database."""
    try:
        result = await session.execute(
            text("""
                SELECT EXISTS (
                    SELECT FROM information_schema.tables
                    WHERE table_schema = 'public'
                    AND table_name = :table_name
                )
            """),
            {"table_name": table_name},
        )
        return result.scalar() or False
    except Exception:
        return False


async def seed_builtin_tool_types():
    async with SessionLocal() as session:
        # Check if tool_types table exists before attempting to seed
        if not await _table_exists(session, "tool_types"):
            logger.warning(
                "tool_types table does not exist. Skipping seeding. "
                "Migrations may not have run yet."
            )
            return

        builtin_types = [
            {
                "name": "code-server",
                "display_name": "VS Code Server",
                "description": "VS Code running in the browser via code-server",
                "compose_template": """version: "3.8"
services:
  code-server:
    image: lscr.io/linuxserver/code-server:latest
    container_name: {{TOOL_NAME}}
    environment:
      - PUID=1000
      - PGID=1000
      - TZ=Europe/London
    volumes:
      - {{REPO_PATH}}:/config/workspace
    ports:
      - "8443:8443"
    restart: unless-stopped""",
                "required_variables": ["REPO_PATH", "TOOL_NAME"],
            },
            {
                "name": "jupyter-notebook",
                "display_name": "Jupyter Notebook",
                "description": "Jupyter Lab for interactive development",
                "compose_template": """version: "3.8"
services:
  jupyter:
    image: jupyter/scipy-notebook:latest
    container_name: {{TOOL_NAME}}
    environment:
      - JUPYTER_ENABLE_LAB=yes
    volumes:
      - {{REPO_PATH}}:/home/jovyan/work
    ports:
      - "8888:8888"
    restart: unless-stopped""",
                "required_variables": ["REPO_PATH", "TOOL_NAME"],
            },
        ]

        for tool_data in builtin_types:
            existing = await session.scalar(select(ToolType).where(ToolType.name == tool_data["name"]))
            if not existing:
                tool_type = ToolType(
                    name=tool_data["name"],
                    display_name=tool_data["display_name"],
                    description=tool_data["description"],
                    compose_template=tool_data["compose_template"],
                    required_variables=tool_data["required_variables"],
                    is_builtin=True,
                )
                session.add(tool_type)

        await session.commit()
        logger.info("Built-in tool types seeded successfully.")


@app.on_event("startup")
async def on_startup():
    logger.info("Starting up Headquarter API...")

    # Initialize database (run migrations)
    db_ready = await init_database()
    if not db_ready:
        logger.error("Database initialization failed. Shutting down.")
        import sys
        sys.exit(1)

    # Seed built-in data
    await seed_builtin_tool_types()
    logger.info("Startup complete.")

@app.get("/health")
async def health_check():
    try:
        from sqlalchemy import text
        async with SessionLocal() as session:
            await session.execute(text("SELECT 1"))
        return {"status": "healthy", "database": "connected"}
    except Exception as exc:
        logger.error("Health check failed: %s", exc)
        return {"status": "unhealthy", "database": "disconnected", "error": str(exc)}

app.include_router(auth_router)
app.include_router(dashboard_router)
app.include_router(projects_router)
app.include_router(users_router)
app.include_router(ssh_keys_router)
app.include_router(git_repositories_router)
app.include_router(user_config_router)
app.include_router(tool_types_router)
app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")