"""Tests for Zellij metadata in spawned agent sessions (bd-30o).

Verifies that:
- amc-hook writes correct zellij_session and zellij_pane from env vars
- amc-hook writes spawn_id from AMC_SPAWN_ID env var
- Non-Zellij agents (missing env vars) get empty strings gracefully
- Spawn mixin passes AMC_SPAWN_ID and correct pane name to Zellij
"""

import json
import os
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch

# ---------------------------------------------------------------------------
# Import hook module (no .py extension)
# ---------------------------------------------------------------------------
hook_path = Path(__file__).parent.parent / "bin" / "amc-hook"
amc_hook = types.ModuleType("amc_hook")
amc_hook.__file__ = str(hook_path)
code = compile(hook_path.read_text(), hook_path, "exec")
exec(code, amc_hook.__dict__)  # noqa: S102 - loading local module

# Import spawn mixin (after hook loading - intentional)
import amc_server.mixins.spawn as spawn_mod  # noqa: E402


# ===========================================================================
# Hook: Zellij env vars -> session file
# ===========================================================================

class TestHookWritesZellijMetadata(unittest.TestCase):
    """Verify amc-hook writes Zellij metadata from environment variables."""

    def _run_hook(self, session_id, event, env_overrides=None, existing_session=None):
        """Helper: run the hook main() with controlled env and stdin."""
        with tempfile.TemporaryDirectory() as tmpdir:
            sessions_dir = Path(tmpdir) / "sessions"
            events_dir = Path(tmpdir) / "events"
            sessions_dir.mkdir()
            events_dir.mkdir()

            if existing_session:
                session_file = sessions_dir / f"{session_id}.json"
                session_file.write_text(json.dumps(existing_session))

            hook_input = json.dumps({
                "hook_event_name": event,
                "session_id": session_id,
                "cwd": "/test/project",
                "last_assistant_message": "",
            })

            env = {
                "CLAUDE_PROJECT_DIR": "/test/project",
            }
            if env_overrides:
                env.update(env_overrides)

            with patch.object(amc_hook, "SESSIONS_DIR", sessions_dir), \
                 patch.object(amc_hook, "EVENTS_DIR", events_dir), \
                 patch("sys.stdin.read", return_value=hook_input), \
                 patch.dict(os.environ, env, clear=False):
                amc_hook.main()

            session_file = sessions_dir / f"{session_id}.json"
            if session_file.exists():
                return json.loads(session_file.read_text())
            return None

    def test_session_start_writes_zellij_session(self):
        """SessionStart with ZELLIJ_SESSION_NAME writes zellij_session."""
        data = self._run_hook("test-sess", "SessionStart", env_overrides={
            "ZELLIJ_SESSION_NAME": "infra",
            "ZELLIJ_PANE_ID": "7",
        })
        self.assertIsNotNone(data)
        self.assertEqual(data["zellij_session"], "infra")

    def test_session_start_writes_zellij_pane(self):
        """SessionStart with ZELLIJ_PANE_ID writes zellij_pane."""
        data = self._run_hook("test-sess", "SessionStart", env_overrides={
            "ZELLIJ_SESSION_NAME": "infra",
            "ZELLIJ_PANE_ID": "7",
        })
        self.assertIsNotNone(data)
        self.assertEqual(data["zellij_pane"], "7")

    def test_session_start_writes_spawn_id(self):
        """SessionStart with AMC_SPAWN_ID writes spawn_id."""
        data = self._run_hook("test-sess", "SessionStart", env_overrides={
            "AMC_SPAWN_ID": "abc-123-def",
            "ZELLIJ_SESSION_NAME": "infra",
            "ZELLIJ_PANE_ID": "3",
        })
        self.assertIsNotNone(data)
        self.assertEqual(data["spawn_id"], "abc-123-def")

    def test_no_zellij_env_vars_writes_empty_strings(self):
        """Without Zellij env vars, zellij_session and zellij_pane are empty."""
        with tempfile.TemporaryDirectory() as tmpdir:
            sessions_dir = Path(tmpdir) / "sessions"
            events_dir = Path(tmpdir) / "events"
            sessions_dir.mkdir()
            events_dir.mkdir()

            hook_input = json.dumps({
                "hook_event_name": "SessionStart",
                "session_id": "no-zellij-sess",
                "cwd": "/test/project",
                "last_assistant_message": "",
            })

            # Ensure Zellij vars are removed from env
            clean_env = os.environ.copy()
            clean_env.pop("ZELLIJ_SESSION_NAME", None)
            clean_env.pop("ZELLIJ_PANE_ID", None)
            clean_env.pop("AMC_SPAWN_ID", None)
            clean_env["CLAUDE_PROJECT_DIR"] = "/test/project"

            with patch.object(amc_hook, "SESSIONS_DIR", sessions_dir), \
                 patch.object(amc_hook, "EVENTS_DIR", events_dir), \
                 patch("sys.stdin.read", return_value=hook_input), \
                 patch.dict(os.environ, clean_env, clear=True):
                amc_hook.main()

            session_file = sessions_dir / "no-zellij-sess.json"
            data = json.loads(session_file.read_text())

        self.assertEqual(data["zellij_session"], "")
        self.assertEqual(data["zellij_pane"], "")
        self.assertNotIn("spawn_id", data)

    def test_no_spawn_id_env_omits_field(self):
        """Without AMC_SPAWN_ID, spawn_id key is absent (not empty)."""
        clean_env = os.environ.copy()
        clean_env.pop("AMC_SPAWN_ID", None)
        clean_env["ZELLIJ_SESSION_NAME"] = "infra"
        clean_env["ZELLIJ_PANE_ID"] = "5"
        clean_env["CLAUDE_PROJECT_DIR"] = "/test/project"

        with tempfile.TemporaryDirectory() as tmpdir:
            sessions_dir = Path(tmpdir) / "sessions"
            events_dir = Path(tmpdir) / "events"
            sessions_dir.mkdir()
            events_dir.mkdir()

            hook_input = json.dumps({
                "hook_event_name": "SessionStart",
                "session_id": "no-spawn-sess",
                "cwd": "/test/project",
                "last_assistant_message": "",
            })

            with patch.object(amc_hook, "SESSIONS_DIR", sessions_dir), \
                 patch.object(amc_hook, "EVENTS_DIR", events_dir), \
                 patch("sys.stdin.read", return_value=hook_input), \
                 patch.dict(os.environ, clean_env, clear=True):
                amc_hook.main()

            session_file = sessions_dir / "no-spawn-sess.json"
            data = json.loads(session_file.read_text())

        self.assertNotIn("spawn_id", data)
        self.assertEqual(data["zellij_session"], "infra")
        self.assertEqual(data["zellij_pane"], "5")

    def test_user_prompt_submit_preserves_zellij_metadata(self):
        """UserPromptSubmit preserves Zellij metadata in session state."""
        existing = {
            "session_id": "test-sess",
            "status": "starting",
            "started_at": "2026-01-01T00:00:00+00:00",
        }
        data = self._run_hook("test-sess", "UserPromptSubmit", env_overrides={
            "ZELLIJ_SESSION_NAME": "infra",
            "ZELLIJ_PANE_ID": "12",
            "AMC_SPAWN_ID": "spawn-uuid-456",
        }, existing_session=existing)
        self.assertIsNotNone(data)
        self.assertEqual(data["zellij_session"], "infra")
        self.assertEqual(data["zellij_pane"], "12")
        self.assertEqual(data["spawn_id"], "spawn-uuid-456")

    def test_stop_event_preserves_zellij_metadata(self):
        """Stop event preserves Zellij metadata in session state."""
        existing = {
            "session_id": "test-sess",
            "status": "active",
            "started_at": "2026-01-01T00:00:00+00:00",
        }
        data = self._run_hook("test-sess", "Stop", env_overrides={
            "ZELLIJ_SESSION_NAME": "infra",
            "ZELLIJ_PANE_ID": "8",
        }, existing_session=existing)
        self.assertIsNotNone(data)
        self.assertEqual(data["zellij_session"], "infra")
        self.assertEqual(data["zellij_pane"], "8")


