"""Exercise cancellation while a real spawned worker is blocked before a result."""
from concurrent.futures import ThreadPoolExecutor
import os
from pathlib import Path
import subprocess
from threading import Event
import time

import pytest

from demo.execution import verification_reports
from demo.session import UISession, begin_ui, cancel_session, reset_session, stream_ui


def blocked_worker(task, connection, directory):
    os.setsid()
    child = subprocess.Popen(["sleep", "30"])
    Path(directory, "unfinished").write_text("temporary build")
    Path(task).write_text(f"{os.getpid()} {child.pid}\n{directory}")
    if not Path(task).name.startswith("build"):
        connection.send(("report", "ready"))
    time.sleep(30)


def completed_worker(task, connection, directory):
    os.setsid()
    connection.send(("report", "completed"))
    connection.send(("done", None))
    connection.close()


def assert_stopped(pid):
    # A killed compiler child may remain briefly as a zombie until reaped by init.
    result = subprocess.run(["ps", "-o", "stat=", "-p", str(pid)], capture_output=True, text=True)
    assert not result.stdout.strip() or result.stdout.strip().startswith("Z")


@pytest.mark.parametrize("stop", [cancel_session, reset_session])
def test_cancel_blocked_build_releases_worker_and_allows_next_run(tmp_path, stop):
    session = UISession()
    marker = tmp_path / "worker"

    def handler(*values, cancelled):
        for value in verification_reports(marker, cancelled=cancelled, _worker=blocked_worker):
            yield value, [], None, None, "", ""

    ticket = begin_ui(session, "source")[-1]
    old_signal = session.cancelled
    updates = stream_ui(handler, (), session, ticket, cancellable=True)
    next(updates)
    assert next(updates)[0] == "ready"
    pids, directory = marker.read_text().splitlines()
    with ThreadPoolExecutor(max_workers=1) as executor:
        waiting = executor.submit(lambda: list(updates))
        time.sleep(0.2)
        assert not waiting.done()
        stop(session)
        new_ticket = begin_ui(session, "new source")[-1]
        assert old_signal.is_set() and not session.cancelled.is_set()
        assert waiting.result(timeout=3) == []
    assert session.busy and session.generation == new_ticket
    assert not Path(directory).exists()
    for pid in map(int, pids.split()):
        assert_stopped(pid)
    assert list(verification_reports(None, _worker=completed_worker)) == ["completed"]




def test_close_stream_cleans_up_worker(tmp_path):
    marker = tmp_path / "worker"
    reports = verification_reports(marker, _worker=blocked_worker)
    assert next(reports) == "ready"
    reports.close()
    pids, directory = marker.read_text().splitlines()
    assert not Path(directory).exists()
    for pid in map(int, pids.split()):
        assert_stopped(pid)


def test_already_cancelled_does_not_start_worker(tmp_path):
    cancelled = Event()
    cancelled.set()
    marker = tmp_path / "worker"
    assert list(verification_reports(marker, cancelled=cancelled, _worker=blocked_worker)) == []
    assert not marker.exists()


def test_cancel_before_first_build_result(tmp_path):
    marker = tmp_path / "build-worker"
    cancelled = Event()
    with ThreadPoolExecutor(max_workers=1) as executor:
        waiting = executor.submit(lambda: list(verification_reports(
            marker, cancelled=cancelled, _worker=blocked_worker)))
        deadline = time.monotonic() + 10
        while not marker.exists():
            assert time.monotonic() < deadline
            time.sleep(0.02)
        assert not waiting.done()
        cancelled.set()
        assert waiting.result(timeout=3) == []
    pids, directory = marker.read_text().splitlines()
    assert not Path(directory).exists()
    for pid in map(int, pids.split()):
        assert_stopped(pid)


def solving_worker(task, connection, directory):
    """Block inside the solver call, rather than before verification starts."""
    os.setsid()
    import angr
    import claripy
    from types import SimpleNamespace
    from verifier.concolic import ConcolicInputs
    provider = ConcolicInputs(0, angr.load_shellcode(b"\xc3", arch="amd64").arch)
    symbol = provider.tracker.symbol("argument", "0", 32, 1)
    provider.tracker.branch(SimpleNamespace(scratch=SimpleNamespace(ins_addr=1)), symbol == 0)
    def blocked_solver(*args, **kwargs):
        Path(task).write_text(f"{os.getpid()}\n{directory}")
        time.sleep(30)
        return False
    claripy.Solver.satisfiable = blocked_solver
    provider.tracker.solve(provider.tracker.branches[0], provider.freeze(), 30000)
    connection.send(("done", None))


def test_cancel_during_concolic_solving(tmp_path):
    marker = tmp_path / "solver-worker"
    cancelled = Event()
    with ThreadPoolExecutor(max_workers=1) as executor:
        waiting = executor.submit(lambda: list(verification_reports(
            marker, cancelled=cancelled, _worker=solving_worker)))
        deadline = time.monotonic() + 10
        while not marker.exists():
            assert time.monotonic() < deadline
            time.sleep(0.02)
        cancelled.set()
        assert waiting.result(timeout=3) == []
    pid, directory = marker.read_text().splitlines()
    assert_stopped(int(pid))
    assert not Path(directory).exists()