"""Tests for unified_session feature.

Covers:
- AgentLoop._dispatch() rewrites session_key to "unified:default" when enabled
- Existing session_key_override is respected (not overwritten)
- Feature is off by default (no behavior change for existing users)
- Config schema serialises unified_session as camelCase "unifiedSession"
- onboard-generated config.json contains "unifiedSession" key
- /new command correctly clears the shared session in unified mode
- /new is NOT a priority command (goes through _dispatch, key rewrite applies)
- Context window consolidation is unaffected by unified_session
"""

import asyncio
import json
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.file_state import FileStateStore
from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.command.builtin import cmd_new, register_builtin_commands
from nanobot.command.router import CommandContext, CommandRouter
from nanobot.config.schema import AgentDefaults, Config
from nanobot.providers.base import GenerationSettings
from nanobot.session.keys import UNIFIED_SESSION_KEY
from nanobot.session.manager import Session, SessionManager
from nanobot.utils.llm_runtime import LLMRuntime

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_loop(tmp_path: Path, unified_session: bool = False) -> AgentLoop:
    """Create a minimal AgentLoop for dispatch-level tests."""
    bus = MessageBus()
    provider = MagicMock()
    provider.get_default_model.return_value = "test-model"

    with patch("nanobot.agent.loop.SessionManager"), \
         patch("nanobot.agent.loop.SubagentManager") as mock_sub_mgr:
        mock_sub_mgr.return_value.cancel_by_session = AsyncMock(return_value=0)
        loop = AgentLoop(
            bus=bus,
            provider=provider,
            workspace=tmp_path,
            unified_session=unified_session,
        )
    return loop


def _runtime(provider) -> LLMRuntime:
    provider.generation = GenerationSettings(max_tokens=100)
    return LLMRuntime.capture(
        provider,
        "test-model",
        context_window_tokens=1000,
    )


def _make_msg(channel: str = "telegram", chat_id: str = "111",
              session_key_override: str | None = None) -> InboundMessage:
    return InboundMessage(
        channel=channel,
        chat_id=chat_id,
        sender_id="user1",
        content="hello",
        session_key_override=session_key_override,
    )


# ---------------------------------------------------------------------------
# TestUnifiedSessionDispatch — core behaviour
# ---------------------------------------------------------------------------

class TestUnifiedSessionDispatch:
    """AgentLoop._dispatch() session key rewriting logic."""

    @pytest.mark.asyncio
    async def test_unified_session_rewrites_key_to_unified_default(self, tmp_path: Path):
        """When unified_session=True, all messages use 'unified:default' as session key."""
        loop = _make_loop(tmp_path, unified_session=True)

        captured: list[str] = []

        async def fake_process(msg, **kwargs):
            captured.append(msg.session_key)
            return None

        loop._process_message = fake_process  # type: ignore[method-assign]

        msg = _make_msg(channel="telegram", chat_id="111")
        await loop._dispatch(msg)

        assert captured == ["unified:default"]

    @pytest.mark.asyncio
    async def test_unified_session_different_channels_share_same_key(self, tmp_path: Path):
        """Messages from different channels all resolve to the same session key."""
        loop = _make_loop(tmp_path, unified_session=True)

        captured: list[str] = []

        async def fake_process(msg, **kwargs):
            captured.append(msg.session_key)
            return None

        loop._process_message = fake_process  # type: ignore[method-assign]

        await loop._dispatch(_make_msg(channel="telegram", chat_id="111"))
        await loop._dispatch(_make_msg(channel="discord", chat_id="222"))
        await loop._dispatch(_make_msg(channel="cli", chat_id="direct"))

        assert captured == ["unified:default", "unified:default", "unified:default"]

    @pytest.mark.asyncio
    async def test_unified_session_disabled_preserves_original_key(self, tmp_path: Path):
        """When unified_session=False (default), session key is channel:chat_id as usual."""
        loop = _make_loop(tmp_path, unified_session=False)

        captured: list[str] = []

        async def fake_process(msg, **kwargs):
            captured.append(msg.session_key)
            return None

        loop._process_message = fake_process  # type: ignore[method-assign]

        msg = _make_msg(channel="telegram", chat_id="999")
        await loop._dispatch(msg)

        assert captured == ["telegram:999"]

    @pytest.mark.asyncio
    async def test_unified_session_respects_existing_override(self, tmp_path: Path):
        """If session_key_override is already set (e.g. Telegram thread), it is NOT overwritten."""
        loop = _make_loop(tmp_path, unified_session=True)

        captured: list[str] = []

        async def fake_process(msg, **kwargs):
            captured.append(msg.session_key)
            return None

        loop._process_message = fake_process  # type: ignore[method-assign]

        msg = _make_msg(channel="telegram", chat_id="111", session_key_override="telegram:thread:42")
        await loop._dispatch(msg)

        assert captured == ["telegram:thread:42"]

    def test_unified_session_default_is_false(self, tmp_path: Path):
        """unified_session defaults to False — no behavior change for existing users."""
        loop = _make_loop(tmp_path)
        assert loop._unified_session is False


