from __future__ import annotations

import re

MARKDOWN_CODE_BLOCK_PATTERN = re.compile(
    r"```python(?::[^\n]*)?\n(.*?)```", re.DOTALL
)

FIRST_CODE_BLOCK_PATTERN = re.compile(
    r"^```python(?::[^\n]*)?\s*\n(.*)\n```[ \t]*$",
    re.MULTILINE | re.DOTALL,
)

FIRST_CODE_BLOCK_FALLBACK_PATTERN = re.compile(
    r"^```python(?::[^\n]*)?\s*\n(.*)",
    re.MULTILINE | re.DOTALL,
)

_PATHOLOGICAL_PATTERN = re.compile(r"(\\['\"]{10,})")


def truncate_pathological_output(code: str) -> str:
    """
    Truncate code at pathological repeated escape sequences.
    """
    if match := _PATHOLOGICAL_PATTERN.search(code):
        return code[: match.start()].rstrip()
    return code


def extract_all_code_from_markdown(markdown: str) -> str:
    """
    Extract all code from markdown code blocks.
    """
    matches = MARKDOWN_CODE_BLOCK_PATTERN.findall(markdown)
    return "\n\n".join(matches)


def extract_code_block(markdown: str) -> str | None:
    """
    Extract the first code block from markdown.
    """
    if match := FIRST_CODE_BLOCK_PATTERN.search(markdown):
        return truncate_pathological_output(match.group(1))
    if match := FIRST_CODE_BLOCK_FALLBACK_PATTERN.search(markdown):
        return truncate_pathological_output(match.group(1))
    return None


def split_markdown_code(
    markdown: str, language: str = "python"
) -> dict[str, str]:
    """
    Split markdown into a dict of filepath to code content.
    """
    if language in ("javascript", "js"):
        lang_pattern = r"(?:javascript|js)"
    elif language in ("typescript", "ts"):
        lang_pattern = r"(?:typescript|ts)"
    else:
        lang_pattern = re.escape(language)

    pattern = re.compile(rf"```{lang_pattern}:([^\n]+)\n(.*?)\n```", re.DOTALL)
    matches = pattern.findall(markdown)
    result: dict[str, str] = {}
    for file_path, code in matches:
        stripped_path = file_path.strip()
        if stripped_path not in result:
            result[stripped_path] = code
    return result


def extract_code_block_with_context(
    text: str, language: str = "python"
) -> tuple[str, str, str] | None:
    """
    Extract a code block and its surrounding context.
    """
    pattern_with_path = rf"(.*?)```{language}:[^\n]+(?:\n|\\n)(.*?)```(.*)"
    if match := re.match(pattern_with_path, text, re.DOTALL | re.MULTILINE):
        return (
            match.group(1).strip(),
            match.group(2),
            match.group(3).strip(),
        )

    pattern = rf"(.*?)```{language}(?::[^\n]*)?(?:\n|\\n)(.*?)```(.*)"
    if match := re.match(pattern, text, re.DOTALL | re.MULTILINE):
        return (
            match.group(1).strip(),
            match.group(2),
            match.group(3).strip(),
        )
    return None


def wrap_code_in_markdown(code: str, language: str = "python") -> str:
    """
    Wrap code in a markdown code block.
    """
    return f"```{language}\n{code}\n```"


def group_code(file_to_code: dict[str, str], language: str = "python") -> str:
    """
    Join file-to-code mapping into grouped markdown blocks.
    """

    def format_block(file_path: str, code: str) -> str:
        normalized = code if code.endswith("\n") else code + "\n"
        return f"```{language}:{file_path}\n{normalized}```"

    return "\n".join(
        format_block(path, code) for path, code in file_to_code.items()
    )


def is_multi_context(code: str) -> bool:
    """
    Check if code is in multi-file markdown format.
    """
    return code.strip().startswith("```python:")
