import csv
import io
import json
import re

from bs4 import BeautifulSoup

def skip_whitespace(text, i):
    """Advance index i past any whitespace."""
    while i < len(text) and text[i].isspace():
        i += 1
    return i

def parse_braced_argument(text, i):
    """
    Given text and an index i that should point at an opening '{',
    return a tuple (argument_content, new_index) where argument_content is the full
    string inside the balanced braces and new_index is the position just after the matching '}'.
    """
    if i >= len(text) or text[i] != '{':
        raise ValueError("Expected '{' at position {}".format(i))
    i += 1  # skip the opening brace
    start = i
    level = 1
    while i < len(text) and level > 0:
        if text[i] == '{':
            level += 1
        elif text[i] == '}':
            level -= 1
        i += 1
    if level != 0:
        raise ValueError("Unbalanced braces starting at position {}".format(start-1))
    # The argument content is from start to i-1 (excluding the closing brace)
    return text[start:i-1], i

def parse_command(text, i):
    """
    Parse a \multirow or \multicolumn command starting at index i.
    This function assumes the command has exactly three braced arguments.

    It processes each argument recursively. For the third argument, after recursive processing,
    it replaces any unescaped & with \&.

    Returns a tuple (command_text, new_index) where command_text is the reconstructed command.
    """
    # Determine which command we have.
    if text.startswith(r"\multirow", i):
        command_name = r"\multirow"
        i += len(r"\multirow")
    elif text.startswith(r"\multicolumn", i):
        command_name = r"\multicolumn"
        i += len(r"\multicolumn")
    else:
        raise ValueError("Expected \\multirow or \\multicolumn at position {}".format(i))

    # Skip whitespace between the command name and the first argument.
    i = skip_whitespace(text, i)
    args = []
    # Expect exactly three arguments
    for arg_index in range(3):
        if i >= len(text) or text[i] != '{':
            raise ValueError("Expected '{' for argument {} at position {}".format(arg_index+1, i))
        arg_content, i = parse_braced_argument(text, i)
        # Process the content recursively to catch nested commands
        processed_arg = clean_multi_cells(arg_content)
        if arg_index == 2:
            # For the cell text (third argument), replace any unescaped &
            processed_arg = re.sub(r'(?<!\\)&', r'\\&', processed_arg)
        args.append(processed_arg)
        # Only skip whitespace between arguments, not after the last one.
        if arg_index < 2:
            i = skip_whitespace(text, i)
    # Reconstruct the full command with its three arguments
    command_text = f"{command_name}{{{args[0]}}}{{{args[1]}}}{{{args[2]}}}"
    return command_text, i

def clean_multi_cells(text):
    """
    Process an arbitrary LaTeX text string and look for occurrences of \multirow or \multicolumn commands.
    When found, the command is parsed (handling nested braces and nested commands) and its third argument is fixed.

    Returns the processed text.
    """
    result = []
    i = 0
    while i < len(text):
        # Find next occurrence of either command.
        idx_multi = text.find(r"\multirow", i)
        idx_multiC = text.find(r"\multicolumn", i)

        # Determine the next index among the two (if any)
        if idx_multi == -1 and idx_multiC == -1:
            result.append(text[i:])
            break
        if idx_multi == -1:
            next_idx = idx_multiC
        elif idx_multiC == -1:
            next_idx = idx_multi
        else:
            next_idx = min(idx_multi, idx_multiC)

        # Append text before the command (preserving any whitespace)
        result.append(text[i:next_idx])
        # Process the command starting at next_idx
        command_text, new_index = parse_command(text, next_idx)
        result.append(command_text)
        i = new_index
    return ''.join(result)

def parse_brace(s, pos):
    """
    Given a string s and an index pos pointing to an opening '{',
    returns a tuple (content, new_pos) where content is the string
    between the matching braces (handling nested braces) and new_pos is
    the index just after the closing '}'.
    """
    if pos >= len(s) or s[pos] != '{':
        raise ValueError("Expected '{' at position %d" % pos)
    pos += 1  # skip the opening brace
    content = ""
    depth = 1
    while pos < len(s) and depth:
        char = s[pos]
        if char == '{':
            depth += 1
            content += char
        elif char == '}':
            depth -= 1
            if depth:
                content += char
        else:
            content += char
        pos += 1
    if depth != 0:
        raise ValueError("Unmatched '{' in string.")
    return content, pos

