from __future__ import annotations

import re
import zipfile
from pathlib import Path
from xml.etree import ElementTree as ET

from revise_report_docx_html import build_bab_4_5, build_state_of_art_note


INPUT_DOCX = Path("/Users/user/Downloads/belum fix.docx")
OUTPUT_DOCX = Path("Laporan_ConfiVoice_Revisi_Bab4_Bab5_FINAL.docx")
WORK_DIR = Path("docx_work")
XML_PATH = "word/document.xml"

W_NS = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
NS = {"w": W_NS}
ET.register_namespace("w", W_NS)


def qn(tag: str) -> str:
    prefix, local = tag.split(":", 1)
    if prefix != "w":
        raise ValueError(tag)
    return f"{{{W_NS}}}{local}"


def text_content(element: ET.Element) -> str:
    return "".join(element.itertext())


def paragraph_text(paragraph: ET.Element) -> str:
    return "".join(t.text or "" for t in paragraph.findall(".//w:t", NS))


def clear_paragraph_runs(paragraph: ET.Element) -> None:
    for child in list(paragraph):
        if child.tag != qn("w:pPr"):
            paragraph.remove(child)


def append_run(paragraph: ET.Element, text: str, bold: bool = False) -> None:
    run = ET.SubElement(paragraph, qn("w:r"))
    if bold:
        rpr = ET.SubElement(run, qn("w:rPr"))
        ET.SubElement(rpr, qn("w:b"))
    t = ET.SubElement(run, qn("w:t"))
    t.set("{http://www.w3.org/XML/1998/namespace}space", "preserve")
    t.text = text


def set_paragraph_text(paragraph: ET.Element, text: str, bold: bool = False) -> None:
    ppr = paragraph.find("w:pPr", NS)
    clear_paragraph_runs(paragraph)
    if ppr is None:
        ppr = ET.Element(qn("w:pPr"))
        paragraph.insert(0, ppr)
    append_run(paragraph, text, bold=bold)


def make_paragraph(text: str, bold: bool = False, center: bool = False) -> ET.Element:
    p = ET.Element(qn("w:p"))
    ppr = ET.SubElement(p, qn("w:pPr"))
    spacing = ET.SubElement(ppr, qn("w:spacing"))
    spacing.set(qn("w:after"), "120")
    if center:
        jc = ET.SubElement(ppr, qn("w:jc"))
        jc.set(qn("w:val"), "center")
    append_run(p, text, bold=bold)
    return p