# ===========================================================================
# Spawn mixin: AMC_SPAWN_ID and pane name passed to Zellij
# ===========================================================================

class TestSpawnPassesZellijMetadata(unittest.TestCase):
    """Verify spawn mixin passes correct metadata to Zellij subprocess."""

    def _make_handler(self, project='test-proj', agent_type='claude'):
        """Create a DummySpawnHandler for testing."""
        from tests.test_spawn import DummySpawnHandler
        body = {'project': project, 'agent_type': agent_type}
        return DummySpawnHandler(body, auth_header='Bearer tok')

    def test_spawn_sets_amc_spawn_id_in_env(self):
        """_spawn_agent_in_project_tab passes AMC_SPAWN_ID in subprocess env."""
        handler = self._make_handler()
        handler._check_zellij_session_exists = MagicMock(return_value=True)
        handler._wait_for_session_file = MagicMock(return_value=True)

        with patch('subprocess.run') as mock_run:
            mock_run.return_value = MagicMock(returncode=0, stdout='', stderr='')
            result = handler._spawn_agent_in_project_tab(
                'test-proj', Path('/fake/test-proj'), 'claude', 'spawn-id-789',
            )

        self.assertTrue(result['ok'])
        # The second subprocess.run call is the pane spawn (first is tab creation)
        pane_call = mock_run.call_args_list[1]
        env_passed = pane_call.kwargs.get('env') or pane_call[1].get('env')
        self.assertIsNotNone(env_passed)
        self.assertEqual(env_passed['AMC_SPAWN_ID'], 'spawn-id-789')

    def test_spawn_creates_pane_with_correct_name(self):
        """Pane name follows '{agent_type}-{project}' convention."""
        handler = self._make_handler()
        handler._check_zellij_session_exists = MagicMock(return_value=True)
        handler._wait_for_session_file = MagicMock(return_value=True)

        with patch('subprocess.run') as mock_run:
            mock_run.return_value = MagicMock(returncode=0, stdout='', stderr='')
            handler._spawn_agent_in_project_tab(
                'my-project', Path('/fake/my-project'), 'claude', 'spawn-001',
            )

        pane_call = mock_run.call_args_list[1]
        cmd = pane_call.args[0] if pane_call.args else pane_call[0][0]
        # Verify --name argument
        name_idx = cmd.index('--name')
        self.assertEqual(cmd[name_idx + 1], 'claude-my-project')

    def test_spawn_creates_pane_with_codex_name(self):
        """Codex agent type produces correct pane name."""
        handler = self._make_handler(agent_type='codex')
        handler._check_zellij_session_exists = MagicMock(return_value=True)
        handler._wait_for_session_file = MagicMock(return_value=True)

        with patch('subprocess.run') as mock_run:
            mock_run.return_value = MagicMock(returncode=0, stdout='', stderr='')
            handler._spawn_agent_in_project_tab(
                'my-project', Path('/fake/my-project'), 'codex', 'spawn-002',
            )

        pane_call = mock_run.call_args_list[1]
        cmd = pane_call.args[0] if pane_call.args else pane_call[0][0]
        name_idx = cmd.index('--name')
        self.assertEqual(cmd[name_idx + 1], 'codex-my-project')

    def test_spawn_uses_correct_zellij_session(self):
        """Spawn commands target the configured ZELLIJ_SESSION."""
        handler = self._make_handler()
        handler._check_zellij_session_exists = MagicMock(return_value=True)
        handler._wait_for_session_file = MagicMock(return_value=True)

        with patch('subprocess.run') as mock_run, \
             patch.object(spawn_mod, 'ZELLIJ_SESSION', 'infra'):
            mock_run.return_value = MagicMock(returncode=0, stdout='', stderr='')
            handler._spawn_agent_in_project_tab(
                'proj', Path('/fake/proj'), 'claude', 'spawn-003',
            )

        # Both calls (tab + pane) should target --session infra
        for c in mock_run.call_args_list:
            cmd = c.args[0] if c.args else c[0][0]
            session_idx = cmd.index('--session')
            self.assertEqual(cmd[session_idx + 1], 'infra')

    def test_spawn_sets_cwd_to_project_path(self):
        """Spawned pane gets --cwd pointing to the project directory."""
        handler = self._make_handler()
        handler._check_zellij_session_exists = MagicMock(return_value=True)
        handler._wait_for_session_file = MagicMock(return_value=True)

        with patch('subprocess.run') as mock_run:
            mock_run.return_value = MagicMock(returncode=0, stdout='', stderr='')
            handler._spawn_agent_in_project_tab(
                'proj', Path('/home/user/projects/proj'), 'claude', 'spawn-004',
            )

        pane_call = mock_run.call_args_list[1]
        cmd = pane_call.args[0] if pane_call.args else pane_call[0][0]
        cwd_idx = cmd.index('--cwd')
        self.assertEqual(cmd[cwd_idx + 1], '/home/user/projects/proj')