def parse_command_merge(s, pos):
    """
    Parse a multirow or multicolumn command starting at s[pos]. If the content
    of the command contains a nested command, then recursively parse the inner
    command and merge its parameters with the outer ones. The merging is done
    so that the outer multirow’s parameters (e.g. rowspan and width) are kept
    while the inner command’s parameters (e.g. colspan, alignment) and its innermost
    content are returned.

    Returns a tuple (merged_dict, new_pos) where merged_dict is a dictionary
    containing the combined parameters and new_pos is the updated index after
    parsing the command.
    """
    if s.startswith(r"\multirow", pos):
        newpos = pos + len(r"\multirow")
        # Parse the three required arguments for multirow: rowspan, width, and content.
        rowspan, newpos = parse_brace(s, newpos)
        width, newpos = parse_brace(s, newpos)
        content, newpos = parse_brace(s, newpos)
        # Look for a nested command (either \multirow or \multicolumn) in the content.
        index_mr = content.find(r"\multirow")
        index_mc = content.find(r"\multicolumn")
        if index_mr == -1 and index_mc == -1:
            # No nested command found; return this command’s details.
            return {"rowspan": rowspan.strip(), "width": width.strip(), "content": content.strip()}, newpos
        else:
            # At least one nested command is present. Pick the first occurrence.
            indices = [i for i in (index_mr, index_mc) if i != -1]
            first_index = min(indices)
            # Parse the inner (nested) command from within the content.
            inner, _ = parse_command_merge(content, first_index)
            # Merge: keep the outer multirow’s parameters and add the inner ones.
            merged = {"rowspan": rowspan.strip(), "width": width.strip()}
            merged.update(inner)
            return merged, newpos

    elif s.startswith(r"\multicolumn", pos):
        newpos = pos + len(r"\multicolumn")
        # Parse the three arguments for multicolumn: colspan, alignment, and content.
        colspan, newpos = parse_brace(s, newpos)
        alignment, newpos = parse_brace(s, newpos)
        content, newpos = parse_brace(s, newpos)
        # Look for a nested command in the content.
        index_mr = content.find(r"\multirow")
        index_mc = content.find(r"\multicolumn")
        if index_mr == -1 and index_mc == -1:
            return {"colspan": colspan.strip(), "alignment": alignment.strip(), "content": content.strip()}, newpos
        else:
            indices = [i for i in (index_mr, index_mc) if i != -1]
            first_index = min(indices)
            inner, _ = parse_command_merge(content, first_index)
            merged = {"colspan": colspan.strip(), "alignment": alignment.strip()}
            merged.update(inner)
            return merged, newpos

    # Not a recognized command starting at pos.
    return None, pos

def extract_merged_commands(s):
    """
    Scan through the LaTeX string s and extract merged multirow/multicolumn commands.
    For each command found, if there is nesting the parser merges the outer and inner
    parameters so that the final result includes both the rowspan (or width) and the colspan
    (or alignment) along with the innermost content.

    Returns a list of dictionaries.
    """
    pos = 0
    results = []
    while pos < len(s):
        if s[pos] == '\\':
            res, newpos = parse_command_merge(s, pos)
            if res is not None:
                results.append(res)
                pos = newpos
                continue
        pos += 1
    return results

def remove_tags(html, tags_to_remove):
    soup = BeautifulSoup(html, "html.parser")
    # Loop through the tags to remove
    for tag_name in tags_to_remove:
        for tag in soup.find_all(tag_name):
            # Move the children of the tag to the parent tag
            tag.unwrap()  # This removes the tag but keeps its contents
    # Return the modified HTML as a string
    return str(soup)

def convert_th_to_td(html):
    """Replace all th tags with td tags
    """
    soup = BeautifulSoup(html)
    for th_tag in soup.find_all('th'):
        th_tag.name = 'td'
    return str(soup)

