from aiservice.common.markdown_utils import (
    extract_all_code_from_markdown,
    extract_code_block,
    extract_code_block_with_context,
    split_markdown_code,
    wrap_code_in_markdown,
)


def test_extract_code_block_basic_python() -> None:
    text = "Some explanation\n```python\nprint('hello')\n```\nMore text"
    result = extract_code_block_with_context(text)
    assert result is not None
    before, code, after = result
    assert before == "Some explanation"
    assert code == "print('hello')\n"
    assert after == "More text"


def test_extract_code_block_with_filepath() -> None:
    text = "Explanation\n```python:src/main.py\ndef foo(): pass\n```\nAfter"
    result = extract_code_block_with_context(text)
    assert result is not None
    before, code, after = result
    assert before == "Explanation"
    assert code == "def foo(): pass\n"
    assert after == "After"


def test_extract_code_block_json_language() -> None:
    text = 'Before\n```json\n{"key": "value"}\n```\nAfter'
    result = extract_code_block_with_context(text, language="json")
    assert result is not None
    before, code, after = result
    assert before == "Before"
    assert code == '{"key": "value"}\n'
    assert after == "After"


def test_extract_code_block_no_block_returns_none() -> None:
    text = "Just some plain text without any code blocks"
    result = extract_code_block_with_context(text)
    assert result is None


def test_extract_code_block_wrong_language_returns_none() -> None:
    text = "```javascript\nconsole.log('hi')\n```"
    result = extract_code_block_with_context(text, language="python")
    assert result is None


def test_extract_code_block_escaped_newline() -> None:
    # Some LLM outputs use literal \n instead of actual newlines
    text = "Before```python\\ncode here```After"
    result = extract_code_block_with_context(text)
    assert result is not None
    before, code, after = result
    assert before == "Before"
    assert code == "code here"
    assert after == "After"


def test_extract_code_block_empty_context() -> None:
    text = "```python\ncode\n```"
    result = extract_code_block_with_context(text)
    assert result is not None
    before, code, after = result
    assert before == ""
    assert code == "code\n"
    assert after == ""


def test_extract_code_block_multiline() -> None:
    text = """Explanation here
```python
def foo():
    return 42

def bar():
    return foo() + 1
```
End of response"""
    result = extract_code_block_with_context(text)
    assert result is not None
    before, code, after = result
    assert before == "Explanation here"
    assert "def foo():" in code
    assert "def bar():" in code
    assert after == "End of response"


def test_extract_code_block_matches_first_block() -> None:
    text = "```python\nfirst\n```\nmiddle\n```python\nsecond\n```"
    result = extract_code_block_with_context(text)
    assert result is not None
    before, code, after = result
    assert before == ""
    assert code == "first\n"
    assert "second" in after


def test_extract_code_block_stops_at_first_close() -> None:
    # Non-greedy matching stops at first ```
    text = "pre\n```python\ncode ``` more\n```\npost"
    result = extract_code_block_with_context(text)
    assert result is not None
    _before, code, after = result
    assert code == "code "
    assert "more" in after


def test_extract_code_block_missing_close_returns_none() -> None:
    text = "```python\ncode without closing"
    result = extract_code_block_with_context(text)
    assert result is None


def test_extract_code_block_strips_whitespace() -> None:
    text = "  \n  Before text  \n  ```python\ncode\n```  \n  After text  \n  "
    result = extract_code_block_with_context(text)
    assert result is not None
    before, _code, after = result
    assert before == "Before text"
    assert after == "After text"


def test_extract_code_block_complete() -> None:
    text = "```python\nprint('hello')\n```"
    result = extract_code_block(text)
    assert result == "print('hello')"


def test_extract_code_block_with_surrounding_text() -> None:
    text = "Some text\n```python\ncode here\n```\nMore text"
    result = extract_code_block(text)
    assert result == "code here"


def test_extract_code_block_incomplete_fallback() -> None:
    # Missing closing ``` - should use fallback pattern
    text = "```python\ncode without closing"
    result = extract_code_block(text)
    assert result == "code without closing"


def test_extract_code_block_no_block() -> None:
    text = "Just plain text"
    result = extract_code_block(text)
    assert result is None


def test_extract_code_block_truncates_pathological_output() -> None:
    # Pathological pattern: backslash followed by 10+ quotes
    text = "```python\nx = " + "\\''''''''''''" + "\nmore\n```"
    result = extract_code_block(text)
    assert result == "x ="


def test_extract_code_block_preserves_code_before_pathological_output() -> None:
    # Valid code followed by pathological pattern - should preserve the valid code
    code = """\
def calculate_sum(numbers):
    total = 0
    for num in numbers:
        total += num
    return total

result = calculate_sum([1, 2, 3])
x = """
    text = "```python\n" + code + "\\''''''''''''" + "\nmore garbage\n```"
    result = extract_code_block(text)
    expected = """\
def calculate_sum(numbers):
    total = 0
    for num in numbers:
        total += num
    return total

result = calculate_sum([1, 2, 3])
x ="""
    assert result == expected


def test_extract_code_block_with_filepath_annotation() -> None:
    text = "```python:src/main.py\ndef foo(): pass\n```"
    result = extract_code_block(text)
    assert result == "def foo(): pass"


def test_extract_code_block_with_filepath_annotation_fallback() -> None:
    text = "```python:src/main.py\ndef foo(): pass"
    result = extract_code_block(text)
    assert result == "def foo(): pass"


