#!/usr/bin/env python3
import argparse
import json
from pathlib import Path

from docx import Document
from docx.enum.text import WD_BREAK, WD_LINE_SPACING
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.shared import Pt


DEFAULT_FONT_SIZE_PT = 10.5
DEFAULT_LINE_SPACING_PT = 13
DEFAULT_VISUAL_WIDTH = 88


def clear_body(document):
    body = document._body._element
    sect_pr = body.sectPr
    for child in list(body):
        if child is not sect_pr:
            body.remove(child)


def strip_line_comment(work):
    quote = ""
    escape = False
    for index in range(len(work) - 1):
        char = work[index]
        if escape:
            escape = False
            continue
        if quote:
            if char == "\\":
                escape = True
            elif char == quote:
                quote = ""
            continue
        if char in ("'", '"', "`"):
            quote = char
            continue
        if char == "/" and work[index + 1] == "/":
            return work[:index].rstrip()
    return work


def strip_comments_and_blanks(raw_lines):
    cleaned = []
    in_block = False
    in_html_block = False
    for raw in raw_lines:
        work = raw.rstrip()
        while True:
            if in_html_block:
                end = work.find("-->")
                if end == -1:
                    work = ""
                    break
                work = work[end + 3 :]
                in_html_block = False
                continue
            html_start = work.find("<!--")
            if html_start != -1:
                html_end = work.find("-->", html_start + 4)
                if html_end == -1:
                    work = work[:html_start].rstrip()
                    in_html_block = True
                    break
                work = (work[:html_start] + work[html_end + 3 :]).rstrip()
                continue
            if in_block:
                end = work.find("*/")
                if end == -1:
                    work = ""
                    break
                work = work[end + 2 :]
                in_block = False
                continue
            start = work.find("/*")
            if start != -1:
                end = work.find("*/", start + 2)
                if end == -1:
                    work = work[:start].rstrip()
                    in_block = True
                    break
                work = (work[:start] + work[end + 2 :]).rstrip()
                continue
            break
        stripped = work.strip()
        if not stripped:
            continue
        if stripped.startswith(("//", "#", "*")):
            continue
        work = strip_line_comment(work)
        if work.strip():
            cleaned.append(work.rstrip())
    return cleaned


def read_source_lines(path):
    return strip_comments_and_blanks(Path(path).read_text(encoding="utf-8", errors="replace").splitlines())


def read_source_range(file_path, start, end):
    raw_lines = Path(file_path).read_text(encoding="utf-8", errors="replace").splitlines()
    return strip_comments_and_blanks(raw_lines[int(start) - 1:int(end)])


def visual_rows(line, visual_width):
    length = len(line.expandtabs(4))
    if length <= visual_width:
        return 1
    return (length + visual_width - 1) // visual_width


def collect_group_lines(group):
    result = []
    group_name = group.get("name", "")
    for file_path in group.get("files", []):
        for line in read_source_lines(file_path):
            result.append((group_name, str(file_path), line))
    for item in group.get("ranges", []):
        file_path = item["file"]
        for line in read_source_range(file_path, item["start"], item["end"]):
            result.append((group_name, str(file_path), line))
    return result


def paginate_group(group, visual_width, capacity):
    source_lines = collect_group_lines(group)
    page_count = int(group["pages"])
    pages = []
    cursor = 0
    for page_index in range(page_count):
        page = []
        used = 0
        while cursor < len(source_lines):
            item = source_lines[cursor]
            rows = visual_rows(item[2], visual_width)
            if page and used + rows > capacity:
                break
            page.append(item)
            used += max(rows, 1)
            cursor += 1
        if not page:
            raise RuntimeError(f"group {group.get('name')} page {page_index + 1} is empty")
        pages.append(page)
    return pages


def paginate_lines(source_lines, page_count, visual_width, capacity):
    pages = []
    cursor = 0
    for page_index in range(page_count):
        page = []
        used = 0
        while cursor < len(source_lines):
            item = source_lines[cursor]
            rows = visual_rows(item[2], visual_width)
            if page and used + rows > capacity:
                break
            page.append(item)
            used += max(rows, 1)
            cursor += 1
        if not page:
            raise RuntimeError(f"page {page_index + 1} is empty")
        pages.append(page)
    return pages


def final_page_from_file(file_path, name="final"):
    return [(name, str(file_path), line) for line in read_source_lines(file_path)]


def final_page_from_range(spec, name="final"):
    file_path = spec["file"]
    start = int(spec.get("start", 1))
    end = int(spec.get("end", start))
    raw_lines = Path(file_path).read_text(encoding="utf-8", errors="replace").splitlines()
    return [(name, str(file_path), line) for line in strip_comments_and_blanks(raw_lines[start - 1:end])]


