"""Reproduce a v0.5.1 Tool Step commit gap using a real worker exit.

Run with the M-Agent test environment:
    python probe-tool-commit-window.py /path/to/M-agent

This diagnostic uses the versioned test fixtures, not production services.
"""

from __future__ import annotations

import asyncio
import json
import os
from pathlib import Path
import subprocess
import sys
import tempfile
from datetime import timedelta

source = Path(sys.argv[1]).resolve()
sys.path[:0] = [str(source / "src"), str(source / "tests")]

from test_m_agent_resolution import (  # noqa: E402
    CountingNotifier,
    RequestNotifyModel,
    build_registry,
)
from m_agent.adapters import FakeClock, PlaintextPayloadCodec, SQLiteRunStore  # noqa: E402
from m_agent.runtime import DEFAULT_LEASE_TTL, Runner, StepStatus, StepType  # noqa: E402


class JournalNotifier(CountingNotifier):
    def __init__(self, journal: Path):
        super().__init__()
        self.journal = journal

    async def invoke(self, request):
        with self.journal.open("a") as stream:
            stream.write("notification\n")
            stream.flush()
            os.fsync(stream.fileno())
        return await super().invoke(request)


class ExitAfterToolStep(SQLiteRunStore):
    async def record_step(self, step, **kwargs):
        await super().record_step(step, **kwargs)
        if step.step_type is StepType.TOOL and step.status is StepStatus.SUCCEEDED:
            os._exit(17)


async def run(directory: Path, worker: bool):
    clock = FakeClock()
    if not worker:
        clock.advance(DEFAULT_LEASE_TTL + timedelta(seconds=1))
    store_class = ExitAfterToolStep if worker else SQLiteRunStore
    store = store_class(
        str(directory / "run.db"),
        payload_codec=PlaintextPayloadCodec(),
        clock=clock,
    )
    runner = Runner(
        registry=build_registry(
            RequestNotifyModel(), JournalNotifier(directory / "journal.txt")
        ),
        store=store,
    )
    try:
        if worker:
            created = await runner.create_run(
                "support_agent", "1.0", "hello", run_id="commit-gap"
            )
            await runner.start_run(created.run_id)
            raise AssertionError("worker did not reach the commit gap")
        steps = await store.get_steps("commit-gap")
        attempts = await store.get_attempts("commit-gap")
        checkpoints = await store.get_checkpoints("commit-gap")
        count_before = len((directory / "journal.txt").read_text().splitlines())
        result = await runner.resume_run("commit-gap")
        count_after = len((directory / "journal.txt").read_text().splitlines())
        evidence = {
            "window": "after Tool Step SUCCEEDED commit, before Attempt/Checkpoint",
            "tool_step_status_before": [
                step.status.value for step in steps if step.step_type is StepType.TOOL
            ],
            "tool_attempt_status_before": [
                attempt.status.value for attempt in attempts
                if any(step.step_id == attempt.step_id and step.step_type is StepType.TOOL
                       for step in steps)
            ],
            "tool_checkpoints_before": sum(
                checkpoint.step_type is StepType.TOOL for checkpoint in checkpoints
            ),
            "external_count_before": count_before,
            "external_count_after": count_after,
            "resumed_status": result.status.value,
        }
        print(json.dumps(evidence, indent=2))
        assert count_before == 1 and count_after == 2, evidence
        assert result.status.value == "SUCCEEDED", evidence
    finally:
        store.close()


if len(sys.argv) > 2:
    asyncio.run(run(Path(sys.argv[2]), worker=True))
else:
    with tempfile.TemporaryDirectory(prefix="durable-run-proofread-") as temp:
        worker = subprocess.run(
            [sys.executable, str(Path(__file__).resolve()), str(source), temp],
            check=False,
        )
        assert worker.returncode == 17, worker.returncode
        asyncio.run(run(Path(temp), worker=False))