from __future__ import annotations

import asyncio
import time
from pathlib import Path

import pytest

from investibook.core import agent_runtime


class FakeSession:
    def __init__(
        self,
        *,
        idle_s: float = 0,
        busy: bool = False,
        workspace_id: str | None = None,
        user_id: str | None = None,
    ):
        self.frontend_id = "frontend"
        self.runtime_session_id = "acp"
        self.workspace_id = workspace_id
        self.user_id = user_id
        self.last_activity = time.monotonic() - idle_s
        self.shutdown_count = 0
        self._busy = busy

    def is_busy(self) -> bool:
        return self._busy

    def touch(self) -> None:
        self.last_activity = time.monotonic()

    async def shutdown(self) -> None:
        self.shutdown_count += 1


class FakeClient:
    def __init__(self, summary: dict[str, object]):
        self._msg_id = "msg-1"
        self.summary = summary
        self.reset_count = 0

    def assistant_completion_summary(self) -> dict[str, object]:
        return dict(self.summary)

    def _mark_open_tools_cancelled(self) -> None:
        pass

    def _seal_open_reasoning_parts(self) -> None:
        pass

    def flush_assistant_message(self) -> None:
        pass

    def reset_assistant_message(self) -> None:
        self._msg_id = None
        self.reset_count += 1


@pytest.fixture(autouse=True)
def clear_acp_sessions():
    agent_runtime.sessions.clear()
    agent_runtime._session_creation_locks.clear()
    agent_runtime._session_prewarm_tasks.clear()
    agent_runtime.shutdown_event.clear()
    yield
    agent_runtime.shutdown_event.set()
    agent_runtime.sessions.clear()
    agent_runtime._session_creation_locks.clear()
    agent_runtime._session_prewarm_tasks.clear()
    agent_runtime.shutdown_event.clear()


@pytest.mark.asyncio
async def test_session_prewarm_coalesces_duplicate_page_loads(monkeypatch):
    started = asyncio.Event()
    release = asyncio.Event()
    calls = []

    async def fake_get_or_create_session(**kwargs):
        calls.append(kwargs)
        started.set()
        await release.wait()

    monkeypatch.setattr(agent_runtime.cfg, "AGENT_SESSION_PREWARM_ENABLED", True)
    monkeypatch.setattr(
        agent_runtime, "get_or_create_session", fake_get_or_create_session
    )

    assert agent_runtime.schedule_session_prewarm(
        workspace_id="ws-1", user_id="user-1", agent_type="opencode", enabled=True
    )
    await started.wait()
    assert not agent_runtime.schedule_session_prewarm(
        workspace_id="ws-1", user_id="user-1", agent_type="opencode", enabled=True
    )
    release.set()
    await agent_runtime.cancel_session_prewarms()

    assert len(calls) == 1


@pytest.mark.asyncio
async def test_session_prewarm_is_disabled_without_user_opt_in(monkeypatch):
    monkeypatch.setattr(agent_runtime.cfg, "AGENT_SESSION_PREWARM_ENABLED", True)
    assert not agent_runtime.schedule_session_prewarm(
        workspace_id="ws-1", user_id="user-1", enabled=False
    )


@pytest.mark.asyncio
async def test_foreground_request_reuses_session_created_by_prewarm(monkeypatch):
    started = asyncio.Event()
    release = asyncio.Event()
    initialized = 0
    warmed = FakeSession(workspace_id="ws-1", user_id="user-1")

    async def fake_create_locked(dict_key, *_args, **_kwargs):
        nonlocal initialized
        existing = agent_runtime.sessions.get(dict_key)
        if existing is not None:
            return existing
        initialized += 1
        started.set()
        await release.wait()
        agent_runtime.sessions[dict_key] = warmed  # type: ignore[assignment]
        return warmed

    monkeypatch.setattr(agent_runtime.cfg, "AGENT_SESSION_PREWARM_ENABLED", True)
    monkeypatch.setattr(
        agent_runtime, "_get_or_create_session_locked", fake_create_locked
    )

    assert agent_runtime.schedule_session_prewarm(
        workspace_id="ws-1", user_id="user-1", agent_type="opencode", enabled=True
    )
    await started.wait()
    foreground = asyncio.create_task(
        agent_runtime.get_or_create_session(
            workspace_id="ws-1", user_id="user-1", agent_type="opencode"
        )
    )
    await asyncio.sleep(0)
    assert not foreground.done()

    release.set()
    assert await foreground is warmed
    pending = list(agent_runtime._session_prewarm_tasks.values())
    if pending:
        await asyncio.gather(*pending)
    assert initialized == 1


