"""Extract tables from PDFs using pdfplumber."""

import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import pdfplumber

from output_utils import output_to_file_or_stdout


def extract(pdf_path: str) -> list[dict]:
    """Return list of table dicts with page number, table index, and rows."""
    results = []
    with pdfplumber.open(pdf_path) as pdf:
        for i, page in enumerate(pdf.pages, start=1):
            tables = page.find_tables()
            for t_idx, table in enumerate(tables):
                rows = table.extract()
                if rows:
                    results.append({
                        "page": i,
                        "table": t_idx + 1,
                        "header": rows[0],
                        "rows": rows[1:],
                    })
    return results


def _tables_to_csv(tables: list[dict]) -> list[list]:
    """Flatten all tables into a single CSV (with a source column)."""
    csv_rows = [["page", "table"] + ["col"]]  # placeholder header
    for t in tables:
        header = t["header"] or []
        # rebuild header row with source columns
        full_header = ["page", "table"] + [str(h) if h else f"col_{j}" for j, h in enumerate(header)]
        if csv_rows[0] == ["page", "table", "col"]:
            csv_rows[0] = full_header
        for row in t["rows"]:
            csv_rows.append([t["page"], t["table"]] + [str(c) if c else "" for c in row])
    return csv_rows


def _tables_to_text(tables: list[dict]) -> list[dict]:
    """Convert tables into page-text dicts for txt/md output."""
    page_texts: dict[int, list[str]] = {}
    for t in tables:
        lines = [f"Table {t['table']}:"]
        header = t["header"] or []
        lines.append("  " + " | ".join(str(h) if h else "" for h in header))
        lines.append("  " + "-" * 40)
        for row in t["rows"]:
            lines.append("  " + " | ".join(str(c) if c else "" for c in row))
        page_texts.setdefault(t["page"], []).append("\n".join(lines))
    return [{"page": p, "text": "\n\n".join(texts)} for p, texts in sorted(page_texts.items())]


def main():
    parser = argparse.ArgumentParser(description="Extract tables from PDFs")
    parser.add_argument("pdf", help="Path to input PDF")
    parser.add_argument("output", nargs="?", default=None, help="Output file path")
    parser.add_argument("--format", choices=["json", "csv", "txt", "md"], default="json",
                        dest="fmt", help="Output format (default: json)")
    args = parser.parse_args()

    tables = extract(args.pdf)

    if args.fmt == "json":
        data = {"extractor": "tables", "tables": tables}
        output_to_file_or_stdout(data, "json", args.output)
    elif args.fmt == "csv":
        csv_rows = _tables_to_csv(tables)
        output_to_file_or_stdout(csv_rows, "csv", args.output)
    else:
        text_pages = _tables_to_text(tables)
        output_to_file_or_stdout(text_pages, args.fmt, args.output)


if __name__ == "__main__":
    main()
