"""Table structure evaluation (TEDS-like scoring).

Provides a simplified Tree-Edit-Distance-based Similarity score for tables
by comparing HTML table structures.

Usage:
    from benchmarks.evaluate_tables import compute_table_score, evaluate_tables
"""

from __future__ import annotations

import re
from html.parser import HTMLParser


def _split_markdown_row(line: str) -> list[str]:
    """Split a Markdown row, honoring optional outer and escaped pipes."""
    stripped = line.strip()
    if stripped.startswith("|"):
        stripped = stripped[1:]
    if stripped.endswith("|") and not stripped.endswith(r"\|"):
        stripped = stripped[:-1]
    return [
        cell.replace(r"\|", "|").strip()
        for cell in re.split(r"(?<!\\)\|", stripped)
    ]


def _is_separator_row(line: str) -> bool:
    cells = _split_markdown_row(line)
    return bool(cells) and all(re.fullmatch(r":?-{3,}:?", cell) for cell in cells)


def _extract_markdown_tables_with_positions(text: str) -> list[tuple[int, str]]:
    """Extract Markdown tables together with their character offsets."""
    lines = text.split("\n")
    offsets = []
    offset = 0
    for line in lines:
        offsets.append(offset)
        offset += len(line) + 1

    tables: list[tuple[int, str]] = []
    index = 0
    while index + 1 < len(lines):
        table_start = index
        header = lines[index].strip()
        separator = lines[index + 1].strip()
        if "|" not in header or not _is_separator_row(separator):
            index += 1
            continue

        header_cells = _split_markdown_row(header)
        separator_cells = _split_markdown_row(separator)
        if len(header_cells) != len(separator_cells):
            index += 1
            continue

        current_table = [header, separator]
        index += 2
        while index < len(lines):
            row = lines[index].strip()
            if not row or "|" not in row:
                break
            current_table.append(row)
            index += 1
        tables.append((offsets[table_start], "\n".join(current_table)))

    return tables


def extract_tables_from_markdown(text: str) -> list[str]:
    """Extract markdown tables from OCR output."""
    return [table for _, table in _extract_markdown_tables_with_positions(text)]


def extract_tables_from_html(text: str) -> list[str]:
    """Extract HTML tables from OCR output."""
    return [table for _, table in _extract_html_tables_with_positions(text)]


def _extract_html_tables_with_positions(text: str) -> list[tuple[int, str]]:
    """Extract balanced outer HTML tables, preserving nested table markup."""
    tag_pattern = re.compile(r"</?table\b[^>]*>", re.IGNORECASE)
    tables: list[tuple[int, str]] = []
    depth = 0
    start: int | None = None
    for match in tag_pattern.finditer(text):
        if re.match(r"<\s*/", match.group(0)):
            if depth == 0:
                continue
            depth -= 1
            if depth == 0 and start is not None:
                tables.append((start, text[start:match.end()]))
                start = None
        else:
            if depth == 0:
                start = match.start()
            depth += 1
    return tables


def _extract_tables_in_document_order(text: str) -> list[str]:
    """Extract Markdown and HTML tables without reordering their occurrences."""
    html_tables = _extract_html_tables_with_positions(text)
    html_ranges = [(start, start + len(table)) for start, table in html_tables]
    located = [
        (start, table)
        for start, table in _extract_markdown_tables_with_positions(text)
        if not any(begin <= start < end for begin, end in html_ranges)
    ]
    located.extend(html_tables)
    located.sort(key=lambda item: item[0])
    return [table for _, table in located]


class _HTMLTableParser(HTMLParser):
    """Collect rows and cells from one HTML table without external parsers."""

    def __init__(self):
        super().__init__(convert_charrefs=True)
        self.table_depth = 0
        self.rows: list[list[str]] = []
        self._row: list[str] | None = None
        self._cell_parts: list[str] | None = None

    def handle_starttag(self, tag: str, attrs):
        tag = tag.lower()
        if tag == "table":
            self.table_depth += 1
        elif self.table_depth == 1 and tag == "tr":
            self._row = []
        elif self.table_depth == 1 and tag in {"td", "th"} and self._row is not None:
            self._cell_parts = []
        elif self.table_depth == 1 and tag == "br" and self._cell_parts is not None:
            self._cell_parts.append("\n")

    def handle_data(self, data: str):
        if self.table_depth == 1 and self._cell_parts is not None:
            self._cell_parts.append(data)

    def handle_endtag(self, tag: str):
        tag = tag.lower()
        if self.table_depth == 1 and tag in {"td", "th"}:
            if self._row is not None and self._cell_parts is not None:
                value = re.sub(r"\s+", " ", "".join(self._cell_parts)).strip()
                self._row.append(value)
            self._cell_parts = None
        elif self.table_depth == 1 and tag == "tr":
            if self._row is not None and self._row:
                self.rows.append(self._row)
            self._row = None
            self._cell_parts = None
        elif tag == "table" and self.table_depth:
            self.table_depth -= 1