def replace_italic(text):
    pattern = re.compile(r'(?<!\\)_(.*?)(?<!\\)_')

    def italic_replacer(match):
        # Get the text inside the underscores.
        content = match.group(1)
        # Remove the escape (backslash) from any escaped underscores inside.
        content = content.replace(r'\_', '_')
        return f"<i>{content}</i>"

    # Replace all occurrences of the pattern using the replacer function.
    return pattern.sub(italic_replacer, text)


def replace_bold(text):
    pattern = re.compile(r'(?<!\\)\*\*(.*?)(?<!\\)\*\*')

    def bold_replacer(match):
        content = match.group(1)
        # Unescape any escaped asterisks within the captured text.
        content = content.replace(r'\*', '*')
        return f"<b>{content}</b>"

    return pattern.sub(bold_replacer, text)

def latex_table_to_html(latex_str, add_head_body=True):
    table_pattern = r'\\begin{tabular}{([^}]*)}\s*(.*?)\\end{tabular}'

    def process_cell(cell):
        cell = cell.strip()

        out = extract_merged_commands(cell)
        if len(out) > 0:
            cell = process_cell(out[0]["content"])["content"]
            rowspan = int(out[0].get("rowspan", "1"))
            colspan = int(out[0].get("colspan", "1"))
            return {
                "content": cell,
                "colspan": colspan,
                "rowspan": rowspan
            }

        cell = re.sub(r'\$([^$]*)\$', r'\1', cell)
        cell = re.sub(r'\\textbf{([^}]*)}', r'<b>\1</b>', cell)
        cell = re.sub(r'\\textit{([^}]*)}', r'<i>\1</i>', cell)
        cell = replace_italic(cell)
        cell = replace_bold(cell)
        cell = cell.replace("\\$", "$").replace("\\%", "%").replace("\\#", "#").replace("\\_", "_")
        cell = cell.replace("\\textasciitilde{}", "~").replace("\\textasciicircum{}", "^")
        cell = cell.replace("\\newline", "\n").replace("\\textless", "<").replace("\\textgreater", ">").replace("\\*", "*").replace("\\backslash", "\\")

        cell = cell.replace(r'\&', '&')
        cell = cell.replace('<tbc>', '')
        cell = cell.replace('\\unknown', '').replace('\\<|unk|\\>', '').replace('<u>', '<underline>').replace('</u>', '</underline>')
        return {
            'content': cell,
            'colspan': 1,
            'rowspan': 1
        }

    def split_row(input_string):
        return re.split(r'(?<!\\)&', input_string)

    def convert_table(match):
        format_spec, content = match.groups()

        # --- Pass 1: collect all processed rows ---
        all_rows = []
        rows = re.split(r'\\\\', content)

        for row in rows:
            if not row.strip():
                continue
            row = row.strip()
            if '\\hline' in row:
                row = row.replace('\\hline', '')
                if not row.strip():
                    continue
            row = clean_multi_cells(row)
            cells = split_row(row)
            processed_cells = [process_cell(cell) for cell in cells]

            def split_lines(text):
                parts = re.split(r'(?:\n|<br\s*/?>)+', text)
                return parts if parts is not None else ['']

            line_lists = [split_lines(cell['content']) for cell in processed_cells]
            max_lines = max(len(lst) for lst in line_lists) if line_lists else 1
            all_rows.append((processed_cells, line_lists, max_lines))

        if not all_rows:
            return '<table></table>'

        # --- Detect header depth from max rowspan in row 0 ---
        first_row_cells = all_rows[0][0]
        header_row_count = max(
            (cell['rowspan'] for cell in first_row_cells),
            default=1
        )

        # --- Strip trailing empty cells from each row ---
        for i in range(len(all_rows)):
            cells, line_lists, ml = all_rows[i]
            while (cells
                   and cells[-1]['content'].strip() == ''
                   and cells[-1]['colspan'] == 1
                   and cells[-1]['rowspan'] == 1):
                cells = cells[:-1]
                line_lists = line_lists[:-1]
            ml = max(len(lst) for lst in line_lists) if line_lists else 1
            all_rows[i] = (cells, line_lists, ml)

        # --- Pass 2: emit HTML ---
        html = ['<table>']
        multirow_tracker = set()
        current_row = 0

        for row_idx, (processed_cells, line_lists, max_lines) in enumerate(all_rows):
            for line_idx in range(max_lines):
                is_header = current_row < header_row_count

                if add_head_body and current_row == 0:
                    html.append(' <thead>')
                if add_head_body and current_row == header_row_count:
                    html.append(' <tbody>')

                html.append('  <tr>')
                current_col = 0

                for col_idx, cell in enumerate(processed_cells):
                    content_segment = line_lists[col_idx][line_idx] if line_idx < len(line_lists[col_idx]) else ''

                    attrs = []
                    if cell['colspan'] > 1:
                        attrs.append(f'colspan="{cell["colspan"]}"')
                    if cell['rowspan'] > 1 and line_idx == 0:
                        attrs.append(f'rowspan="{cell["rowspan"]}"')
                        for r in range(current_row + 1, current_row + cell['rowspan']):
                            for c in range(current_col, current_col + cell['colspan']):
                                multirow_tracker.add((r, c))

                    if cell['rowspan'] > 1 and line_idx > 0:
                        current_col += cell['colspan']
                        continue

                    if (current_row, current_col) in multirow_tracker and content_segment == '' and cell["colspan"] == 1 and cell["rowspan"] == 1:
                        current_col += cell['colspan']
                        continue

                    attr_str = ' ' + ' '.join(attrs) if attrs else ''
                    cell_tag = 'th' if is_header and add_head_body else 'td'
                    html.append(f'    <{cell_tag}{attr_str}>{content_segment}</{cell_tag}>')
                    current_col += cell['colspan']

                html.append('  </tr>')
                if add_head_body and current_row == header_row_count - 1:
                    html.append(' </thead>')
                current_row += 1

        if add_head_body and current_row > header_row_count:
            html.append(' </tbody>')
        html.append('</table>')
        return '\n'.join(html)

    return re.sub(table_pattern, convert_table, latex_str, flags=re.DOTALL)