@pytest.mark.asyncio
async def test_release_acp_session_shutdowns_and_unregisters():
    session = FakeSession()
    agent_runtime.sessions["ws-1"] = session  # type: ignore[assignment]
    agent_runtime._session_creation_locks["ws-1"] = asyncio.Lock()

    released = await agent_runtime.release_agent_session("ws-1")

    assert released is True
    assert session.shutdown_count == 1
    assert "ws-1" not in agent_runtime.sessions
    assert "ws-1" not in agent_runtime._session_creation_locks


def test_agent_event_bridge_summarizes_completion():
    bridge = agent_runtime.AgentEventBridge(workspace_id="ws-1")
    bridge.start_assistant_message("msg-1")
    bridge._msg_text_parts.append("hello")
    bridge._msg_parts["text-0"] = {"type": "text", "id": "text-0", "text": "hello"}
    bridge._msg_parts_order.append("text-0")
    bridge._record_prompt_event("text")
    bridge._record_suppressed_event("session_compacted")

    assert bridge.assistant_completion_summary() == {
        "content_len": 5,
        "part_count": 1,
        "has_text": True,
        "has_parts": True,
        "event_counts": {"text": 1},
        "suppressed_counts": {"session_compacted": 1},
    }


@pytest.mark.asyncio
async def test_live_sessions_are_isolated_by_user_within_workspace():
    first = FakeSession(workspace_id="ws-1", user_id="user-1")
    second = FakeSession(workspace_id="ws-1", user_id="user-2")
    first_key = agent_runtime._session_registry_key("ws-1", "user-1")
    second_key = agent_runtime._session_registry_key("ws-1", "user-2")
    agent_runtime.sessions[first_key] = first  # type: ignore[assignment]
    agent_runtime.sessions[second_key] = second  # type: ignore[assignment]

    assert agent_runtime.get_agent_session("ws-1", "user-1") is first
    assert agent_runtime.get_agent_session("ws-1", "user-2") is second

    released = await agent_runtime.release_agent_session("ws-1", "user-1")

    assert released is True
    assert first.shutdown_count == 1
    assert second.shutdown_count == 0
    assert agent_runtime.get_agent_session("ws-1", "user-1") is None
    assert agent_runtime.get_agent_session("ws-1", "user-2") is second


@pytest.mark.asyncio
async def test_handle_completion_anomaly_invalidates_usage_only_session(tmp_path):
    workspace_dir = tmp_path / "ws-1"
    workspace_dir.mkdir()
    session = agent_runtime.AgentSession(
        "frontend",
        "opencode",
        workspace_dir=workspace_dir,
        user_id="user-1",
    )
    session.runtime_session_id = "runtime-1"

    async def _noop_persist() -> None:
        return None

    async def _stub_invalidate(reason: str, *, clear_persisted_session: bool = True) -> None:
        agent_runtime.sessions.pop("ws-1:user-1", None)
        agent_runtime._session_creation_locks.pop("ws-1:user-1", None)

    session.persist_assistant_message_snapshot = _noop_persist  # type: ignore[method-assign]
    session.invalidate_broken_session = _stub_invalidate  # type: ignore[method-assign]
    session.client = FakeClient(
        {
            "content_len": 0,
            "part_count": 0,
            "has_text": False,
            "has_parts": False,
            "event_counts": {"usage": 1},
            "suppressed_counts": {},
        }
    )
    agent_runtime.sessions["ws-1:user-1"] = session  # type: ignore[assignment]
    agent_runtime._session_creation_locks["ws-1:user-1"] = asyncio.Lock()

    with pytest.raises(agent_runtime.ACPEmptyCompletionError) as exc_info:
        await session._handle_completion_anomaly()

    assert exc_info.value.reason == "usage_only_completion"
    assert "ws-1:user-1" not in agent_runtime.sessions
    assert session.client.reset_count == 1
    assert session.client._msg_id is None