# ---------------------------------------------------------------------------
# TestUnifiedSessionConfig — schema & serialisation
# ---------------------------------------------------------------------------

class TestUnifiedSessionConfig:
    """Config schema and onboard serialisation for unified_session."""

    def test_agent_defaults_unified_session_default_is_false(self):
        """AgentDefaults.unified_session defaults to False."""
        defaults = AgentDefaults()
        assert defaults.unified_session is False

    def test_agent_defaults_unified_session_can_be_enabled(self):
        """AgentDefaults.unified_session can be set to True."""
        defaults = AgentDefaults(unified_session=True)
        assert defaults.unified_session is True

    def test_config_serialises_unified_session_as_camel_case(self):
        """model_dump(by_alias=True) outputs 'unifiedSession' (camelCase) for JSON."""
        config = Config()
        data = config.model_dump(mode="json", by_alias=True)
        agents_defaults = data["agents"]["defaults"]
        assert "unifiedSession" in agents_defaults
        assert agents_defaults["unifiedSession"] is False

    def test_config_parses_unified_session_from_camel_case(self):
        """Config can be loaded from JSON with camelCase 'unifiedSession'."""
        raw = {"agents": {"defaults": {"unifiedSession": True}}}
        config = Config.model_validate(raw)
        assert config.agents.defaults.unified_session is True

    def test_config_parses_unified_session_from_snake_case(self):
        """Config also accepts snake_case 'unified_session' (populate_by_name=True)."""
        raw = {"agents": {"defaults": {"unified_session": True}}}
        config = Config.model_validate(raw)
        assert config.agents.defaults.unified_session is True

    def test_onboard_generated_config_contains_unified_session(self, tmp_path: Path):
        """save_config() writes 'unifiedSession' into config.json (simulates nanobot onboard)."""
        from nanobot.config.loader import save_config

        config = Config()
        config_path = tmp_path / "config.json"
        save_config(config, config_path)

        with open(config_path, encoding="utf-8") as f:
            data = json.load(f)

        agents_defaults = data["agents"]["defaults"]
        assert "unifiedSession" in agents_defaults, (
            "onboard-generated config.json must contain 'unifiedSession' key"
        )
        assert agents_defaults["unifiedSession"] is False


# ---------------------------------------------------------------------------
# TestCmdNewUnifiedSession — /new command behaviour in unified mode
# ---------------------------------------------------------------------------

