from __future__ import annotations

import re
from typing import Iterable

import pandas as pd


HEADING_RE = re.compile(r"^\s{0,3}(#{1,6})\s+(.*?)\s*$")
TABLE_DIVIDER_CELL_RE = re.compile(r"^:?-{3,}:?$")


def extract_markdown_section(markdown: str, heading: str) -> str | None:
    lines = markdown.splitlines()
    heading_normalized = heading.strip().lower()

    start_index: int | None = None
    heading_level: int | None = None

    for index, line in enumerate(lines):
        match = HEADING_RE.match(line)
        if not match:
            continue
        title = match.group(2).strip().lower()
        if title == heading_normalized:
            start_index = index + 1
            heading_level = len(match.group(1))
            break

    if start_index is None or heading_level is None:
        return None

    end_index = len(lines)
    for index in range(start_index, len(lines)):
        match = HEADING_RE.match(lines[index])
        if not match:
            continue
        next_heading_level = len(match.group(1))
        if next_heading_level <= heading_level:
            end_index = index
            break

    section = "\n".join(lines[start_index:end_index]).strip()
    return section or None


def extract_performance_metrics_table(markdown: str, heading: str = "Performance Metrics") -> pd.DataFrame | None:
    section = extract_markdown_section(markdown, heading)
    if not section:
        return None

    blocks = _extract_table_blocks(section)
    for block in blocks:
        table = _parse_markdown_table_block(block)
        if table is None or table.empty:
            continue

        normalized_columns = {str(col).strip().lower() for col in table.columns}
        if "class" in normalized_columns:
            return table

    if blocks:
        fallback = _parse_markdown_table_block(blocks[0])
        if fallback is not None and not fallback.empty:
            return fallback

    return None


def _extract_table_blocks(text: str) -> list[list[str]]:
    lines = text.splitlines()
    blocks: list[list[str]] = []
    current_block: list[str] = []

    for line in lines:
        if "|" in line:
            current_block.append(line)
            continue

        if len(current_block) >= 2:
            blocks.append(current_block)
        current_block = []

    if len(current_block) >= 2:
        blocks.append(current_block)

    return blocks


def _parse_markdown_table_block(lines: Iterable[str]) -> pd.DataFrame | None:
    rows = [_split_markdown_row(line) for line in lines if line.strip()]
    if len(rows) < 2:
        return None

    header_index = None
    for index in range(len(rows) - 1):
        if _is_divider_row(rows[index + 1]):
            header_index = index
            break

    if header_index is None:
        return None

    header = [cell.strip() for cell in rows[header_index]]
    if not header or len(header) < 2:
        return None

    data_rows = rows[header_index + 2 :]
    normalized_rows: list[list[str]] = []

    for row in data_rows:
        if _is_divider_row(row):
            continue

        padded = row[: len(header)] + [""] * max(0, len(header) - len(row))
        if any(cell.strip() for cell in padded):
            normalized_rows.append(padded)

    if not normalized_rows:
        return None

    dataframe = pd.DataFrame(normalized_rows, columns=header)
    dataframe.columns = [str(column).strip() for column in dataframe.columns]
    return dataframe


def _split_markdown_row(line: str) -> list[str]:
    raw = line.strip().strip("|")
    return [cell.strip() for cell in raw.split("|")]


def _is_divider_row(row: Iterable[str]) -> bool:
    cells = [cell.strip() for cell in row]
    if not cells:
        return False
    return all(TABLE_DIVIDER_CELL_RE.match(cell) for cell in cells if cell)