def convert_single_table(table):
    """
    Convert a single HTML table to Markdown format.

    Multi-row headers are flattened into a single header row using
    dot-notation column paths built from ``_expand_spans`` /
    ``_build_column_paths``.

    Args:
        table: BeautifulSoup table element

    Returns:
        str: Markdown table string
    """
    grid, header_rows = _expand_spans(table)
    if not grid:
        return ''

    paths = _build_column_paths(grid, header_rows)
    data_rows = grid[header_rows:]

    markdown_lines = []

    header_data = [p.replace('|', '\\|') for p in paths]
    markdown_lines.append('| ' + ' | '.join(header_data) + ' |')
    markdown_lines.append('| ' + ' | '.join(['---'] * len(paths)) + ' |')

    for row in data_rows:
        row_data = [cell.replace('|', '\\|') for cell in row]
        markdown_lines.append('| ' + ' | '.join(row_data) + ' |')

    return '\n'.join(markdown_lines)
def convert_html_tables_to_markdown(html_content):
    """
    Find all HTML tables and convert them to Markdown while preserving all other content.
    
    Args:
        html_content (str): HTML content that may contain tables
        
    Returns:
        str: HTML content with tables converted to Markdown
    """
    soup = BeautifulSoup(html_content, 'html.parser')
    
    # Find all tables
    tables = soup.find_all('table')
    
    if not tables:
        return html_content  # Return original content unchanged
    
    # Convert each table to markdown and replace it
    for table in tables:
        markdown_table = convert_single_table(table)
        
        # Create a new element to replace the table
        replacement = soup.new_string('\n' + markdown_table + '\n')
        table.replace_with(replacement)
    
    return str(soup)


