#!/usr/bin/env python3

"""Tests for the WriteFile subtool."""

import os
import unittest

from codemcp.testing import MCPEndToEndTestCase


class WriteFileTest(MCPEndToEndTestCase):
    """Test the WriteFile subtool."""

    async def test_write_file(self):
        """Test the WriteFile subtool, which writes to a file and automatically commits the changes."""
        test_file_path = os.path.join(self.temp_dir.name, "new_file.txt")
        content = "New content\nLine 2"

        # First add the file to git to make it tracked
        with open(test_file_path, "w") as f:
            f.write("")

        # Add it to git
        await self.git_run(["add", test_file_path])

        # Commit it
        await self.git_run(["commit", "-m", "Add empty file for WriteFile test"])

        async with self.create_client_session() as session:
            # First initialize project to get chat_id
            init_result_text = await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "InitProject",
                    "path": self.temp_dir.name,
                    "user_prompt": "Test initialization for write_file test",
                    "subject_line": "test: initialize for write file test",
                    "reuse_head_chat_id": False,
                },
            )

            # Extract chat_id from the init result
            chat_id = self.extract_chat_id_from_text(init_result_text)

            # Call the WriteFile tool with chat_id using our new helper method
            result_text = await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "WriteFile",
                    "path": test_file_path,
                    "content": content,
                    "description": "Create new file",
                    "chat_id": chat_id,
                },
            )

            # Verify the success message
            self.assertIn("Successfully wrote to", result_text)

            # Verify the file was created with the correct content
            with open(test_file_path) as f:
                file_content = f.read()
            self.assertEqual(file_content, content + "\n")

            # Verify git state (working tree should be clean after automatic commit)
            status = await self.git_run(["status"], capture_output=True, text=True)

            # Use expect test to verify git status - should show clean working tree
            self.assertExpectedInline(
                status,
                """\
On branch main
nothing to commit, working tree clean""",
            )

            # Get the commit message of the HEAD commit
            commit_message = await self.git_run(
                ["log", "-1", "--pretty=%B"], capture_output=True, text=True
            )

            # Normalize the chat_id in the commit message for expect test
            normalized_commit_message = commit_message.replace(chat_id, "test-chat-id")

            # Use expect test to verify the commit message format
            self.assertExpectedInline(
                normalized_commit_message,
                """\
test: initialize for write file test

Test initialization for write_file test

```git-revs
c9bcf9c  (Base revision)
HEAD     Create new file
```

codemcp-id: test-chat-id""",
            )

            # Second write to the same file
            updated_content = content + "\nAdded third line"

            # Call the WriteFile tool again with updated content
            result_text = await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "WriteFile",
                    "path": test_file_path,
                    "content": updated_content,
                    "description": "Update file with third line",
                    "chat_id": chat_id,
                },
            )

            # Verify the success message for second write
            self.assertIn("Successfully wrote to", result_text)

            # Verify the file was updated with the correct content
            with open(test_file_path) as f:
                file_content = f.read()
            self.assertEqual(file_content, updated_content + "\n")

            # Verify git state after second write
            status = await self.git_run(["status"], capture_output=True, text=True)

            # Use expect test to verify git status - should still show clean working tree
            self.assertExpectedInline(
                status,
                """\
On branch main
nothing to commit, working tree clean""",
            )

            # Get the commit message of the HEAD commit after second write
            commit_message = await self.git_run(
                ["log", "-1", "--pretty=%B"], capture_output=True, text=True
            )

            # Normalize the chat_id in the commit message for expect test
            normalized_commit_message = commit_message.replace(chat_id, "test-chat-id")

            # Use expect test to verify the commit message format for second write
            self.assertExpectedInline(
                normalized_commit_message,
                """\
test: initialize for write file test

Test initialization for write_file test

```git-revs
c9bcf9c  (Base revision)
49bf8ff  Create new file
HEAD     Update file with third line
```

codemcp-id: test-chat-id""",
            )

    async def test_create_new_file_with_write_file(self):
        """Test creating a new file that doesn't exist yet with WriteFile."""
        # Path to a new file that doesn't exist yet, within the git repository
        new_file_path = os.path.join(self.temp_dir.name, "completely_new_file.txt")

        self.assertFalse(
            os.path.exists(new_file_path),
            "Test file should not exist initially",
        )

        async with self.create_client_session() as session:
            # First initialize project to get chat_id
            init_result_text = await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "InitProject",
                    "path": self.temp_dir.name,
                    "user_prompt": "Test initialization for creating new file",
                    "subject_line": "test: initialize for new file test",
                    "reuse_head_chat_id": False,
                },
            )

            # Extract chat_id from the init result
            chat_id = self.extract_chat_id_from_text(init_result_text)

            # Create a new file using our helper method
            result_text = await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "WriteFile",
                    "path": new_file_path,
                    "content": "This is a brand new file",
                    "description": "Create a new file with WriteFile",
                    "chat_id": chat_id,
                },
            )

            # Check that the operation succeeded
            self.assertIn("Successfully wrote to", result_text)

            # Verify the file was created
            self.assertTrue(
                os.path.exists(new_file_path),
                "File was not created even though operation reported success",
            )

            # Check content
            with open(new_file_path) as f:
                content = f.read()
            self.assertEqual(content, "This is a brand new file\n")

            # Verify the file was added to git
            ls_files_output = await self.git_run(
                ["ls-files", new_file_path], capture_output=True, text=True
            )

            # The new file should be tracked in git
            self.assertTrue(
                ls_files_output,
                "New file was created but not added to git",
            )

    async def test_write_to_untracked_file(self):
        """Test that writes to untracked files are rejected."""
        # Create an untracked file (not added to git)
        untracked_file_path = os.path.join(
            self.temp_dir.name,
            "untracked_for_write.txt",
        )
        with open(untracked_file_path, "w") as f:
            f.write("Initial content in untracked file")

        # Verify the file exists but is not tracked by git
        file_exists = os.path.exists(untracked_file_path)
        self.assertTrue(file_exists, "Test file should exist on filesystem")

        # Check that the file is untracked
        ls_files_output = await self.git_run(
            ["ls-files", untracked_file_path], capture_output=True, text=True
        )

        self.assertEqual(ls_files_output, "", "File should not be tracked by git")

        # Save original content and modification time for comparison
        with open(untracked_file_path) as f:
            original_content = f.read()
        original_mtime = os.path.getmtime(untracked_file_path)

        async with self.create_client_session() as session:
            # First initialize project to get chat_id
            init_result_text = await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "InitProject",
                    "path": self.temp_dir.name,
                    "user_prompt": "Test initialization for untracked file test",
                    "subject_line": "test: initialize for untracked file test",
                    "reuse_head_chat_id": False,
                },
            )

            # Extract chat_id from the init result
            chat_id = self.extract_chat_id_from_text(init_result_text)

            # Try to write to the untracked file
            new_content = "This content should not be written to untracked file"
            result_text = await self.call_tool_assert_error(
                session,
                "codemcp",
                {
                    "subtool": "WriteFile",
                    "path": untracked_file_path,
                    "content": new_content,
                    "description": "Attempt to write to untracked file",
                    "chat_id": chat_id,
                },
            )

            # The error message format has changed since we're now directly calling the subtool
            # instead of going through the codemcp tool
            self.assertIn(
                "File is not tracked by git",
                result_text,
                "Error message should indicate the file is not tracked by git",
            )

            # Verify the file content has not changed
            with open(untracked_file_path) as f:
                actual_content = f.read()

            self.assertEqual(
                original_content,
                actual_content,
                "File content should not change when operation is rejected",
            )

            # Verify file modification time was not changed
            current_mtime = os.path.getmtime(untracked_file_path)
            self.assertEqual(
                current_mtime,
                original_mtime,
                "File modification time should not have changed",
            )

    async def test_write_file_outside_tracked_paths(self):
        """Test that codemcp properly handles writing to paths outside tracked paths."""
        # Create a subdirectory but don't add it to git
        subdir_path = os.path.join(self.temp_dir.name, "untrackeddir")
        os.makedirs(subdir_path, exist_ok=True)

        new_file_path = os.path.join(subdir_path, "newfile.txt")

        async with self.create_client_session() as session:
            # First initialize project to get chat_id
            init_result_text = await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "InitProject",
                    "path": self.temp_dir.name,
                    "user_prompt": "Test initialization for untracked directory test",
                    "subject_line": "test: initialize for untracked directory test",
                    "reuse_head_chat_id": False,
                },
            )

            # Extract chat_id from the init result
            chat_id = self.extract_chat_id_from_text(init_result_text)

            # Try to write a new file in the untracked directory
            await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "WriteFile",
                    "path": new_file_path,
                    "content": "New file in untracked directory",
                    "description": "Attempt to create file in untracked directory",
                    "chat_id": chat_id,
                },
            )

            # Since we're using call_tool_assert_success, we know the operation succeeded
            # Check if the directory and file are now tracked in git
            await self.git_run(["status"], capture_output=True, text=True)

            # The file should exist
            self.assertTrue(
                os.path.exists(new_file_path),
                "File was not created even though operation reported success",
            )

            # SECURITY CHECK: If writing to untracked directories succeeds,
            # both the directory and file should be tracked in git
            ls_files_output = await self.git_run(
                ["ls-files", new_file_path], capture_output=True, text=True
            )

            # IMPORTANT: The file should be tracked in git after writing
            self.assertTrue(
                ls_files_output,
                "SECURITY VULNERABILITY: Successfully wrote to untracked directory"
                " but did not add file to git",
            )

    async def test_write_file_with_tilde_path(self):
        """Test that WriteFile correctly handles paths with tilde (~) expansion."""
        # Create a temporary directory in the user's home directory specifically for this test
        home_dir = os.path.expanduser("~")
        home_temp_dir_path = os.path.join(home_dir, f"codemcp_test_tilde_{os.getpid()}")
        os.makedirs(home_temp_dir_path, exist_ok=True)

        try:
            # Initialize a git repository in this directory
            await self.git_run(["init", "-b", "main"], cwd=home_temp_dir_path)
            await self.git_run(
                ["config", "user.email", "test@example.com"], cwd=home_temp_dir_path
            )
            await self.git_run(
                ["config", "user.name", "Test User"], cwd=home_temp_dir_path
            )

            # Create initial commit
            readme_path = os.path.join(home_temp_dir_path, "README.md")
            with open(readme_path, "w") as f:
                f.write("# Test Repository for tilde expansion\n")

            # Create a codemcp.toml file (required for permission checks)
            codemcp_toml_path = os.path.join(home_temp_dir_path, "codemcp.toml")
            with open(codemcp_toml_path, "w") as f:
                f.write("")

            await self.git_run(
                ["add", "README.md", "codemcp.toml"], cwd=home_temp_dir_path
            )
            await self.git_run(
                ["commit", "-m", "Initial commit"], cwd=home_temp_dir_path
            )

            # Create a relative path using tilde
            rel_path = os.path.relpath(home_temp_dir_path, home_dir)
            tilde_path = os.path.join("~", rel_path, "tilde_test_file.txt")

            # Expected absolute path after tilde expansion
            expected_abs_path = os.path.expanduser(tilde_path)
            self.assertEqual(
                os.path.abspath(expected_abs_path),
                os.path.abspath(
                    os.path.join(home_temp_dir_path, "tilde_test_file.txt")
                ),
            )

            # Content to write
            content = "This file was created with a tilde path"

            async with self.create_client_session() as session:
                # First initialize project to get chat_id
                init_result_text = await self.call_tool_assert_success(
                    session,
                    "codemcp",
                    {
                        "subtool": "InitProject",
                        "path": home_temp_dir_path,
                        "user_prompt": "Test initialization for tilde path test",
                        "subject_line": "test: initialize for tilde path test",
                        "reuse_head_chat_id": False,
                    },
                )

                # Extract chat_id from the init result
                chat_id = self.extract_chat_id_from_text(init_result_text)

                # Call the WriteFile tool with the tilde path
                await self.call_tool_assert_success(
                    session,
                    "codemcp",
                    {
                        "subtool": "WriteFile",
                        "path": tilde_path,
                        "content": content,
                        "description": "Create file using tilde path",
                        "chat_id": chat_id,
                    },
                )

                # Verify the file was created
                self.assertTrue(
                    os.path.exists(expected_abs_path),
                    "File should exist after writing with tilde path",
                )

                # Verify content
                with open(expected_abs_path) as f:
                    actual_content = f.read()
                # WriteFile adds a newline at the end of files
                self.assertEqual(actual_content, content + "\n")

                # Verify the file was added to git
                ls_files_output = await self.git_run(
                    ["ls-files", "tilde_test_file.txt"],
                    capture_output=True,
                    text=True,
                    cwd=home_temp_dir_path,
                )

                # The file should be tracked in git
                self.assertTrue(
                    ls_files_output,
                    "New file created with tilde path should be tracked in git",
                )
        finally:
            # Clean up the temporary directory
            import shutil

            if os.path.exists(home_temp_dir_path):
                shutil.rmtree(home_temp_dir_path)

    async def test_user_prompt_with_markdown_code_block(self):
        """Test handling of user prompt that contains a Markdown code block with triple backticks."""
        test_file_path = os.path.join(
            self.temp_dir.name, "markdown_code_block_test.txt"
        )
        content = "File created from a prompt with a code block"

        # Create placeholder file and add to git
        with open(test_file_path, "w") as f:
            f.write("")

        # Add it to git
        await self.git_run(["add", test_file_path])

        # Commit it
        await self.git_run(
            ["commit", "-m", "Add empty file for markdown code block test"]
        )

        # User prompt with Markdown code block
        user_prompt_with_code_block = """Please create a file with this Python code:

```
---
description: Description of when the rule is useful for the LLM
globs: *.js,*.ts
alwaysApply: false
---
Markdown to send to LLM
```

And make sure it runs correctly."""

        async with self.create_client_session() as session:
            # Initialize project to get chat_id
            init_result_text = await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "InitProject",
                    "path": self.temp_dir.name,
                    "user_prompt": user_prompt_with_code_block,
                    "subject_line": "test: user prompt with markdown code block",
                    "reuse_head_chat_id": False,
                },
            )

            # Extract chat_id from the init result
            chat_id = self.extract_chat_id_from_text(init_result_text)

            # Call the WriteFile tool with chat_id
            result_text = await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "WriteFile",
                    "path": test_file_path,
                    "content": content,
                    "description": "Write file from prompt with code block",
                    "chat_id": chat_id,
                },
            )

            # Verify the success message
            self.assertIn("Successfully wrote to", result_text)

            # Verify the file was created with the correct content
            with open(test_file_path) as f:
                file_content = f.read()
            self.assertEqual(file_content, content + "\n")

            # Get the commit message of the HEAD commit
            commit_message = await self.git_run(
                ["log", "-1", "--pretty=%B"], capture_output=True, text=True
            )

            # Normalize the chat_id in the commit message for expect test
            normalized_commit_message = commit_message.replace(chat_id, "test-chat-id")

            # Verify that the commit message contains the code block with triple backticks
            self.assertExpectedInline(
                normalized_commit_message,
                """\
test: user prompt with markdown code block

Please create a file with this Python code:

```
---
description: Description of when the rule is useful for the LLM
globs: *.js,*.ts
alwaysApply: false
---
Markdown to send to LLM
```

And make sure it runs correctly.

```git-revs
6350984  (Base revision)
HEAD     Write file from prompt with code block
```

codemcp-id: test-chat-id""",
            )

            # Now do a second write operation with the same chat_id
            updated_content = content + "\nSecond write with code block in user_prompt"

            # Call the WriteFile tool again with updated content
            result_text = await self.call_tool_assert_success(
                session,
                "codemcp",
                {
                    "subtool": "WriteFile",
                    "path": test_file_path,
                    "content": updated_content,
                    "description": "Update file with second write",
                    "chat_id": chat_id,
                },
            )

            # Verify the success message
            self.assertIn("Successfully wrote to", result_text)

            # Verify the file was updated with the correct content
            with open(test_file_path) as f:
                file_content = f.read()
            self.assertEqual(file_content, updated_content + "\n")

            # Get the commit message after second write
            commit_message = await self.git_run(
                ["log", "-1", "--pretty=%B"], capture_output=True, text=True
            )

            # Normalize the chat_id for expect test
            normalized_commit_message = commit_message.replace(chat_id, "test-chat-id")

            # Verify that the commit message still contains the code block with triple backticks
            self.assertExpectedInline(
                normalized_commit_message,
                """\
test: user prompt with markdown code block

Please create a file with this Python code:

```
---
description: Description of when the rule is useful for the LLM
globs: *.js,*.ts
alwaysApply: false
---
Markdown to send to LLM
```

And make sure it runs correctly.

```git-revs
6350984  (Base revision)
52d0290  Write file from prompt with code block
HEAD     Update file with second write
```

codemcp-id: test-chat-id""",
            )


class OutOfProcessWriteFileTest(WriteFileTest):
    in_process = False


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