import argparse
import re

LATEX_SPECIAL = {
    "\\": r"\textbackslash{}",
    "&": r"\&",
    "%": r"\%",
    "$": r"\$",
    "#": r"\#",
    "_": r"\_",
    "{": r"\{",
    "}": r"\}",
    "~": r"\textasciitilde{}",
    "^": r"\textasciicircum{}",
}


def escape_latex(text):
    return "".join(LATEX_SPECIAL.get(ch, ch) for ch in text)


def slugify(text):
    return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")



def parse_pipe_row(line):
    return [c.strip() for c in line.strip().strip("|").split("|")]


def is_separator_row(line):
    stripped = line.strip()
    return bool(stripped) and set(stripped) <= set("|-: ")


def parse_table(lines, i):
    headers = parse_pipe_row(lines[i])
    i += 1
    if i >= len(lines) or not is_separator_row(lines[i]):
        raise ValueError(f"expected a '|---|' separator row at line {i}, got {lines[i]!r}")
    i += 1
    rows = []
    while i < len(lines) and lines[i].strip().startswith("|"):
        rows.append(parse_pipe_row(lines[i]))
        i += 1
    return headers, rows, i



def convert(report_text):
    lines = report_text.splitlines()
    tables = []
    current_arch = None
    i = 0

    while i < len(lines):
        line = lines[i]

        if line.startswith("## Summary"):
            i += 1
            while i < len(lines) and not lines[i].strip().startswith("|"):
                i += 1
            headers, rows, i = parse_table(lines, i)
            tables.append(("Model comparison summary", "tab:summary", headers, rows))
            continue

        if line.startswith("### "):
            current_arch = line[len("### "):].strip()
            i += 1
            continue

        if line.startswith("**") and current_arch is not None:
            caption_line = line.strip()
            parts = caption_line.split("**")
            breakdown = parts[1] if len(parts) > 1 else "breakdown"
            caption = f"{current_arch} — {caption_line.replace('**', '')}"
            i += 1
            while i < len(lines) and lines[i].strip() == "":
                i += 1
            if i < len(lines) and lines[i].strip().startswith("|"):
                headers, rows, i = parse_table(lines, i)
                label = f"tab:{slugify(current_arch)}-{slugify(breakdown)}"
                tables.append((caption, label, headers, rows))
            continue

        i += 1

    return tables



def table_to_latex(caption, label, headers, rows):
    col_spec = "l" + "c" * (len(headers) - 1)
    out = [
        r"\begin{table}[htbp]",
        r"\centering",
        rf"\caption{{{escape_latex(caption)}}}",
        rf"\label{{{label}}}",
        rf"\begin{{tabular}}{{{col_spec}}}",
        r"\toprule",
        " & ".join(escape_latex(h) for h in headers) + r" \\",
        r"\midrule",
    ]
    for row in rows:
        out.append(" & ".join(escape_latex(c) for c in row) + r" \\")
    out += [r"\bottomrule", r"\end{tabular}", r"\end{table}"]
    return "\n".join(out)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--report", required=True, help="report.md written by evaluate_all_models.py")
    parser.add_argument("--output", required=True, help="path to write the combined .tex file")
    args = parser.parse_args()

    with open(args.report) as f:
        tables = convert(f.read())
    if not tables:
        raise SystemExit(f"No tables found in {args.report} -- "
                         "is this a report.md written by evaluate_all_models.py?")

    with open(args.output, "w") as f:
        f.write(f"% Generated by shared/report_to_latex.py from {args.report}\n"
                "% Requires \\usepackage{booktabs} in the preamble.\n\n")
        for caption, label, headers, rows in tables:
            f.write(table_to_latex(caption, label, headers, rows))
            f.write("\n\n")

    print(f"Wrote {len(tables)} table(s) to {args.output}")


if __name__ == "__main__":
    main()