from __future__ import annotations

import json
import os
import signal
import subprocess
import sys
import time
from pathlib import Path
from unittest.mock import patch

import pytest

from adw_modules import agent_pi
from adw_modules.data_types import PiRequest


class FakeClock:
    def __init__(self) -> None:
        self.value = 0.0
        self.sleeps: list[float] = []

    def __call__(self) -> float:
        return self.value

    def sleep(self, seconds: float) -> None:
        self.sleeps.append(seconds)
        self.value += seconds


class InlineThread:
    def __init__(self, *, target, daemon: bool) -> None:
        self.target = target
        self.daemon = daemon

    def start(self) -> None:
        self.target()

    def join(self, timeout: float | None = None) -> None:
        return None

    def is_alive(self) -> bool:
        return False


class FakeProcess:
    def __init__(self, stdout: list[str] = (), stderr: list[str] = (), *,
                 exits_on_term: bool = True) -> None:
        self.pid = 731
        self.stdout = iter(stdout)
        self.stderr = iter(stderr)
        self.returncode: int | None = None
        self.exits_on_term = exits_on_term
        self.signals: list[int] = []

    def poll(self) -> int | None:
        return self.returncode

    def wait(self, timeout: float | None = None) -> int:
        if self.returncode is not None:
            return self.returncode
        if timeout is not None:
            raise subprocess.TimeoutExpired("fake-pi", timeout)
        raise AssertionError("unbounded wait before process termination")

    def signal(self, value: int) -> None:
        self.signals.append(value)
        if value == signal.SIGTERM and self.exits_on_term:
            self.returncode = -signal.SIGTERM
        if value == signal.SIGKILL:
            self.returncode = -signal.SIGKILL


def _wait_for_child_ready(process: subprocess.Popen, ready_path: Path) -> None:
    deadline = time.monotonic() + 5
    while not ready_path.exists():
        returncode = process.poll()
        if returncode is not None:
            pytest.fail(
                f"child readiness failed: leader exited with {returncode}; "
                f"ready_path={ready_path}"
            )
        if time.monotonic() >= deadline:
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            pytest.fail(
                f"child readiness timed out after 5s; leader_pid={process.pid}; "
                f"ready_path={ready_path}"
            )
        time.sleep(0.01)


def _request(tmp_path: Path) -> PiRequest:
    return PiRequest(prompt="prompt", system_prompt="system", model="fake",
                     session_id="session-1", session_dir=str(tmp_path / "sessions"),
                     raw_output_path=str(tmp_path / "raw.jsonl"), cwd=str(tmp_path))


def _run(tmp_path: Path, process: FakeProcess, clock: FakeClock, *,
         timeout: float, idle_timeout: float):
    events: list[dict] = []
    attempts: list[dict] = []

    def killpg(pid: int, value: int) -> None:
        assert pid == process.pid
        process.signal(value)

    with patch.object(agent_pi, "resolve_model", return_value=("fake", "fake")), \
         patch.object(agent_pi, "context_window", return_value=1000):
        with pytest.raises(TimeoutError) as raised:
            agent_pi.run(_request(tmp_path), on_event=events.append,
                         on_attempt_end=lambda _attempt, info: attempts.append(info),
                         timeout_seconds=timeout, idle_timeout_seconds=idle_timeout,
                         clock=clock, sleeper=clock.sleep,
                         popen_factory=lambda *_args, **_kwargs: process,
                         kill_process_group=killpg, thread_factory=InlineThread)
    return raised.value, events, attempts


def test_starting_uses_wall_timeout_without_arming_idle(tmp_path: Path) -> None:
    clock = FakeClock()
    process = FakeProcess(exits_on_term=True)

    error, events, attempts = _run(tmp_path, process, clock, timeout=1.0, idle_timeout=0.1)

    assert str(error) == "pi wall timeout: timed out after 1s"
    assert clock.value == 1.0 + agent_pi.TERMINATE_GRACE_SECONDS
    assert process.signals == [signal.SIGTERM, signal.SIGKILL]
    assert attempts == [{"returncode": -signal.SIGKILL, "signal": signal.SIGKILL,
                         "timed_out": True, "timeout_kind": "wall",
                         "stderr_path": str(tmp_path / "stderr.log"), "tokens": None,
                         "usage": None, "error": str(error),
                         "provider_failure": {"kind": "timeout", "detail": str(error),
                                              "retry_after_seconds": None,
                                              "resume_at": None}}]
    terminal = [event for event in events if event["type"] == "agent_attempt_end"]
    assert len(terminal) == 1
    assert terminal[0]["lifecycle"] == "terminal"


