import modulus.container
import modulus.build_context
import modulus.clangd
import ast
import contextlib
import io
import json
import os
import pathlib
import subprocess
import tempfile
import sys
import unittest
from unittest.mock import patch

import modulus.cli
import modulus.config
import modulus.export
import modulus.format_layout
import modulus.formatting

class formatting_test(unittest.TestCase):
  def formatted(self, source, suffix):
    result = modulus.format_layout.format_source(source, suffix)
    self.assertEqual(modulus.format_layout.format_source(result, suffix), result)
    return result
  def test_cpp_units_namespaces_and_include_groups(self):
    source = '#include <vector>\n#include <string>\n\n#include "local.hpp"\n\nnamespace sample {\nint one() {\n\nint x=1;\n\nx++;\nreturn x;\n}\nint two() {return 2;}\n}\n'
    result = self.formatted(source, ".cpp")
    self.assertTrue(result.startswith(source.split("namespace")[0]))
    self.assertIn("namespace sample {\n\nint one()", result)
    self.assertIn("int x = 1;\n  x++;\n  return x;", result)
    self.assertIn("}\n\nint two()", result)
    self.assertTrue(result.endswith("}\n\n}\n"))
  def test_cpp_no_column_limit_and_simple_lists(self):
    name = "argument_" + "a" * 180
    result = self.formatted(f"void function(\nint {name},\nint b) {{\n  call(\n{name},\nb);\n}}\n", ".cpp")
    self.assertIn(f"void function(int {name}, int b)", result)
    self.assertIn(f"call({name}, b);", result)
  def test_cpp_class_and_struct_members_have_no_blank_lines(self):
    for kind in ("class", "struct"):
      with self.subTest(kind=kind):
        source = f'{kind} source_location {{\npublic:\n\n  source_location() = default;\n\n  source_location(source_location const&) = delete;\n\n  int line() const;\n\n  int value() const {{\n\n    return line_;\n\n  }}\n\n  struct nested {{\n\n    int value;\n\n  }};\n\nprivate:\n\n  int line_;\n\n}};\n\nstruct next {{ int value; }};\n'
        result = self.formatted(source, ".hpp")
        body = result.split(f"{kind} source_location {{\n", 1)[1].split("\n};", 1)[0]
        self.assertNotIn("\n\n", body)
        self.assertIn("};\n\nstruct next", result)
  def test_cpp_class_keeps_literal_and_comment_blank_lines(self):
    literal = 'R"raw(first\n\nsecond)raw"'
    comment = '/* first\n\nsecond */'
    source = f'struct example {{\n\n  const char* text = {literal};\n\n  {comment}\n\n  int value;\n}};\n'
    result = self.formatted(source, ".hpp")
    self.assertIn(literal, result)
    self.assertIn(comment, result)
    self.assertNotIn("\n\n", result.replace(literal, "literal").replace(comment, "comment"))
  def test_cpp_literals_macros_templates_and_lambdas(self):
    literal = 'R"raw(first\n\n  second\n)raw"'
    source = f'template<class T>\nvoid test(T v) {{\n  const char* s = {literal};\n\n  visit(v, [](auto a) {{\n\nreturn a;\n}});\n}}\nBOOST_AUTO_TEST_CASE(example) {{\n\nCHECK(true);\n\n}}\n'
    result = self.formatted(source, ".cpp")
    self.assertIn(literal, result)
    self.assertIn("BOOST_AUTO_TEST_CASE(example) {\n  CHECK(true);\n}", result)
    self.assertIn("[](auto a) {\n", result)
  def test_python_ast_literals_decorators_and_nested_units(self):
    source = 'import os\n\nimport pathlib\n\n@decorator\ndef outer(\n    first,\n    second,\n):\n    text = """first\n\n    second\n"""\n\n    first += 1\n    def nested():\n        return 1\n    return call(\n        first,\n        second,\n    )\n'
    result = self.formatted(source, ".py")
    self.assertEqual(ast.dump(ast.parse(source)), ast.dump(ast.parse(result)))
    self.assertIn('"""first\n\n    second\n"""', result)
    self.assertIn("@decorator\ndef outer(first, second):", result)
    self.assertIn("first += 1\n\n  def nested()", result)
    self.assertIn("return 1\n\n  return call(first, second)", result)
    self.assertTrue(result.startswith("import os\n\nimport pathlib\n"))
  def test_python_no_column_limit_and_structured_literals(self):
    name = "a" * 250
    source = f"def test({name},\n other):\n return call({name},\nother)\n"
    result = self.formatted(source, ".py")
    self.assertIn(f"def test({name}, other):", result)
    data = 'def test():\n  return call({\n    "a": 1,\n    "b": 2,\n  })\n'
    self.assertIn('{\n', self.formatted(data, ".py"))
  def test_python_classes_keep_members_together(self):
    source = 'class Example:\n\n  """first\n\nsecond"""\n\n  field = 1\n\n  @staticmethod\n  def value():\n\n    return 1\n\n  class Nested:\n\n    field = 2\n\n  async def other(self):\n\n    return 2\n\nclass Next:\n\n  field = 3\n'
    result = self.formatted(source, ".py")
    self.assertEqual(ast.dump(ast.parse(source)), ast.dump(ast.parse(result)))
    self.assertIn('"""first\n\nsecond"""', result)
    self.assertNotIn("\n\n", result.split("\n\nclass Next", 1)[0].replace('"""first\n\nsecond"""', 'docstring'))
    self.assertIn("return 2\n\nclass Next:\n  field = 3", result)
  def test_python_tuple_continuations_use_two_spaces(self):
    source = 'runtime_ignore_patterns = (\n    ".git",\n    ".tmp",\n    "__pycache__",\n)\n'
    expected = 'runtime_ignore_patterns = (\n  ".git",\n  ".tmp",\n  "__pycache__",\n)\n'
    self.assertEqual(self.formatted(source, ".py"), expected)
  def test_python_nested_collections_use_two_spaces_per_opening_line(self):
    source = 'def example():\n    result = build({\n        "items": [\n            (\n                1,\n                2,\n            ),\n        ],\n        "names": {\n            "a",\n            "b",\n        },\n    })\n    return result\n'
    expected = 'def example():\n  result = build({\n    "items": [\n      (\n        1,\n        2,\n      ),\n    ],\n    "names": {\n      "a",\n      "b",\n    },\n  })\n  return result\n'
    self.assertEqual(self.formatted(source, ".py"), expected)
    self.assertEqual(self.formatted("values = [[\n    1,\n    2,\n]]\n", ".py"), "values = [[\n  1,\n  2,\n]]\n")
  def test_python_continued_imports_headers_and_comments(self):
    source = 'from somewhere import (\n    first,\n    second,\n)\n\n@decorator(\n    first,  # comment\n    second,\n)\ndef example(\n    first,  # comment\n    second,\n):\n    if (\n        first  # condition\n        and second\n    ):\n        return first\n'
    expected = 'from somewhere import (\n  first,\n  second,\n)\n\n@decorator(\n  first,  # comment\n  second,\n)\ndef example(\n  first,  # comment\n  second,\n):\n  if (\n    first  # condition\n    and second\n  ):\n    return first\n'
    self.assertEqual(self.formatted(source, ".py"), expected)
  def test_python_explicit_continuations_use_two_spaces(self):
    source = 'def example():\n    with first() as one, \\\n            second() as two:\n        return one + \\\n            two\n'
    expected = 'def example():\n  with first() as one, \\\n    second() as two:\n    return one + \\\n      two\n'
    self.assertEqual(self.formatted(source, ".py"), expected)
  def test_python_continuations_preserve_multiline_literal_contents(self):
    source = 'def example():\n    return [\n        # first value\n        f"{call({1: 2})}",\n        """first\n    second\n\n        third""",\n        f"""start {name}\n    {value}\n\n        end""",\n        "next",\n    ]\n'
    expected = 'def example():\n  return [\n    # first value\n    f"{call({1: 2})}",\n    """first\n    second\n\n        third""",\n    f"""start {name}\n    {value}\n\n        end""",\n    "next",\n  ]\n'
    self.assertEqual(self.formatted(source, ".py"), expected)
  def test_python_adjacent_literals_are_continuations_not_literal_contents(self):
    source = 'def example():\n    text = (\n        "first"\n        "second"\n        f"{value}"\n        f"{other}"\n    )\n    return text\n'
    expected = 'def example():\n  text = (\n    "first"\n    "second"\n    f"{value}"\n    f"{other}"\n  )\n  return text\n'
    self.assertEqual(self.formatted(source, ".py"), expected)
  def test_aliases_and_portability_exception(self):
    cpp = "namespace std { using bool_t = decltype(true); }\nusing integer = int;\nnamespace fs = std::filesystem;\ntypedef int count;\n"
    self.assertEqual(len(list(modulus.format_layout.alias_diagnostics(cpp, ".hpp"))), 4)
    self.assertEqual(len(list(modulus.format_layout.alias_diagnostics(cpp, ".hpp", True))), 3)
    python = "import pathlib as p\nfrom typing import Literal\nRole = Literal['a', 'b']\n"
    self.assertEqual(len(list(modulus.format_layout.alias_diagnostics(python, ".py"))), 2)
  def test_stdin_failure_has_no_output_and_never_builds_image(self):
    with tempfile.TemporaryDirectory() as directory:
      root = pathlib.Path(directory)
      context = modulus.build_context.build_context(root, modulus.config.normalize({"package": {"name": "sample"}}))
      context.docker_marker = root / "not-in-container"
      output, errors = io.StringIO(), io.StringIO()
      with patch.dict(os.environ, {"MODULUS_DIRECT": "1"}), patch("sys.stdin", io.StringIO("def broken(")), contextlib.redirect_stdout(output), contextlib.redirect_stderr(errors):
        self.assertEqual(modulus.formatting.launch(context, ["--stdin-filename", "test.py"]), 1)
      self.assertEqual(output.getvalue(), "")
      with patch.dict(os.environ, {"MODULUS_DIRECT": "0"}), patch.object(modulus.container, "current_image", return_value=None), patch.object(modulus.container, "ensure_build_image") as build, contextlib.redirect_stderr(errors):
        self.assertEqual(modulus.formatting.launch(context, ["--check"]), 1)
      build.assert_not_called()
  def test_discovery_excludes_generated_and_symlinks(self):
    with tempfile.TemporaryDirectory() as directory:
      root = pathlib.Path(directory)
      subprocess.run(["git", "init", "-q", str(root)], check=True)
      for name in ["keep.cpp", "skip.cpp", "build/generated.cpp", "vendor/library.cpp"]:
        path = root / name
        path.parent.mkdir(parents=True, exist_ok=True)
        path.touch()
      (root / "link.cpp").symlink_to(root / "keep.cpp")
      self.assertEqual(modulus.formatting.source_files(root, [], ["skip.cpp"]), [root / "keep.cpp"])
  def test_nested_lists_join_in_one_pass_without_changing_string_spaces(self):
    cpp = 'void example() { outer(\n inner(\n 1,\n 2),\n "a  b"); }\n'
    self.assertIn('outer(inner(1, 2), "a  b");', self.formatted(cpp, ".cpp"))
    python = 'def example():\n  return outer(\n    inner(\n      1,\n      2),\n    "a  b")\n'
    self.assertIn('outer(inner(1, 2), "a  b")', self.formatted(python, ".py"))
  def test_multiline_fstrings_preserve_blank_lines_on_new_python_tokenizers(self):
    source = 'def example(name):\n  return f"""first {name}\n\nsecond\n\nthird"""\n'
    self.assertEqual(self.formatted(source, ".py"), source)
  def test_alias_check_distinguishes_module_aliases_from_saved_callbacks(self):
    source = 'import modulus.cli\nsaved = modulus.container.docker_run\nconfig = modulus.cli\nItems = list[int]\n'
    self.assertEqual([line for line, _ in modulus.format_layout.alias_diagnostics(source, ".py")], [3, 4])
  def test_exported_stdin_formatter_needs_no_compile_database(self):
    with tempfile.TemporaryDirectory() as directory:
      root = pathlib.Path(directory)
      (root / "modulus.json").write_text(json.dumps({"package": {"name": "sample"}}))
      modulus.export.sync_project(root, pathlib.Path(__file__).resolve().parent.parent)
      result = subprocess.run([sys.executable, str(root / "build/build_tools/format.py"), "--stdin-filename", "sample.py"], cwd=root, env=dict(os.environ, MODULUS_DIRECT="1", PYTHONDONTWRITEBYTECODE="1"), input="def sample():\n    values = (\n        1,\n        2,\n    )\n    return values\n", text=True, capture_output=True)
      self.assertEqual(result.returncode, 0, result.stderr)
      self.assertEqual(result.stdout, "def sample():\n  values = (\n    1,\n    2,\n  )\n  return values\n")
      self.assertEqual(result.stderr, "")
  @unittest.skipUnless(os.environ.get("MODULUS_FORMAT_REPO") and os.environ.get("MODULUS_CONFORM_PATH"), "set MODULUS_FORMAT_REPO and MODULUS_CONFORM_PATH for the Neovim integration test")
  def test_neovim_container_formatting(self):
    root = pathlib.Path(__file__).resolve().parent.parent
    env = dict(os.environ, MODULUS_FORMAT_HELPER=str(root / "tools/nvim/modulus.lua"))
    result = subprocess.run(["nvim", "--headless", "-u", "NONE", "-l", str(root / "tests/nvim_format.lua")], env=env, text=True, capture_output=True, timeout=60)
    self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
    self.assertIn("MODULUS_FORMAT_OK", result.stdout + result.stderr)

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