"""Generate Jupyter Notebook (.ipynb) from proof.py."""

import re

SECTION_MARKER_RE = re.compile(
    r"^# [-=]{3,}\n# \d+\.\s+(.+)\n# [-=]{3,}$",
    re.MULTILINE,
)
SECTION_PLAIN_RE = re.compile(r"^# (\d+\.\s+.+)$", re.MULTILINE)


def generate_notebook(proof_py_text: str, proof_data: dict, slug: str, canonical_url: str) -> dict:
    verdict = proof_data.get("verdict", {})
    verdict_str = verdict.get("value", "") if isinstance(verdict, dict) else verdict
    claim = proof_data.get("claim_natural", slug)
    date = proof_data.get("generator", {}).get("generated_at", "")
    claim_type = proof_data.get("claim_formal", {}).get("claim_type")
    is_theorem = claim_type == "theorem"

    if is_theorem:
        intro = (
            "This notebook re-runs the implementation regression checks. "
            "The verdict above is established by the deductive argument in proof.md, "
            "not by re-executing this notebook."
        )
    else:
        intro = (
            "This notebook is an interactive version of the proof script. "
            "Run each cell to re-verify the claim."
        )

    cells = []
    cells.append(_markdown_cell(
        f"# Proof: {claim}\n\n"
        f"**Verdict:** {verdict_str}  \n**Date:** {date}  \n"
        f"**Source:** [{canonical_url}]({canonical_url})\n\n---\n\n"
        f"{intro}"
    ))

    sections = _split_on_markers(proof_py_text)
    if len(sections) == 1:
        cells.append(_code_cell(sections[0]["code"]))
    else:
        for section in sections:
            if section["heading"]:
                cells.append(_markdown_cell(f"## {section['heading']}"))
            cells.append(_code_cell(section["code"]))

    cells.append(_markdown_cell(
        f"---\n\n*Generated by [Proof Engine](https://github.com/yaniv-golan/proof-engine). "
        f"View the [full proof page]({canonical_url}).*"
    ))

    return {
        "nbformat": 4, "nbformat_minor": 5,
        "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
                      "language_info": {"name": "python", "version": "3.11.0"}},
        "cells": cells,
    }


def _split_on_markers(text: str) -> list[dict]:
    matches = list(SECTION_MARKER_RE.finditer(text))
    if not matches:
        matches = list(SECTION_PLAIN_RE.finditer(text))
    if not matches:
        return [{"heading": None, "code": text.strip()}]

    sections = []
    preamble = text[: matches[0].start()].strip()
    if preamble:
        sections.append({"heading": "Setup", "code": preamble})
    for i, match in enumerate(matches):
        heading = match.group(1).strip()
        start = match.end()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
        code = text[start:end].strip()
        if code:
            sections.append({"heading": heading, "code": code})
    return sections


def _markdown_cell(source: str) -> dict:
    return {"cell_type": "markdown", "metadata": {}, "source": source}


def _code_cell(source: str) -> dict:
    return {"cell_type": "code", "metadata": {}, "source": source, "outputs": [], "execution_count": None}