class TestCmdNewUnifiedSession:
    """/new command routing and session-clear behaviour in unified mode."""

    def test_new_is_not_a_priority_command(self):
        """/new must NOT be in the priority table — it must go through _dispatch()
        so the unified session key rewrite applies before cmd_new runs."""
        router = CommandRouter()
        register_builtin_commands(router)
        assert router.is_priority("/new") is False

    def test_new_is_an_exact_command(self):
        """/new must be registered as an exact command."""
        router = CommandRouter()
        register_builtin_commands(router)
        assert "/new" in router._exact

    @pytest.mark.asyncio
    async def test_cmd_new_clears_unified_session(self, tmp_path: Path):
        """cmd_new called with key='unified:default' clears the shared session."""
        sessions = SessionManager(tmp_path)

        # Pre-populate the shared session with some messages
        shared = sessions.get_or_create("unified:default")
        shared.add_message("user", "hello from telegram")
        shared.add_message("assistant", "hi there")
        sessions.save(shared)
        assert len(sessions.get_or_create("unified:default").messages) == 2
        expected_snapshot = list(shared.messages)

        # schedule_background is a *sync* method that schedules a coroutine via
        # asyncio.create_task().  Mirror that exactly so the coroutine is consumed
        # and no RuntimeWarning is emitted.
        admitted_runtime = MagicMock(name="admitted_runtime")
        file_state_store = FileStateStore()
        previous_file_state = file_state_store.for_session("unified:default")
        tracked_file = tmp_path / "tracked.txt"
        tracked_file.write_text("tracked", encoding="utf-8")
        previous_file_state.record_read(tracked_file)
        loop = SimpleNamespace(
            sessions=sessions,
            consolidator=SimpleNamespace(archive_session=AsyncMock(return_value=True)),
            _cancel_active_tasks=AsyncMock(return_value=0),
            discard_session_file_state=file_state_store.discard,
            llm_runtime=MagicMock(return_value=MagicMock()),
            schedule_background=lambda coro: asyncio.ensure_future(coro),
        )

        msg = InboundMessage(
            channel="telegram", sender_id="user1", chat_id="111", content="/new",
            session_key_override="unified:default",  # as _dispatch() would set it
        )
        ctx = CommandContext(
            msg=msg,
            session=None,
            key="unified:default",
            raw="/new",
            loop=loop,
            runtime=admitted_runtime,
        )

        result = await cmd_new(ctx)

        assert "New session started" in result.content
        # Invalidate cache and reload from disk to confirm persistence
        sessions.invalidate("unified:default")
        reloaded = sessions.get_or_create("unified:default")
        assert reloaded.messages == []
        reset_file_state = file_state_store.for_session("unified:default")
        assert reset_file_state is not previous_file_state
        assert reset_file_state.is_unchanged(tracked_file) is False
        archived = loop.consolidator.archive_session.call_args.args[0]
        assert archived.key == "unified:default"
        assert archived.messages == expected_snapshot
        assert archived.last_consolidated == 0
        loop.consolidator.archive_session.assert_called_once_with(
            archived,
            archive_end=len(expected_snapshot),
            runtime=admitted_runtime,
        )
        loop.llm_runtime.assert_not_called()

    @pytest.mark.asyncio
    async def test_cmd_new_in_unified_mode_does_not_affect_other_sessions(self, tmp_path: Path):
        """Clearing unified:default must not touch other sessions on disk."""
        sessions = SessionManager(tmp_path)

        other = sessions.get_or_create("discord:999")
        other.add_message("user", "discord message")
        sessions.save(other)

        shared = sessions.get_or_create("unified:default")
        shared.add_message("user", "shared message")
        sessions.save(shared)

        loop = SimpleNamespace(
            sessions=sessions,
            consolidator=SimpleNamespace(archive_session=AsyncMock(return_value=True)),
            _cancel_active_tasks=AsyncMock(return_value=0),
            discard_session_file_state=MagicMock(),
            runtime_for_session=MagicMock(return_value=MagicMock()),
            schedule_background=lambda coro: asyncio.ensure_future(coro),
        )

        msg = InboundMessage(
            channel="telegram", sender_id="user1", chat_id="111", content="/new",
            session_key_override="unified:default",
        )
        ctx = CommandContext(msg=msg, session=None, key="unified:default", raw="/new", loop=loop)
        await cmd_new(ctx)

        sessions.invalidate("unified:default")
        sessions.invalidate("discord:999")
        assert sessions.get_or_create("unified:default").messages == []
        assert len(sessions.get_or_create("discord:999").messages) == 1


# ---------------------------------------------------------------------------
# TestConsolidationUnaffectedByUnifiedSession — consolidation is key-agnostic
# ---------------------------------------------------------------------------

