#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# ///
"""
Markdown Table Justifier

ASCII character-justifies markdown tables by padding all cells
to match the maximum width in each column.

Usage:
    ./table_justify.py <file>
    ./table_justify.py <file> -o <output>
    cat <file> | table-justify
"""

import argparse
import re
import sys
from pathlib import Path


def parse_table(lines: list[str]) -> tuple[list[list[str]], list[int]]:
    """
    Parse a markdown table into rows and calculate column widths.
    Returns (rows, max_widths).
    """
    rows = []
    max_widths = []

    for line in lines:
        line = line.rstrip()
        if not line.startswith("|"):
            continue

        # Split by | and strip whitespace
        cells = [cell.strip() for cell in line.split("|")]
        # Remove empty first/last cells from leading/trailing |
        cells = [c for c in cells if c or c == ""]
        if len(cells) > 0 and cells[0] == "":
            cells = cells[1:]
        if len(cells) > 0 and cells[-1] == "":
            cells = cells[:-1]

        if not cells:
            continue

        rows.append(cells)

        # Update max widths
        while len(max_widths) < len(cells):
            max_widths.append(0)

        for i, cell in enumerate(cells):
            max_widths[i] = max(max_widths[i], len(cell))

    return rows, max_widths


def is_separator_row(row: list[str]) -> bool:
    """Check if a row is a header separator (all dashes)."""
    if not row:
        return False
    return all(re.match(r"^-+$", cell.strip().replace(" ", "")) or cell.strip() == "" for cell in row)


def format_separator(widths: list[int]) -> str:
    """Format the separator row."""
    cells = ["-" * w for w in widths]
    return "|" + "|".join(f" {c} " for c in cells) + "|"


def format_row(row: list[str], widths: list[int]) -> str:
    """Format a data row with proper padding."""
    padded = []
    for i, cell in enumerate(row):
        if i < len(widths):
            padded.append(cell.ljust(widths[i]))
        else:
            padded.append(cell)
    return "|" + "|".join(f" {c} " for c in padded) + "|"


def justify_table(lines: list[str]) -> list[str]:
    """Justify a markdown table."""
    rows, max_widths = parse_table(lines)

    if not rows:
        return lines

    result = []
    for i, row in enumerate(rows):
        if i == 1 and is_separator_row(row):
            result.append(format_separator(max_widths))
        else:
            result.append(format_row(row, max_widths))

    return result


def process_content(content: str) -> str:
    """Process content and justify all tables found."""
    lines = content.split("\n")
    result = []
    table_lines = []
    in_table = False

    for line in lines:
        if line.strip().startswith("|"):
            table_lines.append(line)
            in_table = True
        else:
            if in_table:
                # End of table, process it
                result.extend(justify_table(table_lines))
                table_lines = []
                in_table = False
            result.append(line)

    # Process last table if file ends with one
    if table_lines:
        result.extend(justify_table(table_lines))

    return "\n".join(result)


def main():
    parser = argparse.ArgumentParser(
        description="ASCII character-justify markdown tables"
    )
    parser.add_argument("input", nargs="?", help="Input file (default: stdin)")
    parser.add_argument("-o", "--output", help="Output file (default: stdout)")
    args = parser.parse_args()

    # Read input
    if args.input:
        content = Path(args.input).read_text()
    else:
        content = sys.stdin.read()

    # Process
    result = process_content(content)

    # Write output
    if args.output:
        Path(args.output).write_text(result)
    else:
        print(result)


if __name__ == "__main__":
    main()