"""
Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
SPDX-License-Identifier: MIT
"""

import re
from typing import List, Dict, Any


"""
    Example input:
    [
        {"label": "tab", "bbox": [0.176, 0.74, 0.824, 0.82], "text": "<table><tr><td></td><td>HellaSwag</td><td>Obqa</td><td>WinoGrande</td><td>ARC-c</td><td>ARC-e</td><td>boolq</td><td>piqa</td><td>Avg</td></tr><tr><td>OPT-1.3B</td><td>53.65</td><td>33.40</td><td>59.59</td><td>29.44</td><td>50.80</td><td>60.83</td><td>72.36</td><td>51.44</td></tr><tr><td>Pythia-1.0B</td><td>47.16</td><td>31.40</td><td>53.43</td><td>27.05</td><td>48.99</td><td>57.83</td><td>69.21</td><td>48.30</td></tr><tr><td>Pythia-1.4B</td><td>52.01</td><td>33.20</td><td>57.38</td><td>28.50</td><td>54.00</td><td>63.27</td><td>70.95</td><td>51.33</td></tr><tr><td>TinyLlama-1.1B</td><td>59.20</td><td>36.00</td><td>59.12</td><td>30.10</td><td>55.25</td><td>57.83</td><td>73.29</td><td>52.99</td></tr></table>", "reading_order": 6},
        {"label": "cap", "bbox": [0.28, 0.729, 0.711, 0.74], "text": "Table 2: Zero-shot performance on commonsense reasoning tasks", "reading_order": 7},
        {"label": "para", "bbox": [0.176, 0.848, 0.826, 0.873], "text": "We of performance during training We tracked the accuracy of TinyLlama on common-\nsense reasoning benchmarks during its pre-training, as shown in Fig. 2 . Generally, the performance of", "reading_order": 8},
        {"label": "fnote", "bbox": [0.176, 0.88, 0.824, 0.912], "text": "${ }^{4}$ Due to a bug in the config file, the learning rate did not decrease immediately after warmup and remained at\nthe maximum value for several steps before we fixed this.", "reading_order": 9},
        {"label": "foot", "bbox": [0.496, 0.939, 0.501, 0.95], "text": "14", "reading_order": 10}
    ]
"""


def extract_table_from_html(html_string):
    """Extract and clean table tags from HTML string"""
    try:
        table_pattern = re.compile(r"<table.*?>.*?</table>", re.DOTALL)
        tables = table_pattern.findall(html_string)
        tables = [re.sub(r"<table[^>]*>", "<table>", table) for table in tables]
        return "\n".join(tables)
    except Exception as e:
        print(f"extract_table_from_html error: {str(e)}")
        return f"<table><tr><td>Error extracting table: {str(e)}</td></tr></table>"