@pytest.mark.asyncio
async def test_reaper_releases_idle_sessions_but_skips_busy_sessions():
    idle = FakeSession(idle_s=120)
    busy = FakeSession(idle_s=120, busy=True)
    fresh = FakeSession(idle_s=1)
    agent_runtime.sessions["idle"] = idle  # type: ignore[assignment]
    agent_runtime.sessions["busy"] = busy  # type: ignore[assignment]
    agent_runtime.sessions["fresh"] = fresh  # type: ignore[assignment]

    task = asyncio.create_task(agent_runtime.reap_idle_agent_sessions(idle_timeout_s=30, interval_s=0))
    await asyncio.sleep(0.05)
    agent_runtime.shutdown_event.set()
    await task

    assert idle.shutdown_count == 1
    assert "idle" not in agent_runtime.sessions
    assert busy.shutdown_count == 0
    assert fresh.shutdown_count == 0
    assert "busy" in agent_runtime.sessions
    assert "fresh" in agent_runtime.sessions


@pytest.mark.asyncio
async def test_lru_eviction_prefers_oldest_idle_session(monkeypatch):
    monkeypatch.setattr(agent_runtime, "AGENT_MAX_SESSIONS", 2)
    oldest = FakeSession(idle_s=100)
    newer = FakeSession(idle_s=10)
    busy_old = FakeSession(idle_s=200, busy=True)
    agent_runtime.sessions["oldest"] = oldest  # type: ignore[assignment]
    agent_runtime.sessions["newer"] = newer  # type: ignore[assignment]
    agent_runtime.sessions["busy-old"] = busy_old  # type: ignore[assignment]

    await agent_runtime._evict_lru_sessions_locked()

    assert oldest.shutdown_count == 1
    assert newer.shutdown_count == 1
    assert busy_old.shutdown_count == 0
    assert "oldest" not in agent_runtime.sessions
    assert "newer" not in agent_runtime.sessions
    assert "busy-old" in agent_runtime.sessions


def test_acp_session_exposes_workspace_id_from_workspace_dir():
    session = agent_runtime.AgentSession(
        "frontend",
        "opencode",
        workspace_dir=Path("/tmp/ws-123"),
    )

    assert session.workspace_id == "ws-123"


@pytest.mark.asyncio
async def test_opencode_sessions_use_isolated_temporary_directories(
    monkeypatch, tmp_path
):
    monkeypatch.setattr(
        agent_runtime,
        "sandbox_provider_for_workspace",
        lambda workspace_id, *, user_id=None: None,
    )
    first_workspace = tmp_path / "ws-1"
    second_workspace = tmp_path / "ws-2"
    first_workspace.mkdir()
    second_workspace.mkdir()
    first = agent_runtime.AgentSession(
        "frontend-1", "opencode", workspace_dir=first_workspace
    )
    second = agent_runtime.AgentSession(
        "frontend-2", "opencode", workspace_dir=second_workspace
    )
    first_env: dict[str, str] = {}
    second_env: dict[str, str] = {}

    first._configure_runtime_temp_env(first_env)
    second._configure_runtime_temp_env(second_env)

    first_temp = Path(first_env["TMPDIR"])
    second_temp = Path(second_env["TMPDIR"])
    assert first_temp != second_temp
    assert first_temp.is_dir()
    assert second_temp.is_dir()
    for name in ("TMP", "TEMP", "BUN_TMPDIR"):
        assert first_env[name] == str(first_temp)
        assert second_env[name] == str(second_temp)

    await first.shutdown()
    await second.shutdown()

    assert not first_temp.exists()
    assert not second_temp.exists()