class TestConsolidationUnaffectedByUnifiedSession:
    """maybe_consolidate_by_tokens() behaviour is identical regardless of session key."""

    @pytest.mark.asyncio
    async def test_consolidation_skips_empty_session_for_unified_key(self):
        """Empty unified:default session → consolidation exits immediately, archive not called."""
        from nanobot.agent.memory import Consolidator, MemoryStore

        store = MagicMock(spec=MemoryStore)
        mock_provider = MagicMock()
        mock_provider.chat_with_retry = AsyncMock(return_value=MagicMock(content="summary"))
        runtime = _runtime(mock_provider)
        # Use spec= so MagicMock doesn't auto-generate AsyncMock for non-async methods,
        # which would leave unawaited coroutines and trigger RuntimeWarning.
        sessions = MagicMock(spec=SessionManager)

        consolidator = Consolidator(
            store=store,
            sessions=sessions,
            build_messages=MagicMock(return_value=[]),
            get_tool_definitions=MagicMock(return_value=[]),
        )
        consolidator.archive_session = AsyncMock()

        session = Session(key="unified:default")
        session.messages = []
        sessions.get_or_create.return_value = session

        await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)

        consolidator.archive_session.assert_not_called()

    @pytest.mark.asyncio
    async def test_consolidation_behaviour_identical_for_any_key(self):
        """Archive call count is the same for 'telegram:123' and 'unified:default'
        under identical token conditions."""
        from nanobot.agent.memory import Consolidator, MemoryStore

        archive_calls: dict[str, int] = {}

        for key in ("telegram:123", "unified:default"):
            store = MagicMock(spec=MemoryStore)
            mock_provider = MagicMock()
            mock_provider.chat_with_retry = AsyncMock(return_value=MagicMock(content="summary"))
            runtime = _runtime(mock_provider)
            sessions = MagicMock(spec=SessionManager)

            consolidator = Consolidator(
                store=store,
                sessions=sessions,
                build_messages=MagicMock(return_value=[]),
                get_tool_definitions=MagicMock(return_value=[]),
            )

            session = Session(key=key)
            session.messages = []  # empty → exits immediately for both keys
            sessions.get_or_create.return_value = session

            consolidator.archive_session = AsyncMock()
            await consolidator.maybe_consolidate_by_tokens(
                session,
                runtime=runtime,
            )
            archive_calls[key] = consolidator.archive_session.call_count

        assert archive_calls["telegram:123"] == archive_calls["unified:default"] == 0

    @pytest.mark.asyncio
    async def test_consolidation_triggers_when_over_budget_unified_key(self):
        """When tokens exceed budget, consolidation attempts to find a boundary —
        behaviour is identical to any other session key."""
        from nanobot.agent.memory import Consolidator, MemoryStore

        store = MagicMock(spec=MemoryStore)
        mock_provider = MagicMock()
        runtime = _runtime(mock_provider)
        sessions = MagicMock(spec=SessionManager)

        consolidator = Consolidator(
            store=store,
            sessions=sessions,
            build_messages=MagicMock(return_value=[]),
            get_tool_definitions=MagicMock(return_value=[]),
        )

        session = Session(key="unified:default")
        session.messages = [{"role": "user", "content": "msg"}]
        sessions.get_or_create.return_value = session

        # Simulate over-budget: estimated > budget
        consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(950, "tiktoken"))
        # No valid boundary found → returns gracefully without archiving
        consolidator.pick_consolidation_boundary = MagicMock(return_value=None)
        consolidator.archive_session = AsyncMock()

        await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)

        # estimate was called (consolidation was attempted)
        consolidator.estimate_session_prompt_tokens.assert_called_once_with(
            session,
            runtime=runtime,
        )
        # but archive was not called (no valid boundary)
        consolidator.archive_session.assert_not_called()


# ---------------------------------------------------------------------------
# TestStopCommandWithUnifiedSession — /stop command integration
# ---------------------------------------------------------------------------