def test_streaming_idle_timeout_drains_output_and_escalates(tmp_path: Path) -> None:
    stdout = json.dumps({"type": "tool_execution_start", "toolCallId": "call-1"}) + "\n"
    clock = FakeClock()
    process = FakeProcess([stdout], ["provider stderr\n"], exits_on_term=False)

    error, events, attempts = _run(tmp_path, process, clock, timeout=2.0, idle_timeout=0.2)

    assert str(error) == "pi idle timeout: timed out after 0.2s"
    assert clock.value == 0.2 + agent_pi.TERMINATE_GRACE_SECONDS
    assert process.signals == [signal.SIGTERM, signal.SIGKILL]
    assert (tmp_path / "raw.jsonl").read_text() == stdout
    assert (tmp_path / "stderr.log").read_text() == "provider stderr\n"
    assert [event["type"] for event in events] == ["agent_attempt_start", "tool_execution_start", "agent_attempt_end"]
    assert attempts[0]["timeout_kind"] == "idle"
    assert attempts[0]["returncode"] == -signal.SIGKILL


def test_timeout_kills_group_and_drains_child_streams(tmp_path: Path) -> None:
    child = """import json
from pathlib import Path
import signal
import sys
import time
signal.signal(signal.SIGTERM, signal.SIG_IGN)
print(json.dumps({'type': 'child_stdout'}), flush=True)
print('child stderr', file=sys.stderr, flush=True)
Path(sys.argv[1]).write_text('ready')
while True:
    time.sleep(1)
"""
    ready_path = tmp_path / "child-ready"
    leader = f"""import json
from pathlib import Path
import subprocess
import sys
import time
subprocess.Popen([sys.executable, '-u', '-c', {child!r}, {str(ready_path)!r}])
while not Path({str(ready_path)!r}).exists():
    time.sleep(.01)
print(json.dumps({{'type': 'leader_stdout'}}), flush=True)
print('leader stderr', file=sys.stderr, flush=True)
while True:
    time.sleep(1)
"""
    signals: list[int] = []
    events: list[dict] = []

    def popen_factory(_cmd, **kwargs):
        process = subprocess.Popen([sys.executable, "-u", "-c", leader], **kwargs)
        _wait_for_child_ready(process, ready_path)
        return process

    def killpg(pid: int, value: int) -> None:
        signals.append(value)
        os.killpg(pid, value)

    started = time.monotonic()
    with patch.object(agent_pi, "resolve_model", return_value=("fake", "fake")), \
         patch.object(agent_pi, "context_window", return_value=1000):
        with pytest.raises(TimeoutError, match="pi wall timeout"):
            agent_pi.run(_request(tmp_path), on_event=events.append,
                         timeout_seconds=1, idle_timeout_seconds=5,
                         popen_factory=popen_factory, kill_process_group=killpg)

    assert time.monotonic() - started < 5
    assert signals == [signal.SIGTERM, signal.SIGKILL]
    assert '"type": "leader_stdout"' in (tmp_path / "raw.jsonl").read_text()
    assert '"type": "child_stdout"' in (tmp_path / "raw.jsonl").read_text()
    stderr = (tmp_path / "stderr.log").read_text()
    assert "leader stderr" in stderr
    assert "child stderr" in stderr
    assert [event["type"] for event in events].count("agent_attempt_end") == 1


def test_timeout_preserves_grace_when_leader_exits_before_child(tmp_path: Path) -> None:
    child = """import json
from pathlib import Path
import signal
import sys
import time
signal.signal(signal.SIGTERM, signal.SIG_IGN)
print(json.dumps({'type': 'child_stdout'}), flush=True)
Path(sys.argv[1]).write_text('ready')
while True:
    time.sleep(1)
"""
    ready_path = tmp_path / "child-ready"
    leader = f"""import json
from pathlib import Path
import subprocess
import sys
import time
subprocess.Popen([sys.executable, '-u', '-c', {child!r}, {str(ready_path)!r}])
while not Path({str(ready_path)!r}).exists():
    time.sleep(.01)
print(json.dumps({{'type': 'leader_stdout'}}), flush=True)
while True:
    time.sleep(1)
"""
    signals: list[tuple[int, float]] = []

    def popen_factory(_cmd, **kwargs):
        process = subprocess.Popen([sys.executable, "-u", "-c", leader], **kwargs)
        _wait_for_child_ready(process, ready_path)
        return process

    def killpg(pid: int, value: int) -> None:
        signals.append((value, time.monotonic()))
        os.killpg(pid, value)

    with patch.object(agent_pi, "resolve_model", return_value=("fake", "fake")), \
         patch.object(agent_pi, "context_window", return_value=1000):
        with pytest.raises(TimeoutError, match="pi wall timeout"):
            agent_pi.run(_request(tmp_path), timeout_seconds=1, idle_timeout_seconds=5,
                         popen_factory=popen_factory, kill_process_group=killpg)

    assert [value for value, _time in signals] == [signal.SIGTERM, signal.SIGKILL]
    assert signals[1][1] - signals[0][1] >= agent_pi.TERMINATE_GRACE_SECONDS - 0.1
    assert '"type": "leader_stdout"' in (tmp_path / "raw.jsonl").read_text()
    assert '"type": "child_stdout"' in (tmp_path / "raw.jsonl").read_text()