# ===========================================================================
# Spawn mixin: _wait_for_session_file finds session by spawn_id
# ===========================================================================

class TestWaitForSessionFile(unittest.TestCase):
    """Verify _wait_for_session_file matches session files by spawn_id."""

    def _make_handler(self):
        from tests.test_spawn import DummySpawnHandler
        return DummySpawnHandler()

    def test_finds_session_with_matching_spawn_id(self):
        """Session file with matching spawn_id is found."""
        handler = self._make_handler()

        with tempfile.TemporaryDirectory() as tmpdir:
            sessions = Path(tmpdir)
            session_file = sessions / "abc123.json"
            session_file.write_text(json.dumps({
                "session_id": "abc123",
                "spawn_id": "target-spawn-id",
                "zellij_session": "infra",
                "zellij_pane": "5",
            }))

            with patch.object(spawn_mod, 'SESSIONS_DIR', sessions):
                found = handler._wait_for_session_file("target-spawn-id", timeout=1.0)

            self.assertTrue(found)

    def test_ignores_session_with_different_spawn_id(self):
        """Session file with different spawn_id is not matched."""
        handler = self._make_handler()

        with tempfile.TemporaryDirectory() as tmpdir:
            sessions = Path(tmpdir)
            session_file = sessions / "abc123.json"
            session_file.write_text(json.dumps({
                "session_id": "abc123",
                "spawn_id": "other-spawn-id",
            }))

            with patch.object(spawn_mod, 'SESSIONS_DIR', sessions):
                found = handler._wait_for_session_file("target-spawn-id", timeout=0.5)

            self.assertFalse(found)

    def test_ignores_session_without_spawn_id(self):
        """Session file without spawn_id is not matched."""
        handler = self._make_handler()

        with tempfile.TemporaryDirectory() as tmpdir:
            sessions = Path(tmpdir)
            session_file = sessions / "abc123.json"
            session_file.write_text(json.dumps({
                "session_id": "abc123",
                "zellij_session": "infra",
            }))

            with patch.object(spawn_mod, 'SESSIONS_DIR', sessions):
                found = handler._wait_for_session_file("target-spawn-id", timeout=0.5)

            self.assertFalse(found)

    def test_found_session_has_correct_zellij_metadata(self):
        """A found session file contains the expected Zellij metadata."""
        with tempfile.TemporaryDirectory() as tmpdir:
            sessions = Path(tmpdir)
            session_data = {
                "session_id": "sess-001",
                "spawn_id": "spawn-xyz",
                "zellij_session": "infra",
                "zellij_pane": "9",
                "status": "starting",
            }
            (sessions / "sess-001.json").write_text(json.dumps(session_data))

            # Verify the file directly (simulating what dashboard would do)
            for f in sessions.glob("*.json"):
                data = json.loads(f.read_text())
                if data.get("spawn_id") == "spawn-xyz":
                    self.assertEqual(data["zellij_session"], "infra")
                    self.assertEqual(data["zellij_pane"], "9")
                    return

            self.fail("Session file with spawn_id not found")