def test_extract_code_block_nested_code_fence_in_triple_quote() -> None:
    # LLM embeds function definition in a triple-quoted string containing ```
    text = '```python\nimport pytest\n_source = """```python:file.py\ndef foo(): pass\n```"""\ndef test_foo():\n    assert True\n```'
    result = extract_code_block(text)
    assert (
        result
        == 'import pytest\n_source = """```python:file.py\ndef foo(): pass\n```"""\ndef test_foo():\n    assert True'
    )


def test_extract_code_block_nested_code_fence_block() -> None:
    # LLM nests a ```python:filepath block inside the main block
    text = "```python\nimport pytest\n```python:src/mod.py\ndef foo(): pass\n```\ndef test_foo():\n    assert True\n```"
    result = extract_code_block(text)
    assert result == "import pytest\n```python:src/mod.py\ndef foo(): pass\n```\ndef test_foo():\n    assert True"


def test_extract_all_code_single_block() -> None:
    text = "```python\ncode1\n```"
    result = extract_all_code_from_markdown(text)
    assert result == "code1\n"


def test_extract_all_code_multiple_blocks() -> None:
    text = "```python\ncode1\n```\ntext\n```python\ncode2\n```"
    result = extract_all_code_from_markdown(text)
    assert "code1" in result
    assert "code2" in result


def test_extract_all_code_with_filepath() -> None:
    text = "```python:file.py\ncode\n```"
    result = extract_all_code_from_markdown(text)
    assert result == "code\n"


def test_split_markdown_code_single_file() -> None:
    text = "```python:src/main.py\ndef main(): pass\n```"
    result = split_markdown_code(text)
    assert result == {"src/main.py": "def main(): pass"}


def test_split_markdown_code_multiple_files() -> None:
    text = "```python:a.py\ncode_a\n```\n```python:b.py\ncode_b\n```"
    result = split_markdown_code(text)
    assert result == {"a.py": "code_a", "b.py": "code_b"}


def test_split_markdown_code_ignores_plain_blocks() -> None:
    # Plain ```python blocks without filepath are ignored
    text = "```python\nignored\n```\n```python:file.py\nincluded\n```"
    result = split_markdown_code(text)
    assert result == {"file.py": "included"}


def test_split_markdown_code_strips_filepath() -> None:
    text = "```python:  spaced.py  \ncode\n```"
    result = split_markdown_code(text)
    assert "spaced.py" in result


def test_wrap_code_in_markdown_python() -> None:
    code = "print('hello')"
    result = wrap_code_in_markdown(code)
    assert result == "```python\nprint('hello')\n```"


def test_wrap_code_in_markdown_json() -> None:
    code = '{"key": "value"}'
    result = wrap_code_in_markdown(code, language="json")
    assert result == '```json\n{"key": "value"}\n```'


def test_wrap_code_in_markdown_multiline() -> None:
    code = "def foo():\n    return 42"
    result = wrap_code_in_markdown(code)
    assert result == "```python\ndef foo():\n    return 42\n```"


def test_wrap_and_extract_roundtrip() -> None:
    original_code = "print('hello')"
    wrapped = wrap_code_in_markdown(original_code)
    extracted = extract_code_block(wrapped)
    assert extracted == original_code


def test_extract_code_block_with_context_prefers_filepath_block() -> None:
    """Test that filepath-annotated blocks are preferred over plain blocks.

    This handles LLM responses that include explanation code blocks before
    the actual optimized code with a filepath annotation. This test is based on
    an actual failing JIT rewrite response (trace 59be9279-dd59-4fad-aaf4-6bc541e44cb8)
    where the LLM returned an explanation with example code before the real optimized code.
    """
    # Actual response pattern from gpt-5-mini JIT rewrite that caused extraction failure
    text = """**Optimization Explanation:**
I applied targeted micro-optimizations that reduce attribute lookups.

```python
# Example of the original slow pattern:
def slow_convert(sequence):
    result = []
    for item in sequence:
        result.append(tuple(item))
    return tuple(result)
```

The optimized code uses list comprehension for better performance:

```python:unstructured/documents/elements.py
from __future__ import annotations
import dataclasses as dc

@dc.dataclass
class DataSourceMetadata:
    url: str | None = None
```

These yield measurable runtime improvements."""
    result = extract_code_block_with_context(text)
    assert result is not None
    before, code, after = result
    # Should extract the filepath-annotated block, not the example block
    assert (
        before
        == """**Optimization Explanation:**
I applied targeted micro-optimizations that reduce attribute lookups.

```python
# Example of the original slow pattern:
def slow_convert(sequence):
    result = []
    for item in sequence:
        result.append(tuple(item))
    return tuple(result)
```

The optimized code uses list comprehension for better performance:"""
    )
    assert (
        code
        == """from __future__ import annotations
import dataclasses as dc

@dc.dataclass
class DataSourceMetadata:
    url: str | None = None
"""
    )
    assert after == "These yield measurable runtime improvements."


def test_extract_code_block_with_context_falls_back_to_plain_block() -> None:
    """Test fallback to plain block when no filepath-annotated block exists."""
    text = """Here's the code:

```python
def simple_function():
    return 42
```
"""
    result = extract_code_block_with_context(text)
    assert result is not None
    before, code, after = result
    assert before == "Here's the code:"
    assert code == "def simple_function():\n    return 42\n"
    assert after == ""
