#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import asyncio
import dataclasses
import socket
import traceback


@dataclasses.dataclass
class ScgiRequest:
    headers: dict[bytes, bytes]
    body: bytes


async def parse_scgi_request(reader: asyncio.StreamReader) -> ScgiRequest:
    hlen = int((await reader.readuntil(b':'))[:-1])
    header_raw = await reader.readexactly(hlen + 1)
    assert len(header_raw) >= 16, "invalid request: too short (< 16)"
    assert header_raw[-2:] == b'\0,', \
        fr"Invalid request: missing header/netstring terminator '\x00,', got {header_raw[-2:]!r}"
    header_list = header_raw[:-2].split(b'\0')
    assert len(header_list) % 2 == 0, \
        f"Invalid request: odd numbers of header entries (must be pairs), got {len(header_list)}"
    assert header_list[0] == b'CONTENT_LENGTH', \
        f"Invalid request: first header entry must be 'CONTENT_LENGTH', got {header_list[0]!r}"
    clen = int(header_list[1])
    headers = {}
    i = 0
    while i < len(header_list):
        key = header_list[i]
        value = header_list[i+1]
        i += 2
        assert not key in headers, f"Invalid request: duplicate header key {key!r}"
        headers[key] = value
    assert headers.get(b'SCGI') == b'1', "Invalid request: missing SCGI=1 header"
    body = await reader.readexactly(clen)
    return ScgiRequest(headers=headers, body=body)


async def handle_scgi(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    print("scgi-envcheck: Incoming connection", flush=True)
    try:
        req = await parse_scgi_request(reader)
        envvar = req.headers[b'QUERY_STRING']
        result = req.headers[envvar]
    except KeyboardInterrupt:
        raise
    except Exception as e:
        print(traceback.format_exc())
        writer.write(b"Status: 500\r\nContent-Type: text/plain\r\n\r\n" + str(e).encode())
    else:
        writer.write(b"Status: 200\r\nContent-Type: text/plain\r\n\r\n" + result)
    await writer.drain()
    writer.close()
    await writer.wait_closed()


async def main() -> None:
    sock = socket.socket(fileno=0)

    if sock.type == socket.AF_UNIX:
        server = await asyncio.start_unix_server(handle_scgi, sock=sock, start_serving=False)
    else:
        server = await asyncio.start_server(handle_scgi, sock=sock, start_serving=False)

    addr = server.sockets[0].getsockname()
    print(f'Serving on {addr}', flush=True)

    async with server:
        await server.serve_forever()


try:
    asyncio.run(main())
except KeyboardInterrupt:
    pass