def _expand_spans(table):
    """
    Expand colspan/rowspan in a BeautifulSoup <table> into a regular grid.
    Returns a tuple ``(grid, header_rows)`` where *grid* is a list-of-lists
    of cell text strings and *header_rows* is the number of rows inside
    ``<thead>`` (defaults to 1 when no ``<thead>`` is present).
    """
    thead = table.find('thead')
    header_rows = len(thead.find_all('tr')) if thead else 1

    rows = table.find_all('tr')
    if not rows:
        return [], 0

    num_cols = 0
    for row in rows:
        col_count = sum(int(c.get('colspan', 1)) for c in row.find_all(['td', 'th']))
        num_cols = max(num_cols, col_count)

    grid = [[None] * num_cols for _ in range(len(rows))]

    for r, row in enumerate(rows):
        cells = row.find_all(['td', 'th'])
        c = 0
        for cell in cells:
            while c < num_cols and grid[r][c] is not None:
                c += 1
            if c >= num_cols:
                break

            colspan = int(cell.get('colspan', 1))
            rowspan = int(cell.get('rowspan', 1))
            text = cell.get_text(separator=' ', strip=True)

            for dr in range(rowspan):
                for dc in range(colspan):
                    rr, cc = r + dr, c + dc
                    if rr < len(grid) and cc < num_cols:
                        grid[rr][cc] = text
            c += colspan

    for row in grid:
        for i in range(len(row)):
            if row[i] is None:
                row[i] = ''

    return grid, header_rows


def _build_column_paths(grid, header_rows):
    """Build hierarchical column names from multi-row headers.

    For each column, walk down rows ``0..header_rows-1``, collect distinct
    non-empty labels (skipping duplicates produced by rowspan fill), and
    join with ``'.'``.
    """
    num_cols = len(grid[0]) if grid else 0
    paths = []
    for c in range(num_cols):
        parts = []
        prev = None
        for r in range(header_rows):
            val = grid[r][c].strip() if r < len(grid) else ""
            if val and val != prev:
                parts.append(val)
                prev = val
        paths.append(".".join(parts) if parts else f"column_{c}")
    return paths


def _deduplicate_headers(raw_headers):
    """Append ``_N`` suffixes to make every header string unique."""
    headers = []
    seen = {}
    for h in raw_headers:
        key = h if h else "column"
        if key in seen:
            seen[key] += 1
            key = f"{key}_{seen[key]}"
        else:
            seen[key] = 0
        headers.append(key)
    return headers


def _grid_to_records(grid, header_rows=1):
    """
    Convert a grid (list-of-lists) to a list of dicts.  When *header_rows* > 1
    the column names are built by joining the distinct labels down each column
    with ``'.'``.  Duplicate or empty headers get disambiguated with a suffix.
    """
    if len(grid) < header_rows + 1:
        return grid

    raw_headers = _build_column_paths(grid, header_rows)
    headers = _deduplicate_headers(raw_headers)
    return [dict(zip(headers, row)) for row in grid[header_rows:]]


def convert_html_table_to_json(html_content, records=True):
    """
    Convert HTML tables produced by ``latex_table_to_html`` into JSON.

    Parameters
    ----------
    html_content : str
        HTML string (may contain one or more ``<table>`` elements and
        surrounding non-table text).
    records : bool
        If *True* (default) the header rows are used to build hierarchical
        column names (dot-separated) and each subsequent row becomes a
        ``{header: value, ...}`` dict.
        If *False* the table is returned as a plain list-of-lists.

    Returns
    -------
    str
        A JSON string.  If the input contains a single table the top-level
        value is a list (of dicts or lists).  If there are multiple tables
        the top-level value is a list of such lists.
    """
    soup = BeautifulSoup(html_content, 'html.parser')
    tables = soup.find_all('table')

    if not tables:
        return html_content

    results = []
    for table in tables:
        grid, header_rows = _expand_spans(table)
        if records:
            results.append(_grid_to_records(grid, header_rows))
        else:
            results.append(grid)

    payload = results[0] if len(results) == 1 else results
    return json.dumps(payload, ensure_ascii=False, indent=2)