class MarkdownConverter:
    """Convert structured recognition results to Markdown format"""

    def __init__(self):
        # Define heading levels for different section types
        self.heading_levels = {"title": "#", "sec": "##", "sub_sec": "###"}

        # Define which labels need special handling
        self.special_labels = {
            "tab",
            "fig",
            "title",
            "sec",
            "sub_sec",
            "list",
            "formula",
            "reference",
            "alg",
        }

    def try_remove_newline(self, text: str) -> str:
        try:
            # Preprocess text to handle line breaks
            text = text.strip()
            text = text.replace("-\n", "")

            # Handle Chinese text line breaks
            def is_chinese(char):
                return "\u4e00" <= char <= "\u9fff"

            lines = text.split("\n")
            processed_lines = []

            # Process all lines except the last one
            for i in range(len(lines) - 1):
                current_line = lines[i].strip()
                next_line = lines[i + 1].strip()

                # Always add the current line, but determine if we need a newline
                if current_line:  # If current line is not empty
                    if next_line:  # If next line is not empty
                        # For Chinese text handling
                        if is_chinese(current_line[-1]) and is_chinese(next_line[0]):
                            processed_lines.append(current_line)
                        else:
                            processed_lines.append(current_line + " ")
                    else:
                        # Next line is empty, add current line with newline
                        processed_lines.append(current_line + "\n")
                else:
                    # Current line is empty, add an empty line
                    processed_lines.append("\n")

            # Add the last line
            if lines and lines[-1].strip():
                processed_lines.append(lines[-1].strip())

            text = "".join(processed_lines)

            return text
        except Exception as e:
            print(f"try_remove_newline error: {str(e)}")
            return text  # Return original text on error

    def _handle_text(self, text: str) -> str:
        """
        Process regular text content, preserving paragraph structure
        """
        try:
            if not text:
                return ""

            if text.strip().startswith("\\begin{array}") and text.strip().endswith(
                "\\end{array}"
            ):
                text = "$$" + text + "$$"
            elif (
                (
                    "_{" in text
                    or "^{" in text
                    or "\\" in text
                    or "_ {" in text
                    or "^ {" in text
                )
                and ("$" not in text)
                and ("\\begin" not in text)
            ):
                text = "$" + text + "$"

            # Process formulas in text before handling other text processing
            text = self._process_formulas_in_text(text)

            text = self.try_remove_newline(text)

            # Return processed text
            return text
        except Exception as e:
            print(f"_handle_text error: {str(e)}")
            return text  # Return original text on error

    def _process_formulas_in_text(self, text: str) -> str:
        """
        Process mathematical formulas in text by iteratively finding and replacing formulas.
        - Identify inline and block formulas
        - Replace newlines within formulas with \\
        """
        try:
            # Define formula delimiters and their corresponding patterns
            delimiters = [
                ("$$", "$$"),  # Block formula with $$
                ("\\[", "\\]"),  # Block formula with \[ \]
                ("$", "$"),  # Inline formula with $
                ("\\(", "\\)"),  # Inline formula with \( \)
            ]

            # Process the text by iterating through each delimiter type
            result = text

            for start_delim, end_delim in delimiters:
                # Create a pattern that matches from start to end delimiter
                # Using a custom approach to avoid issues with nested delimiters
                current_pos = 0
                processed_parts = []

                while current_pos < len(result):
                    # Find the next start delimiter
                    start_pos = result.find(start_delim, current_pos)
                    if start_pos == -1:
                        # No more formulas of this type
                        processed_parts.append(result[current_pos:])
                        break

                    # Add text before the formula
                    processed_parts.append(result[current_pos:start_pos])

                    # Find the matching end delimiter
                    end_pos = result.find(end_delim, start_pos + len(start_delim))
                    if end_pos == -1:
                        # No matching end delimiter, treat as regular text
                        processed_parts.append(result[start_pos:])
                        break

                    # Extract the formula content (without delimiters)
                    formula_content = result[start_pos + len(start_delim) : end_pos]

                    # Process the formula content - replace newlines with \\
                    processed_formula = formula_content.replace("\n", " \\\\ ")

                    # Add the processed formula with its delimiters
                    processed_parts.append(
                        f"{start_delim}{processed_formula}{end_delim}"
                    )

                    # Move past this formula
                    current_pos = end_pos + len(end_delim)

                # Update the result with processed text
                result = "".join(processed_parts)
            return result
        except Exception as e:
            print(f"_process_formulas_in_text error: {str(e)}")
            return text  # Return original text on error

    def _remove_newline_in_heading(self, text: str) -> str:
        """
        Remove newline in heading
        """
        try:
            # Handle Chinese text line breaks
            def is_chinese(char):
                return "\u4e00" <= char <= "\u9fff"

            # Check if the text contains Chinese characters
            if any(is_chinese(char) for char in text):
                return text.replace("\n", "")
            else:
                return text.replace("\n", " ")

        except Exception as e:
            print(f"_remove_newline_in_heading error: {str(e)}")
            return text

    def _handle_heading(self, text: str, label: str) -> str:
        """
        Convert section headings to appropriate markdown format
        """
        try:
            level = self.heading_levels.get(label, "#")
            text = text.strip()
            text = self._remove_newline_in_heading(text)
            text = self._handle_text(text)
            return f"{level} {text}\n\n"
        except Exception as e:
            print(f"_handle_heading error: {str(e)}")
            return f"# Error processing heading: {text}\n\n"

    def _handle_list_item(self, text: str) -> str:
        """
        Convert list items to markdown list format
        """
        try:
            return f"- {text.strip()}\n"
        except Exception as e:
            print(f"_handle_list_item error: {str(e)}")
            return f"- Error processing list item: {text}\n"

    def _handle_figure(self, text: str, section_count: int) -> str:
        """
        Handle figure content
        """
        try:
            # Check if it's a file path starting with "figures/"
            if text.startswith("figures/"):
                # Convert to relative path from markdown directory to figures directory
                relative_path = f"../{text}"
                return f"![Figure {section_count}]({relative_path})\n\n"

            # Check if it's already a markdown format image link
            if text.startswith("!["):
                # Already in markdown format, return directly
                return f"{text}\n\n"

            # If it's still base64 format, maintain original logic
            if text.startswith("data:image/"):
                return f"![Figure {section_count}]({text})\n\n"
            elif ";" in text and "," in text:
                return f"![Figure {section_count}]({text})\n\n"
            else:
                # Assume it's raw base64, convert to data URI
                img_format = "png"
                data_uri = f"data:image/{img_format};base64,{text}"
                return f"![Figure {section_count}]({data_uri})\n\n"

        except Exception as e:
            print(f"_handle_figure error: {str(e)}")
            return f"*[Error processing figure: {str(e)}]*\n\n"

    def _handle_table(self, text: str) -> str:
        """
        Convert table content to markdown format
        """
        try:
            markdown_content = []
            if "<table" in text.lower() or "<tr" in text.lower():
                markdown_table = extract_table_from_html(text)
                markdown_content.append(markdown_table + "\n")
            else:
                table_lines = text.split("\n")
                if table_lines:
                    col_count = len(table_lines[0].split()) if table_lines[0] else 1
                    header = "| " + " | ".join(table_lines[0].split()) + " |"
                    markdown_content.append(header)
                    markdown_content.append(
                        "| " + " | ".join(["---"] * col_count) + " |"
                    )
                    for line in table_lines[1:]:
                        cells = line.split()
                        while len(cells) < col_count:
                            cells.append("")
                        markdown_content.append("| " + " | ".join(cells) + " |")
            return "\n".join(markdown_content) + "\n\n"
        except Exception as e:
            print(f"_handle_table error: {str(e)}")
            return f"*[Error processing table: {str(e)}]*\n\n"

    def _handle_algorithm(self, text: str) -> str:
        """
        Process algorithm blocks with proper formatting
        """
        try:
            # Remove algorithm environment tags if present
            text = re.sub(
                r"\\begin\{algorithm\}(.*?)\\end\{algorithm\}",
                r"\1",
                text,
                flags=re.DOTALL,
            )
            text = text.replace("\\begin{algorithm}", "").replace(
                "\\end{algorithm}", ""
            )
            text = text.replace("\\begin{algorithmic}", "").replace(
                "\\end{algorithmic}", ""
            )

            # Process the algorithm text
            lines = text.strip().split("\n")

            # Check if there's a caption or label
            caption = ""
            algorithm_text = []

            for line in lines:
                if "\\caption" in line:
                    # Extract caption text
                    caption_match = re.search(r"\\caption\{(.*?)\}", line)
                    if caption_match:
                        caption = f"**{caption_match.group(1)}**\n\n"
                    continue
                elif "\\label" in line:
                    continue  # Skip label lines
                else:
                    algorithm_text.append(line)

            # Join the algorithm text and wrap in code block
            formatted_text = "\n".join(algorithm_text)

            # Return the formatted algorithm with caption
            return f"{caption}```\n{formatted_text}\n```\n\n"
        except Exception as e:
            print(f"_handle_algorithm error: {str(e)}")
            return f"*[Error processing algorithm: {str(e)}]*\n\n{text}\n\n"

    def _handle_formula(self, text: str) -> str:
        """
        Handle formula-specific content
        """
        try:
            # Process the formula content
            processed_text = self._process_formulas_in_text(text)

            # For formula blocks, ensure they're properly formatted in markdown
            if "$$" not in processed_text and "\\[" not in processed_text:
                # If no block formula delimiters are present, wrap in $$ for block formula
                processed_text = f"$${processed_text}$$"

            return f"{processed_text}\n\n"
        except Exception as e:
            print(f"_handle_formula error: {str(e)}")
            return f"*[Error processing formula: {str(e)}]*\n\n"

    def convert(self, recognition_results: List[Dict[str, Any]]) -> str:
        """
        Convert recognition results to markdown format
        """
        try:
            markdown_content = []

            for section_count, result in enumerate(recognition_results):
                try:
                    label = result.get("label", "")
                    text = result.get("text", "").strip()

                    # Skip empty text
                    if not text:
                        continue

                    # Handle different content types
                    if label in {"title", "sec", "sub_sec"}:
                        markdown_content.append(self._handle_heading(text, label))
                    elif label == "list":
                        markdown_content.append(self._handle_list_item(text))
                    elif label == "fig":
                        markdown_content.append(
                            self._handle_figure(text, section_count)
                        )
                    elif label == "tab":
                        markdown_content.append(self._handle_table(text))
                    elif label == "alg":
                        markdown_content.append(self._handle_algorithm(text))
                    elif label == "formula":
                        markdown_content.append(self._handle_formula(text))
                    elif label not in self.special_labels:
                        # Handle regular text (paragraphs, etc.)
                        processed_text = self._handle_text(text)
                        markdown_content.append(f"{processed_text}\n\n")
                except Exception as e:
                    print(f"Error processing item {section_count}: {str(e)}")
                    # Add a placeholder for the failed item
                    markdown_content.append("*[Error processing content]*\n\n")

            # Join all content and apply post-processing
            result = "".join(markdown_content)
            return self._post_process(result)
        except Exception as e:
            print(f"convert error: {str(e)}")
            return f"Error generating markdown content: {str(e)}"

    def _post_process(self, markdown_content: str) -> str:
        """
        Apply post-processing fixes to the generated markdown content
        """
        try:
            # Handle author information
            author_pattern = re.compile(r"\\author\{(.*?)\}", re.DOTALL)

            def process_author_match(match):
                # Extract author content
                author_content = match.group(1)
                # Process the author content
                return self._handle_text(author_content)

            # Replace \author{...} with processed content
            markdown_content = author_pattern.sub(
                process_author_match, markdown_content
            )

            # Handle special case where author is inside math environment
            math_author_pattern = re.compile(r"\$(\\author\{.*?\})\$", re.DOTALL)
            match = math_author_pattern.search(markdown_content)
            if match:
                # Extract the author command
                author_cmd = match.group(1)
                # Extract content from author command
                author_content_match = re.search(
                    r"\\author\{(.*?)\}", author_cmd, re.DOTALL
                )
                if author_content_match:
                    # Get author content and process it
                    author_content = author_content_match.group(1)
                    processed_content = self._handle_text(author_content)
                    # Replace the entire $\author{...}$ block with processed content
                    markdown_content = markdown_content.replace(
                        match.group(0), processed_content
                    )

            # Replace LaTeX abstract environment with plain text
            markdown_content = re.sub(
                r"\\begin\{abstract\}(.*?)\\end\{abstract\}",
                r"**Abstract** \1",
                markdown_content,
                flags=re.DOTALL,
            )

            # Replace standalone \begin{abstract} (without matching end)
            markdown_content = re.sub(
                r"\\begin\{abstract\}", r"**Abstract**", markdown_content
            )

            # Replace LaTeX equation numbers with tag format, handling cases with extra backslashes
            markdown_content = re.sub(
                r"\\eqno\{\((.*?)\)\}", r"\\tag{\1}", markdown_content
            )

            # Find the starting tag of the formula
            markdown_content = markdown_content.replace("\[ \\\\", "$$ \\\\")

            # Find the ending tag of the formula (ensure this is the only ending tag)
            markdown_content = markdown_content.replace("\\\\ \]", "\\\\ $$")

            # Fix other common LaTeX issues
            replacements = [
                # Fix spacing issues in subscripts and superscripts
                (r"_ {", r"_{"),
                (r"^ {", r"^{"),
                # Fix potential issues with multiple consecutive newlines
                (r"\n{3,}", r"\n\n"),
            ]

            for old, new in replacements:
                markdown_content = re.sub(old, new, markdown_content)

            return markdown_content
        except Exception as e:
            print(f"_post_process error: {str(e)}")
            return markdown_content  # Return original content if post-processing fails