def final_page_from_method(spec, name="final"):
    file_path = Path(spec["file"])
    marker = spec["marker"]
    raw_lines = file_path.read_text(encoding="utf-8", errors="replace").splitlines()
    start_index = next((idx for idx, line in enumerate(raw_lines) if marker in line), None)
    if start_index is None:
        raise RuntimeError(f"final_page_method marker not found: {marker}")
    selected = []
    brace_depth = 0
    seen_open = False
    for line in raw_lines[start_index:]:
        selected.append(line)
        code = line.split("//", 1)[0]
        brace_depth += code.count("{")
        if "{" in code:
            seen_open = True
        brace_depth -= code.count("}")
        if seen_open and brace_depth <= 0:
            break
    return [(name, str(file_path), line) for line in strip_comments_and_blanks(selected)]


def add_code_paragraph(document, line, font_size_pt, line_spacing_pt):
    para = document.add_paragraph()
    para.paragraph_format.space_before = Pt(0)
    para.paragraph_format.space_after = Pt(0)
    para.paragraph_format.line_spacing_rule = WD_LINE_SPACING.EXACTLY
    para.paragraph_format.line_spacing = Pt(line_spacing_pt)
    p_pr = para._p.get_or_add_pPr()
    snap = OxmlElement("w:snapToGrid")
    snap.set(qn("w:val"), "0")
    p_pr.append(snap)
    run = para.add_run(line.expandtabs(4))
    run.font.size = Pt(font_size_pt)
    return para


def build_docx(template, output, pages, font_size_pt, line_spacing_pt):
    document = Document(template)
    clear_body(document)
    for section in document.sections:
        for paragraph in section.header.paragraphs:
            paragraph.clear()
    style = document.styles["Normal"]
    style.font.size = Pt(font_size_pt)
    style.paragraph_format.space_before = Pt(0)
    style.paragraph_format.space_after = Pt(0)
    style.paragraph_format.line_spacing_rule = WD_LINE_SPACING.EXACTLY
    style.paragraph_format.line_spacing = Pt(line_spacing_pt)
    for page_no, page in enumerate(pages, 1):
        for _, _, line in page:
            add_code_paragraph(document, line, font_size_pt, line_spacing_pt)
        if page_no < len(pages):
            document.paragraphs[-1].add_run().add_break(WD_BREAK.PAGE)
    document.save(output)


def validate_docx(path, expected_pages):
    from zipfile import ZipFile
    from xml.etree import ElementTree as ET

    ns = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
    path = Path(path)
    with ZipFile(path) as archive:
        doc = ET.fromstring(archive.read("word/document.xml"))
        texts = []
        empty_body_paras = 0
        for paragraph in doc.findall(".//w:body/w:p", ns):
            text = "".join(node.text or "" for node in paragraph.findall(".//w:t", ns))
            has_page_break = bool(paragraph.findall('.//w:br[@w:type="page"]', ns))
            if text:
                texts.append(text)
            elif not has_page_break:
                empty_body_paras += 1
        page_breaks = len(doc.findall('.//w:br[@w:type="page"]', ns))
        header_text = ""
        if "word/header1.xml" in archive.namelist():
            header = ET.fromstring(archive.read("word/header1.xml"))
            header_text = "".join(node.text or "" for node in header.findall(".//w:t", ns))
    comment_hits = sum(
        1
        for text in texts
        if text.lstrip().startswith(("//", "/*", "*", "#"))
        or " //" in text
        or "/*" in text
        or "*/" in text
    )
    prefix_hits = sum(1 for text in texts if " | " in text and ("/" in text or "\\" in text))
    return {
        "page_breaks": page_breaks,
        "expected_page_breaks": max(expected_pages - 1, 0),
        "body_lines": len(texts),
        "empty_body_paras": empty_body_paras,
        "header_text": header_text,
        "comment_hits": comment_hits,
        "prefix_hits": prefix_hits,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--template", required=True)
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--pages", type=int, default=30)
    parser.add_argument("--font-size", type=float, default=DEFAULT_FONT_SIZE_PT)
    parser.add_argument("--line-spacing", type=float, default=DEFAULT_LINE_SPACING_PT)
    parser.add_argument("--visual-width", type=int, default=DEFAULT_VISUAL_WIDTH)
    parser.add_argument("--capacity", type=int, default=50)
    args = parser.parse_args()

    manifest = json.loads(Path(args.manifest).read_text(encoding="utf-8"))
    groups = manifest["groups"]
    if sum(int(group["pages"]) for group in groups) != args.pages:
        raise RuntimeError("sum(groups[].pages) must equal --pages")
    if manifest.get("continuous_pages"):
        source_lines = []
        for group in groups:
            source_lines.extend(collect_group_lines(group))
        pages = paginate_lines(source_lines, args.pages, args.visual_width, args.capacity)
    else:
        pages = []
        for group in groups:
            pages.extend(paginate_group(group, args.visual_width, args.capacity))
    if manifest.get("final_page_method"):
        pages[-1] = final_page_from_method(manifest["final_page_method"], groups[-1].get("name", "final"))
    elif manifest.get("final_page_range"):
        pages[-1] = final_page_from_range(manifest["final_page_range"], groups[-1].get("name", "final"))
    elif manifest.get("final_page_file"):
        pages[-1] = final_page_from_file(manifest["final_page_file"], groups[-1].get("name", "final"))
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    build_docx(args.template, args.output, pages, args.font_size, args.line_spacing)
    report = validate_docx(args.output, args.pages)
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()