def _extract_cell_objects(table):
    """Extract rows of cell-objects that preserve span metadata.

    Each cell is represented as a dict with ``text``, ``rowspan``,
    ``colspan``, and ``is_header`` fields.  Cells covered by a prior
    span are **not** duplicated — only the originating cell appears.
    This makes the representation lossless: you can reconstruct the
    HTML (or LaTeX) ``rowspan``/``colspan`` structure from it.

    Returns ``(header_cell_rows, body_cell_rows, num_cols)`` where each
    ``*_cell_rows`` is a list of lists of cell-dicts.
    """
    thead = table.find('thead')
    tbody = table.find('tbody')

    def rows_from(section):
        return section.find_all('tr') if section else []

    header_trs = rows_from(thead)
    if tbody:
        body_trs = rows_from(tbody)
    else:
        all_trs = table.find_all('tr')
        body_trs = all_trs[len(header_trs):]

    num_cols = 0
    for row in table.find_all('tr'):
        col_count = sum(int(c.get('colspan', 1)) for c in row.find_all(['td', 'th']))
        num_cols = max(num_cols, col_count)

    def process_rows(trs):
        out = []
        for tr in trs:
            row_cells = []
            for cell in tr.find_all(['td', 'th']):
                colspan = int(cell.get('colspan', 1))
                rowspan = int(cell.get('rowspan', 1))
                text = cell.get_text(separator=' ', strip=True)
                is_header = cell.name == 'th'
                obj = {"text": text}
                if rowspan != 1:
                    obj["rowspan"] = rowspan
                if colspan != 1:
                    obj["colspan"] = colspan
                if is_header:
                    obj["is_header"] = True
                row_cells.append(obj)
            out.append(row_cells)
        return out

    return process_rows(header_trs), process_rows(body_trs), num_cols


def convert_html_table_to_json_hierarchical(html_content):
    """
    Convert HTML tables into a structured JSON format that explicitly
    separates header rows from data rows and preserves cell-level span
    metadata for lossless round-tripping.

    Returns
    -------
    str
        A JSON string.  Each table becomes an object with:

        - ``header_rows`` (int): number of header rows.
        - ``columns`` (list[str]): pre-flattened dot-notation column
          names (e.g. ``"International.EMEA.Q1"``).
        - ``headers`` (list[list[cell]]): header rows where each cell
          is either a plain string or an object with ``text``,
          ``rowspan``, and/or ``colspan`` keys.  Cells spanned by a
          prior cell are omitted (not duplicated), so the original
          merge structure is fully recoverable.
        - ``data`` (list[list[str]]): body rows (flat string arrays).

        If the input contains a single table the top-level value is a
        single object; multiple tables produce a list of objects.
    """
    soup = BeautifulSoup(html_content, 'html.parser')
    tables = soup.find_all('table')

    if not tables:
        return html_content

    results = []
    for table in tables:
        header_cells, body_cells, num_cols = _extract_cell_objects(table)
        grid, header_row_count = _expand_spans(table)
        columns = _build_column_paths(grid, header_row_count)
        columns = _deduplicate_headers(columns)

        def simplify_cell(obj):
            if obj.get("rowspan", 1) == 1 and obj.get("colspan", 1) == 1:
                return obj["text"]
            return {k: v for k, v in obj.items() if k != "is_header"}

        headers = [[simplify_cell(c) for c in row] for row in header_cells]
        data = [[c["text"] for c in row] for row in body_cells]

        results.append({
            "header_rows": header_row_count,
            "columns": columns,
            "headers": headers,
            "data": data,
        })

    payload = results[0] if len(results) == 1 else results
    return json.dumps(payload, ensure_ascii=False, indent=2)


def convert_html_table_to_csv(html_content):
    """
    Convert HTML tables produced by ``latex_table_to_html`` into CSV.

    Multi-row headers are collapsed into a single CSV header row using
    dot-notation column paths.  Each table becomes a CSV block separated
    by a blank line.

    Returns
    -------
    str
        CSV-formatted string.
    """
    soup = BeautifulSoup(html_content, 'html.parser')
    tables = soup.find_all('table')

    if not tables:
        return html_content

    csv_blocks = []
    for table in tables:
        grid, header_rows = _expand_spans(table)
        paths = _build_column_paths(grid, header_rows)
        data_rows = grid[header_rows:]

        buf = io.StringIO()
        writer = csv.writer(buf)
        writer.writerow(paths)
        writer.writerows(data_rows)
        csv_blocks.append(buf.getvalue().rstrip('\r\n'))

    return '\n\n'.join(csv_blocks)