server.py

import os
from datetime import datetime, timezone
from typing import Union

import docker
from dotenv import load_dotenv
from fastapi import Body, Depends, FastAPI, Header, HTTPException
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker

from compat.migrate import get_client, list_runs, migrate_all, migrate_run_v1
from python.env import get_database_url, get_smtp_config
from python.docker import start_server, stop_server, stop_all
from python.models import Run, RunStatus, RunTriggers, RunTriggerType
from python.server import check_run, send_alert, check_api_key

load_dotenv()

SMTP_CONFIG = get_smtp_config()
DATABASE_URL = get_database_url()
DOMAIN = os.getenv("W_DOMAIN", "localhost")
if not DATABASE_URL:
    raise ValueError("DATABASE_URL environment variable is not set")

client = docker.from_env()
engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

app = FastAPI()


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def require_api_key(
    authorization: str = Header(None),
    session: Session = Depends(get_db),
):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Authorization header missing or invalid")
    raw_key = authorization.removeprefix("Bearer ")
    api_key = check_api_key(session, raw_key)
    if not api_key:
        raise HTTPException(status_code=401, detail="Invalid API key")
    return api_key


@app.post("/api/docker/start")
async def start_docker(
    api_key=Depends(require_api_key),
):
    host = os.getenv('D_DOMAIN', 'localhost')
    port, password, url, private_key, ssh_port = start_server(client, host=host)

    cmd_save = f"echo -e '{private_key}' > id_ed25519; chmod 600 id_ed25519"
    cmd_connect = f"{cmd_save}; ssh -i id_ed25519 -p {ssh_port} mlop@{host}"
    cmd_ssh = f"{cmd_save}; echo -e '\nHost {password}\n  HostName {host}\n  Port {ssh_port}\n  User mlop\n  IdentityFile' $(realpath id_ed25519) >> ~/.ssh/config"
    cmd_code = f"{cmd_ssh}; code --remote ssh-remote+{password} /home/mlop"
    return {"port": port, "password": password, "url": url, "connect": cmd_connect, "code": cmd_code}  # "key": private_key


@app.post("/api/docker/stop")
async def stop_docker(
    port: int = Body(..., embed=True),
    api_key=Depends(require_api_key),
):
    try:
        stop_server(client, int(port))
        return {"status": "success"}
    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Failed to stop server: {e}")


@app.post("/api/docker/stop-all")
async def stop_all_docker(
    api_key=Depends(require_api_key),
):
    try:
        stop_all(client)
        return {"status": "success"}
    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Failed to stop all servers: {e}")


@app.post("/api/runs/trigger")
async def get_run_triggers(
    runId: int = Body(..., embed=True),
    session: Session = Depends(get_db),
    authorization: str = Header(None),
):
    run = check_run(session, runId, authorization)

    if not run.status == RunStatus.CANCELLED:
        triggers = session.query(RunTriggers).filter(
            RunTriggers.runId == runId).all()
        for trigger in triggers:
            if trigger.triggerType == RunTriggerType.CANCEL:
                run.status = RunStatus.CANCELLED
                run.statusUpdated = datetime.now(timezone.utc)
                session.commit()
                session.refresh(run)
    else:
        triggers = []

    return {
        "status": run.status,
        "triggers": [
            {
                "trigger": trigger.trigger,
            }
            for trigger in triggers
        ]
        if (False and triggers is not None)
        else None,
    }


@app.post("/api/runs/alert")
async def set_run_alerts(
    runId: int = Body(..., embed=True),
    alert: dict[str, Union[str, int, bool, None]] = Body(..., embed=True),
    session: Session = Depends(get_db),
    authorization: str = Header(None),
):
    run = check_run(session, runId, authorization)

    if not isinstance(alert, dict):  # TODO: add more checks
        raise HTTPException(status_code=400, detail="Invalid alert")

    try:
        send_alert(
            session,
            run,
            SMTP_CONFIG,
            last_update_time=datetime.fromtimestamp(
                alert.get("timestamp") / 1000, tz=timezone.utc
            )
            if alert.get("timestamp")
            else datetime.now(timezone.utc),
            title=alert.get("title", "Status Update"),
            body=alert.get("body", "alert"),
            level=alert.get("level", "INFO"),
            email=alert.get("email", True),
        )

        if alert.get("url"):
            # TODO: add webhook support
            raise HTTPException(status_code=302, detail=alert.get("url"))
        else:
            return {"status": "success"}
    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Failed to send alert: {e}")


@app.post("/api/compat/w/viewer")  # TODO: protect
async def _viewer(key: str = Body(..., embed=True)):
    c = get_client(key, DOMAIN)
    return c.viewer()


@app.post("/api/compat/w/list-runs")
async def _list_runs(
    auth: str = Body(..., embed=True),
    key: str = Body(..., embed=True),
    entity: str = Body(..., embed=True),
):
    c = get_client(key, DOMAIN)
    return list_runs(c, entity)


@app.post("/api/compat/w/migrate-all")
async def _migrate_all(
    auth: str = Body(..., embed=True),
    key: str = Body(..., embed=True),
    entity: str = Body(..., embed=True),
):
    if migrate_all(auth, key, entity, DOMAIN):
        return {"status": "success"}
    else:
        raise HTTPException(status_code=500, detail="Failed to migrate runs")


@app.post("/api/compat/w/migrate-run")
async def _migrate_run(
    auth: str = Body(..., embed=True),
    key: str = Body(..., embed=True),
    entity: str = Body(..., embed=True),
    project: str = Body(..., embed=True),
    run: str = Body(..., embed=True),
):
    c = get_client(key, DOMAIN)
    if migrate_run_v1(auth, c, entity, project, run):
        return {"status": "success"}
    else:
        raise HTTPException(status_code=500, detail="Failed to migrate run")


if __name__ == "__main__":
    import uvicorn

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