# ===========================================================================
# End-to-end: hook + env -> session file with full metadata
# ===========================================================================

class TestEndToEndZellijMetadata(unittest.TestCase):
    """End-to-end: simulate spawned agent hook writing full metadata."""

    def test_spawned_agent_session_file_has_all_fields(self):
        """A spawned agent's session file has session, pane, and spawn_id."""
        with tempfile.TemporaryDirectory() as tmpdir:
            sessions_dir = Path(tmpdir) / "sessions"
            events_dir = Path(tmpdir) / "events"
            sessions_dir.mkdir()
            events_dir.mkdir()

            hook_input = json.dumps({
                "hook_event_name": "SessionStart",
                "session_id": "spawned-agent-001",
                "cwd": "/home/user/projects/amc",
                "last_assistant_message": "",
            })

            # Simulate the environment a spawned agent would see:
            # AMC_SPAWN_ID set by spawn mixin, Zellij vars set by Zellij itself
            spawn_env = os.environ.copy()
            spawn_env["AMC_SPAWN_ID"] = "spawn-uuid-e2e"
            spawn_env["ZELLIJ_SESSION_NAME"] = "infra"
            spawn_env["ZELLIJ_PANE_ID"] = "14"
            spawn_env["CLAUDE_PROJECT_DIR"] = "/home/user/projects/amc"

            with patch.object(amc_hook, "SESSIONS_DIR", sessions_dir), \
                 patch.object(amc_hook, "EVENTS_DIR", events_dir), \
                 patch("sys.stdin.read", return_value=hook_input), \
                 patch.dict(os.environ, spawn_env, clear=True):
                amc_hook.main()

            session_file = sessions_dir / "spawned-agent-001.json"
            self.assertTrue(session_file.exists(), "Session file not created")
            data = json.loads(session_file.read_text())

            # All metadata present and correct
            self.assertEqual(data["zellij_session"], "infra")
            self.assertEqual(data["zellij_pane"], "14")
            self.assertEqual(data["spawn_id"], "spawn-uuid-e2e")
            self.assertEqual(data["status"], "starting")
            self.assertEqual(data["project"], "amc")

    def test_non_zellij_agent_session_file_graceful(self):
        """Agent started outside Zellij has empty metadata, no crash."""
        with tempfile.TemporaryDirectory() as tmpdir:
            sessions_dir = Path(tmpdir) / "sessions"
            events_dir = Path(tmpdir) / "events"
            sessions_dir.mkdir()
            events_dir.mkdir()

            hook_input = json.dumps({
                "hook_event_name": "SessionStart",
                "session_id": "non-zellij-agent",
                "cwd": "/home/user/projects/amc",
                "last_assistant_message": "",
            })

            # No Zellij vars, no spawn ID
            clean_env = {
                "HOME": os.environ.get("HOME", "/tmp"),
                "PATH": os.environ.get("PATH", "/usr/bin"),
                "CLAUDE_PROJECT_DIR": "/home/user/projects/amc",
            }

            with patch.object(amc_hook, "SESSIONS_DIR", sessions_dir), \
                 patch.object(amc_hook, "EVENTS_DIR", events_dir), \
                 patch("sys.stdin.read", return_value=hook_input), \
                 patch.dict(os.environ, clean_env, clear=True):
                amc_hook.main()

            session_file = sessions_dir / "non-zellij-agent.json"
            self.assertTrue(session_file.exists(), "Session file not created")
            data = json.loads(session_file.read_text())

            # Metadata present but empty (graceful degradation)
            self.assertEqual(data["zellij_session"], "")
            self.assertEqual(data["zellij_pane"], "")
            self.assertNotIn("spawn_id", data)
            # Session is still valid
            self.assertEqual(data["status"], "starting")
            self.assertEqual(data["session_id"], "non-zellij-agent")


if __name__ == "__main__":
    unittest.main()
