"""Utilities for extracting code from markdown code blocks.

This module provides patterns and functions for working with Python code blocks
in markdown format, supporting both plain ```python and ```python:filepath variants.
"""

from __future__ import annotations

import re

from aiservice.common.llm_output_utils import truncate_pathological_output

# Matches both ```python and ```python:filepath blocks, captures content only
MARKDOWN_CODE_BLOCK_PATTERN = re.compile(r"```python(?::[^\n]*)?\n(.*?)```", re.DOTALL)

# Matches first ```python block (with optional :filepath), captures content.
# Uses greedy (.*) to handle LLM outputs with nested code fences (e.g. ```python:filepath
# blocks inside the main block). Requires closing ``` to be alone on its line.
FIRST_CODE_BLOCK_PATTERN = re.compile(r"^```python(?::[^\n]*)?\s*\n(.*)\n```[ \t]*$", re.MULTILINE | re.DOTALL)

# Fallback for incomplete code blocks (missing closing ```)
FIRST_CODE_BLOCK_FALLBACK_PATTERN = re.compile(r"^```python(?::[^\n]*)?\s*\n(.*)", re.MULTILINE | re.DOTALL)


def extract_all_code_from_markdown(markdown: str) -> str:
    """Extract all code from markdown code blocks.

    Handles both ```python and ```python:filepath blocks.
    Returns all code blocks joined with double newlines.
    """
    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.

    Returns the content of the first ```python block, or None if not found.
    Falls back to incomplete blocks (missing closing ```).
    Automatically truncates pathological LLM output patterns.
    """
    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 -> code content.

    Only works with ```<language>:filepath blocks. Plain ```<language> blocks are ignored.

    Args:
        markdown: The markdown text to parse.
        language: The language tag to match (default: "python"). For JavaScript/TypeScript,
                  this will also match common aliases (js, ts).

    """
    # Build pattern based on language - handle common aliases
    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.

    Returns (before_text, code_content, after_text) or None if not found.
    Handles optional filepath syntax (e.g., ```python:filepath.py).

    Prefers code blocks with filepath annotations over plain code blocks,
    as these are more likely to contain the actual code in multi-block responses
    where the LLM may include example or explanation blocks before the real code.
    """
    # First try to find a code block with filepath (e.g., ```python:file.py)
    # These are more likely to contain the actual optimized code
    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()

    # Fall back to any code block (with or without filepath)
    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.

    Args:
        code: The code to wrap.
        language: The language identifier for syntax highlighting (default: "python").

    Returns:
        The code wrapped in markdown code block syntax.

    """
    return f"```{language}\n{code}\n```"