def make_table(rows: list[list[str]]) -> ET.Element:
    tbl = ET.Element(qn("w:tbl"))
    tbl_pr = ET.SubElement(tbl, qn("w:tblPr"))
    borders = ET.SubElement(tbl_pr, qn("w:tblBorders"))
    for side in ["top", "left", "bottom", "right", "insideH", "insideV"]:
        border = ET.SubElement(borders, qn(f"w:{side}"))
        border.set(qn("w:val"), "single")
        border.set(qn("w:sz"), "4")
        border.set(qn("w:space"), "0")
        border.set(qn("w:color"), "000000")
    width = ET.SubElement(tbl_pr, qn("w:tblW"))
    width.set(qn("w:w"), "9360")
    width.set(qn("w:type"), "dxa")

    col_count = max(len(row) for row in rows) if rows else 1
    col_width = str(max(1200, 9360 // col_count))
    grid = ET.SubElement(tbl, qn("w:tblGrid"))
    for _ in range(col_count):
        col = ET.SubElement(grid, qn("w:gridCol"))
        col.set(qn("w:w"), col_width)

    for row_index, row in enumerate(rows):
        tr = ET.SubElement(tbl, qn("w:tr"))
        for item in row:
            tc = ET.SubElement(tr, qn("w:tc"))
            tc_pr = ET.SubElement(tc, qn("w:tcPr"))
            tcw = ET.SubElement(tc_pr, qn("w:tcW"))
            tcw.set(qn("w:w"), col_width)
            tcw.set(qn("w:type"), "dxa")
            tc.append(make_paragraph(item, bold=row_index == 0))
    return tbl


def html_fragment_to_ooxml(fragment: str) -> list[ET.Element]:
    cleaned = re.sub(r"<br\s*/?>", "", fragment)
    root = ET.fromstring(f"<root>{cleaned}</root>")
    output: list[ET.Element] = []

    for child in list(root):
        if child.tag == "p":
            cls = child.attrib.get("class", "")
            txt = text_content(child).strip()
            if not txt:
                output.append(make_paragraph(""))
                continue
            is_heading = cls == "p7" or txt.startswith("BAB ")
            is_subheading = bool(child.find(".//b") is not None) and not is_heading
            output.append(make_paragraph(txt, bold=is_heading or is_subheading, center=is_heading))
        elif child.tag == "ul":
            for li in child.iter("li"):
                txt = text_content(li).strip()
                if txt:
                    output.append(make_paragraph(f"- {txt}"))
        elif child.tag == "table":
            rows: list[list[str]] = []
            for tr in child.iter("tr"):
                cells = [text_content(td).strip() for td in tr.iter("td")]
                if cells:
                    rows.append(cells)
            if rows:
                output.append(make_table(rows))
                output.append(make_paragraph(""))
    return output


def patch_document_xml(xml_bytes: bytes) -> bytes:
    root = ET.fromstring(xml_bytes)
    body = root.find(".//w:body", NS)
    if body is None:
        raise RuntimeError("word/body tidak ditemukan")

    for text_node in root.findall(".//w:t", NS):
        if text_node.text == "PROPOSAL TUGAS AKHIR":
            text_node.text = "LAPORAN TUGAS AKHIR"

    for paragraph in body.findall("w:p", NS):
        txt = paragraph_text(paragraph)
        if txt.startswith("Disini Peneliti akan mengimplementasikan"):
            set_paragraph_text(
                paragraph,
                "Pada tahap ini, peneliti mengimplementasikan model aplikasi dengan membangun aplikasi mobile "
                "menggunakan framework Flutter. Aplikasi ini digunakan untuk merekam atau memilih audio, "
                "mengirimkan audio ke backend, serta menampilkan hasil klasifikasi tingkat percaya diri "
                "berdasarkan model machine learning yang telah dibuat. Hasil prediksi ditampilkan setelah "
                "audio selesai dikirim dan diproses oleh sistem.",
            )
        elif txt.startswith("Flowchart sistem ini dimulai dengan Button Lanjut"):
            set_paragraph_text(
                paragraph,
                "Flowchart sistem dimulai dari pengguna membuka aplikasi, kemudian melakukan login atau registrasi. "
                "Setelah berhasil masuk, pengguna diarahkan ke halaman utama. Pada halaman utama, pengguna dapat "
                "memasukkan nama siswa, merekam audio secara langsung atau memilih file audio yang sudah tersedia. "
                "Setelah audio siap, pengguna menekan tombol Prediksi untuk mengirimkan audio ke backend. Backend "
                "memproses audio menggunakan model machine learning dan mengembalikan hasil klasifikasi berupa "
                "Percaya Diri atau Tidak Percaya Diri beserta persentase probabilitasnya. Setelah hasil prediksi "
                "ditampilkan, pengguna dapat menyimpan hasil analisis ke database atau melihat riwayat analisis "
                "yang telah tersimpan. Flowchart dari aplikasi untuk pengguna dapat dilihat melalui Gambar 3.3.",
            )

    body_children = list(body)
    daftar_pustaka_indexes = [
        index
        for index, child in enumerate(body_children)
        if child.tag == qn("w:p") and paragraph_text(child).strip() == "DAFTAR PUSTAKA"
    ]
    if not daftar_pustaka_indexes:
        raise RuntimeError("Paragraf DAFTAR PUSTAKA tidak ditemukan")
    insert_index = daftar_pustaka_indexes[-1]

    addition = "\n".join([build_bab_4_5(), build_state_of_art_note()])
    elements = html_fragment_to_ooxml(addition)
    for offset, element in enumerate(elements):
        body.insert(insert_index + offset, element)

    # Refresh list after insertion, then add extra bibliography item after Zalukhu.
    for index, child in enumerate(list(body)):
        if child.tag == qn("w:p") and paragraph_text(child).startswith("Zalukhu, A., Purba"):
            body.insert(
                index + 1,
                make_paragraph(
                    "Jain, M., Narayan, S., Balaji, P., Bharath, K. P., Bhowmick, A., Karthik, R., "
                    "& Muthu, R. K. (2020). Speech Emotion Recognition using Support Vector Machine. "
                    "arXiv. https://arxiv.org/abs/2002.07590"
                ),
            )
            break

    return ET.tostring(root, encoding="utf-8", xml_declaration=True)


def main() -> None:
    OUTPUT_DOCX.unlink(missing_ok=True)
    with zipfile.ZipFile(INPUT_DOCX, "r") as zin, zipfile.ZipFile(OUTPUT_DOCX, "w", zipfile.ZIP_DEFLATED) as zout:
        for item in zin.infolist():
            data = zin.read(item.filename)
            if item.filename == XML_PATH:
                data = patch_document_xml(data)
            zout.writestr(item, data)


if __name__ == "__main__":
    main()