class FakeSandbox:
    name = "test-sandbox"
    workspace_root = "/workspace"

    def command_timeout(self, local_timeout_s: float) -> float:
        return max(local_timeout_s, 120.0)


def test_sandbox_execution_switch_defaults_to_local_and_supports_provider(
    monkeypatch,
):
    provider = None
    monkeypatch.setattr(
        agent_runtime,
        "sandbox_provider_for_workspace",
        lambda workspace_id, *, user_id=None: provider,
    )
    session = agent_runtime.AgentSession(
        "frontend",
        "opencode",
        workspace_dir=Path("/tmp/ws-123"),
        user_id="user-1",
    )
    assert session.uses_sandbox is False

    provider = FakeSandbox()
    session = agent_runtime.AgentSession(
        "frontend",
        "opencode",
        workspace_dir=Path("/tmp/ws-123"),
        user_id="user-1",
    )
    assert session.uses_sandbox is True
    assert session.sandbox_provider_name == "test-sandbox"


def test_sandbox_disables_backend_filesystem_capability(monkeypatch):
    monkeypatch.setattr(
        agent_runtime,
        "sandbox_provider_for_workspace",
        lambda workspace_id, *, user_id=None: FakeSandbox(),
    )
    session = agent_runtime.AgentSession(
        "frontend",
        "opencode",
        workspace_dir=Path("/tmp/ws-123"),
        user_id="user-1",
    )
    capabilities = session._client_capabilities()

    assert capabilities.fs.read_text_file is False
    assert capabilities.fs.write_text_file is False


def test_pi_rpc_uses_provider_command_timeout_in_sandbox():
    connection = object.__new__(agent_runtime.PiRpcConnection)
    connection.agent_session = type("Session", (), {"sandbox": FakeSandbox()})()

    assert connection._command_timeout(30.0) == 120.0

    connection.agent_session = type("Session", (), {"sandbox": None})()
    assert connection._command_timeout(30.0) == 30.0


@pytest.mark.asyncio
async def test_acp_prompt_injects_static_protocol_once(tmp_path):
    workspace_dir = tmp_path / "plain-workspace"
    workspace_dir.mkdir()
    session = agent_runtime.AgentSession("frontend", "opencode", workspace_dir=workspace_dir)
    session._initialized = True
    session.runtime_session_id = "acp-session"

    class FakeConn:
        def __init__(self):
            self.prompts = []

        async def prompt(self, *, session_id, prompt):
            self.prompts.append((session_id, prompt))
            class Response:
                field_meta = None

            return Response()

    conn = FakeConn()
    session.conn = conn  # type: ignore[assignment]

    await session.prompt("第一条用户消息")
    await session.prompt("第二条用户消息")

    first_text = conn.prompts[0][1][0].text
    second_text = conn.prompts[1][1][0].text

    assert "## Investibook Runtime Protocol" in first_text
    assert "Reply in the user's language" in first_text
    assert "Do not introduce yourself" not in first_text
    assert "Do not start TiDB diagnosis" not in first_text
    assert "Plan first" not in first_text
    assert "No links in text output" not in first_text
    assert "## Investibook Runtime Protocol" not in second_text
    assert second_text == "第二条用户消息"


@pytest.mark.asyncio
async def test_await_agent_rpc_or_exit_fails_fast_when_process_exits():
    session = agent_runtime.AgentSession("frontend", "opencode", workspace_dir=Path("/tmp/ws-123"))

    class FakeProc:
        async def wait(self) -> int:
            return 1

    session.proc = FakeProc()  # type: ignore[assignment]

    async def _never_returns() -> str:
        await asyncio.sleep(60)
        return "unreachable"

    with pytest.raises(RuntimeError, match=r"opencode exited during initialize \(rc=1\): fatal startup error"):
        await session._await_agent_rpc_or_exit(
            _never_returns(),
            phase="initialize",
            label="opencode",
            stderr_lines=["fatal startup error"],
        )