advent-of-code-ui/backend/src/holt59/aoc/app.py

import asyncio
import datetime
import json
import logging
import subprocess
from pathlib import Path
from typing import Any, Literal, TypeAlias

from fastapi import FastAPI, Form, HTTPException
from fastapi.responses import FileResponse
from sse_starlette import EventSourceResponse

DataMode: TypeAlias = Literal["holt59", "tests"]

LOGGER = logging.getLogger("uvicorn.info")

GIT_AOC_PATH = Path("/code")
AOC_FILES_PATH = Path("/files")


def update_repository():
    subprocess.run(["git", "fetch"], cwd=GIT_AOC_PATH)
    subprocess.run(["git", "reset", "--hard", "origin/master"], cwd=GIT_AOC_PATH)


def read_head_commit():
    # read git commit
    process = subprocess.run(
        ["git", "log", "-1", '--format="%H"'], cwd=GIT_AOC_PATH, stdout=subprocess.PIPE
    )
    return process.stdout.decode("utf-8").strip()[1:-1]


def read_test_file(year: int, day: int, mode: DataMode) -> str:
    path = GIT_AOC_PATH.joinpath(f"src/holt59/aoc/inputs/{mode}/{year}/day{day}.txt")
    if not path.exists():
        return ""

    with open(path, "r") as fp:
        return fp.read()


app = FastAPI()


@app.post("/update")
def update():
    update_repository()
    head_commit = read_head_commit()
    LOGGER.info(f"repository now at {head_commit}")
    return {"HEAD": head_commit}


@app.get("/data")
async def data(year: int, day: int, mode: DataMode = "tests"):
    return {"content": read_test_file(year, day, mode)}


@app.get("/file")
async def file(name: str):
    path = AOC_FILES_PATH.joinpath(name).resolve()

    if path.parent != AOC_FILES_PATH:
        raise HTTPException(403)

    return FileResponse(path)


@app.post("/submit-sse")
async def submit_sse(year: int = Form(), day: int = Form(), input: str = Form()):
    data = input.rstrip().replace("\r\n", "\n")

    process = await asyncio.create_subprocess_exec(
        "holt59-aoc",
        "--stdin",
        "--api",
        "--output",
        AOC_FILES_PATH,
        "--year",
        str(year),
        str(day),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        stdin=asyncio.subprocess.PIPE,
    )

    assert process.stdin is not None
    assert process.stdout is not None
    assert process.stderr is not None

    message_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

    async def watch_stream(
        fp: asyncio.StreamReader, stream: Literal["stdout", "stderr"]
    ):
        while True:
            output = await fp.readline()
            if output == b"":
                break
            await message_queue.put(
                {
                    "type": "stream",
                    "stream": stream,
                    "data": json.loads(output)
                    if stream == "stdout"
                    else {"type": "error", "message": output.decode("utf-8")},
                }
            )

    async def run():
        assert process.stdin is not None
        assert process.stdout is not None
        assert process.stderr is not None

        await message_queue.put(
            {
                "type": "start",
                "data": {
                    "time": datetime.datetime.now().isoformat(),
                    "commit": read_head_commit(),
                },
            }
        )

        process.stdin.write(data.encode("utf-8"))
        await process.stdin.drain()
        process.stdin.close()

        await asyncio.gather(
            # process.stdin.drain(),
            watch_stream(process.stdout, "stdout"),
            watch_stream(process.stderr, "stderr"),
        )

        returncode: int | None = None
        try:
            returncode = await asyncio.wait_for(process.wait(), timeout=300)
        except asyncio.TimeoutError:
            ...

        await message_queue.put(
            {
                "type": "complete",
                "data": {
                    "status": "success" if returncode is not None else "error",
                    "returncode": -1 if returncode is None else -1,
                },
            }
        )

        await message_queue.join()

    loop = asyncio.get_event_loop()
    loop.create_task(run())

    async def stream_queue():
        while True:
            message = await message_queue.get()

            yield json.dumps(
                {
                    "event": message["type"],
                    "id": 0,
                    "retry": 0,
                    "data": message["data"],
                }
            )

            message_queue.task_done()

            if message["type"] == "complete":
                break

    return EventSourceResponse(stream_queue())


@app.post("/submit")
async def submit(
    year: int = Form(), day: int = Form(), verbose: bool = Form(), input: str = Form()
) -> dict[str, Any]:
    data = input.strip().replace("\r\n", "\n")

    args: list[str] = ["holt59-aoc", "--stdin"]
    if verbose:
        args += ["-v"]

    start = datetime.datetime.now()
    process = subprocess.run(
        args + ["--year", str(year), str(day)],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        input=data.encode("utf-8"),
        timeout=300,
    )
    time = datetime.datetime.now() - start

    return {
        "commit": read_head_commit(),
        "input": data,
        "error": process.returncode != 0,
        "stdout": process.stdout.decode("utf-8"),
        "stderr": process.stderr.decode("utf-8"),
        "time": time.total_seconds(),
    }


@app.get("/")
async def root():
    return {"message": "Hello World"}