# agent/markdown_utils.py
import re
from typing import List, Dict, Optional


def extract_section(md_text: str, section_title: str) -> str:
    """Extract content of a ## or ### section from markdown text.
    Stops at the next heading of equal or higher level (fewer or equal #'s).
    For ## sections, includes ### sub-sections.
    """
    heading_re = rf'^(#{{2,3}})\s+{re.escape(section_title)}[^\S\n]*$'
    h_match = re.search(heading_re, md_text, re.MULTILINE)
    if not h_match:
        return ''

    level = len(h_match.group(1))
    start = h_match.end()

    next_re = rf'^#{{1,{level}}}\s'
    n_match = re.search(next_re, md_text[start:], re.MULTILINE)
    end = start + n_match.start() if n_match else len(md_text)

    return md_text[start:end]


def parse_table_rows(text: str) -> List[Dict[str, str]]:
    """Parse a markdown table from text, return list of row dicts.
    Handles tables with exactly one header row and one separator row.
    """
    lines = []
    for line in text.split('\n'):
        stripped = line.strip()
        if stripped.startswith('|') and stripped.endswith('|'):
            if re.match(r'^\|[\s\-:]+\|', stripped):
                continue
            lines.append(stripped)

    if len(lines) < 1:
        return []

    headers = [cell.strip() for cell in lines[0].split('|')[1:-1]]
    rows = []
    for line in lines[1:]:
        cells = [cell.strip() for cell in line.split('|')[1:-1]]
        if len(cells) == len(headers):
            rows.append(dict(zip(headers, cells)))
    return rows


def parse_table_from_section(md_text: str, section_title: str) -> List[Dict[str, str]]:
    """Find a ### section and parse its first table."""
    section_text = extract_section(md_text, section_title)
    if not section_text:
        return []
    return parse_table_rows(section_text)


def find_row_by_key(rows: List[Dict[str, str]], key_col: str, key_value: str) -> Optional[Dict[str, str]]:
    """Find a table row where a specific column matches the key value."""
    for row in rows:
        if row.get(key_col, '').strip() == key_value:
            return row
    return None