def _normalize_html_table(table_text: str) -> dict:
    parser = _HTMLTableParser()
    parser.feed(table_text)
    parser.close()
    rows = parser.rows
    if not rows:
        return {"num_rows": 0, "num_cols": 0, "cells": []}
    return {
        "num_rows": len(rows),
        "num_cols": max(len(row) for row in rows),
        "cells": [cell for row in rows for cell in row],
    }


def normalize_table(table_text: str) -> dict:
    """Parse a markdown table into a normalized structure.

    Returns dict with: num_rows, num_cols, cells (flattened list).
    """
    if re.search(r"<table\b", table_text, re.IGNORECASE):
        return _normalize_html_table(table_text)

    lines = [line.strip() for line in table_text.strip().split("\n") if line.strip()]

    # Filter out the Markdown header separator.
    data_lines = [line for line in lines if not _is_separator_row(line)]

    rows = []
    for line in data_lines:
        rows.append(_split_markdown_row(line))

    if not rows:
        return {"num_rows": 0, "num_cols": 0, "cells": []}

    num_cols = max(len(r) for r in rows) if rows else 0
    cells = []
    for row in rows:
        cells.extend(row)

    return {
        "num_rows": len(rows),
        "num_cols": num_cols,
        "cells": cells,
    }


def compute_table_score(reference_table: str, hypothesis_table: str) -> float:
    """Compute a simplified TEDS-like score between two tables.

    Score is based on:
    - Structure match (rows x cols): 40% weight
    - Cell content match: 60% weight

    Returns score in [0, 1] where 1 = perfect match.
    """
    ref = normalize_table(reference_table)
    hyp = normalize_table(hypothesis_table)

    if ref["num_rows"] == 0 and hyp["num_rows"] == 0:
        return 1.0
    if ref["num_rows"] == 0 or hyp["num_rows"] == 0:
        return 0.0

    # Structure score
    row_match = 1.0 - abs(ref["num_rows"] - hyp["num_rows"]) / max(ref["num_rows"], hyp["num_rows"])
    col_match = 1.0 - abs(ref["num_cols"] - hyp["num_cols"]) / max(ref["num_cols"], hyp["num_cols"])
    structure_score = (row_match + col_match) / 2

    # Cell content score (compare flattened cells)
    ref_cells = ref["cells"]
    hyp_cells = hyp["cells"]

    if not ref_cells:
        return structure_score * 0.4

    # Simple cell-by-cell comparison
    max_len = max(len(ref_cells), len(hyp_cells))
    matches = 0
    for i in range(min(len(ref_cells), len(hyp_cells))):
        # Normalize whitespace for comparison
        rc = re.sub(r"\s+", " ", ref_cells[i]).strip().lower()
        hc = re.sub(r"\s+", " ", hyp_cells[i]).strip().lower()
        if rc == hc:
            matches += 1
        elif rc and hc and (rc in hc or hc in rc):
            matches += 0.5

    cell_score = matches / max_len if max_len > 0 else 0.0

    # Weighted combination
    return 0.4 * structure_score + 0.6 * cell_score


def evaluate_tables(reference_text: str, hypothesis_text: str) -> dict:
    """Evaluate table extraction quality between reference and hypothesis.

    Returns dict with: num_ref_tables, num_hyp_tables, mean_score, per_table_scores.
    """
    ref_tables = _extract_tables_in_document_order(reference_text)
    hyp_tables = _extract_tables_in_document_order(hypothesis_text)

    if not ref_tables:
        return {
            "num_ref_tables": 0,
            "num_hyp_tables": len(hyp_tables),
            "mean_score": 1.0 if not hyp_tables else 0.0,
            "per_table_scores": [0.0] * len(hyp_tables),
        }

    # Match tables by position (simplified)
    scores = []
    for i, ref_table in enumerate(ref_tables):
        if i < len(hyp_tables):
            score = compute_table_score(ref_table, hyp_tables[i])
        else:
            score = 0.0  # Missing table
        scores.append(score)

    # Hallucinated extra tables must count against the result as well.
    scores.extend([0.0] * max(0, len(hyp_tables) - len(ref_tables)))

    return {
        "num_ref_tables": len(ref_tables),
        "num_hyp_tables": len(hyp_tables),
        "mean_score": float(sum(scores) / len(scores)) if scores else 0.0,
        "per_table_scores": scores,
    }