class TestStopCommandWithUnifiedSession:
    """Verify /stop command works correctly with unified session enabled."""

    @pytest.mark.asyncio
    async def test_active_tasks_use_effective_key_in_unified_mode(self, tmp_path: Path):
        """When unified_session=True, tasks are stored under UNIFIED_SESSION_KEY."""
        loop = _make_loop(tmp_path, unified_session=True)

        # Create a message from telegram channel
        msg = _make_msg(channel="telegram", chat_id="123456")

        # Mock _dispatch to complete immediately
        async def fake_dispatch(m):
            pass

        loop._dispatch = fake_dispatch  # type: ignore[method-assign]

        # Simulate the task creation flow (from _run loop)
        effective_key = UNIFIED_SESSION_KEY if loop._unified_session and not msg.session_key_override else msg.session_key
        task = asyncio.create_task(loop._dispatch(msg))
        loop._active_tasks.setdefault(effective_key, set()).add(task)

        # Wait for task to complete
        await task

        # Verify the task is stored under UNIFIED_SESSION_KEY, not the original channel:chat_id
        assert UNIFIED_SESSION_KEY in loop._active_tasks
        assert "telegram:123456" not in loop._active_tasks

    @pytest.mark.asyncio
    async def test_stop_command_finds_task_in_unified_mode(self, tmp_path: Path):
        """cmd_stop can cancel tasks when unified_session=True."""
        from nanobot.command.builtin import cmd_stop

        loop = _make_loop(tmp_path, unified_session=True)

        # Create a long-running task stored under UNIFIED_SESSION_KEY
        async def long_running():
            await asyncio.sleep(10)  # Will be cancelled

        task = asyncio.create_task(long_running())
        loop._active_tasks[UNIFIED_SESSION_KEY] = {task}

        # Create a message that would have session_key=UNIFIED_SESSION_KEY after dispatch
        msg = InboundMessage(
            channel="telegram",
            chat_id="123456",
            sender_id="user1",
            content="/stop",
            session_key_override=UNIFIED_SESSION_KEY,  # Simulate post-dispatch state
        )

        ctx = CommandContext(msg=msg, session=None, key=UNIFIED_SESSION_KEY, raw="/stop", loop=loop)

        # Execute /stop
        result = await cmd_stop(ctx)

        # Verify task was cancelled
        assert task.cancelled() or task.done()
        assert "Stopped 1 task" in result.content

    @pytest.mark.asyncio
    async def test_stop_command_uses_effective_key_without_session_override(self, tmp_path: Path):
        """Priority /stop must cancel the unified session even before dispatch rewrites the message."""
        from nanobot.command.builtin import cmd_stop

        loop = _make_loop(tmp_path, unified_session=True)

        async def long_running():
            await asyncio.sleep(10)

        task = asyncio.create_task(long_running())
        loop._active_tasks[UNIFIED_SESSION_KEY] = {task}
        msg = InboundMessage(
            channel="telegram",
            chat_id="123456",
            sender_id="user1",
            content="/stop",
        )
        ctx = CommandContext(msg=msg, session=None, key=UNIFIED_SESSION_KEY, raw="/stop", loop=loop)

        result = await cmd_stop(ctx)

        assert task.cancelled() or task.done()
        assert "Stopped 1 task" in result.content

    @pytest.mark.asyncio
    async def test_stop_command_cross_channel_in_unified_mode(self, tmp_path: Path):
        """In unified mode, /stop from one channel cancels tasks from another channel."""
        from nanobot.command.builtin import cmd_stop

        loop = _make_loop(tmp_path, unified_session=True)

        # Create tasks from different channels, all stored under UNIFIED_SESSION_KEY
        async def long_running():
            await asyncio.sleep(10)

        task1 = asyncio.create_task(long_running())
        task2 = asyncio.create_task(long_running())
        loop._active_tasks[UNIFIED_SESSION_KEY] = {task1, task2}

        # /stop from discord should cancel tasks started from telegram
        msg = InboundMessage(
            channel="discord",
            chat_id="789012",
            sender_id="user2",
            content="/stop",
            session_key_override=UNIFIED_SESSION_KEY,
        )

        ctx = CommandContext(msg=msg, session=None, key=UNIFIED_SESSION_KEY, raw="/stop", loop=loop)

        result = await cmd_stop(ctx)

        # Both tasks should be cancelled
        assert "Stopped 2 task" in result.content