import argparse
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import shutil
import signal
import subprocess
import sys
import threading
import time

RUN_TIMEOUT_SECONDS = 600
CLOSE_TIMEOUT_SECONDS = 600


def log_event(message):
    print(f"{datetime.now(timezone.utc).isoformat()}: {message}", flush=True)


def forward_logs(stream, logfile, failures):
    try:
        while True:
            chunk = stream.read1(65536)
            if not chunk:
                break
            logfile.write(chunk)
            logfile.flush()
            sys.stdout.buffer.write(chunk)
            sys.stdout.buffer.flush()
    except Exception as error:
        failures.append(str(error))
    finally:
        stream.close()


def _stop_process_group(process):
    try:
        os.killpg(process.pid, signal.SIGKILL)
    except ProcessLookupError:
        pass
    return process.wait()


def execute_request(request, command):
    logfile_path = request / "logfile"
    failures = []
    timed_out = False
    return_code = None
    log_event(f"starting {request.name}")
    try:
        with logfile_path.open("wb") as logfile:
            deadline = time.monotonic() + RUN_TIMEOUT_SECONDS
            process = subprocess.Popen(
                command,
                cwd=request / "data",
                stdin=subprocess.DEVNULL,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                # Foundry는 setsid()를 차단하므로 현재 세션에서 그룹만 분리합니다.
                process_group=0,
            )
            forwarder = threading.Thread(
                target=forward_logs,
                args=(process.stdout, logfile, failures),
                daemon=True,
            )
            forwarder.start()
            try:
                return_code = process.wait(timeout=max(0, deadline - time.monotonic()))
            except subprocess.TimeoutExpired:
                timed_out = True
                return_code = _stop_process_group(process)
            forwarder.join(timeout=max(0, deadline - time.monotonic()))
            if forwarder.is_alive():
                timed_out = True
                return_code = _stop_process_group(process)
                forwarder.join()
    except OSError as error:
        failures.append(str(error))
        logfile_path.write_text(str(error), encoding="utf-8")

    status = {
        "return_code": return_code,
        "timed_out": timed_out,
        "errors": failures,
    }
    pending = request / "done_flag.pending"
    pending.write_text(json.dumps(status), encoding="utf-8")
    pending.replace(request / "done_flag")
    log_event(f"finished {request.name}: {json.dumps(status)}")


def serve(command, shared_directory):
    requests_directory = shared_directory / "requests"
    requests_directory.mkdir(parents=True, exist_ok=True)
    while True:
        for request in requests_directory.iterdir():
            if not request.name.startswith("request-") or request.is_symlink():
                continue
            if not request.is_dir():
                continue
            if (request / "close_flag").is_file():
                shutil.rmtree(request)
                continue
            if not (request / "start_flag").is_file() or (request / "done_flag").exists():
                continue
            execute_request(request, command)
            deadline = time.monotonic() + CLOSE_TIMEOUT_SECONDS
            while not (request / "close_flag").exists():
                if time.monotonic() >= deadline:
                    log_event(f"close acknowledgement timed out for {request.name}")
                    break
                time.sleep(0.1)
            shutil.rmtree(request)
        time.sleep(0.1)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("command", nargs=argparse.REMAINDER)
    arguments = parser.parse_args()
    command = arguments.command
    if command and command[0] == "--":
        command = command[1:]
    if not command:
        parser.error("A foreground command and its arguments are required")
    serve(command, Path(os.environ["SHARED_DIR"]))


if __name__ == "__main__":
    main()