"""
format_contract.py
------------------
Phục hồi Heading style cho file DOCX sau khi Google Drive OCR export.
Google Drive OCR hay xoá sạch heading styles VÀ bold formatting.

Logic detect heading (không cần bold):
  - ALL CAPS + < 120 ký tự           →  Heading 1
  - Regex: "ĐIỀU X." / "Chương X"    →  Heading 1
  - Regex: "X.Y." (sub-section)      →  Heading 2
  - Bold + Mixed case + < 120 ký tự  →  Heading 2 (fallback)
"""

import re
import shutil
from io import BytesIO
from pathlib import Path

from docx import Document
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from lxml import etree


# ───────────────────────────────────────────────
# Regex patterns cho cấu trúc pháp lý VN
# ───────────────────────────────────────────────

# Heading 1: "ĐIỀU 1.", "Điều 1.", "CHƯƠNG I", "Phần 1", etc.
H1_PATTERNS = re.compile(
    r"^(?:"
    r"(?:ĐIỀU|Điều)\s+\d+[\.\:\s]"       # Điều 1. / ĐIỀU 1:
    r"|(?:CHƯƠNG|Chương)\s+[IVXLCDM\d]+"  # Chương I / CHƯƠNG 2
    r"|(?:PHẦN|Phần)\s+[IVXLCDM\d]+"      # Phần I / PHẦN 2
    r"|(?:MỤC|Mục)\s+[IVXLCDM\d]+"        # Mục I / MỤC 2
    r")",
    re.IGNORECASE | re.UNICODE,
)

# Heading 2: "1.1.", "2.3.", "10.26." etc.
H2_PATTERNS = re.compile(
    r"^\d+\.\d+\.?\s",
    re.UNICODE,
)

# Heading 3: "1.1.1.", "2.3.4." etc.
H3_PATTERNS = re.compile(
    r"^\d+\.\d+\.\d+\.?\s",
    re.UNICODE,
)


# ───────────────────────────────────────────────
# Helpers
# ───────────────────────────────────────────────

def _is_all_caps(text: str) -> bool:
    """True nếu toàn bộ chữ cái trong text đều là hoa."""
    letters = [c for c in text if c.isalpha()]
    return len(letters) > 0 and all(c == c.upper() for c in letters)


def _para_is_bold(para) -> bool:
    """True nếu đoạn văn có ít nhất 1 run in đậm."""
    return any(run.bold for run in para.runs if run.text.strip())


def _ensure_heading_style(doc, style_id: str, style_name: str, outline_level: int, font_size: int):
    """Thêm heading style vào document nếu chưa có."""
    styles_el = doc.styles.element
    for s in styles_el.findall(qn("w:style")):
        if s.get(qn("w:styleId")) == style_id:
            return

    xml = f"""<w:style xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"
        w:type="paragraph" w:styleId="{style_id}">
      <w:name w:val="{style_name}"/>
      <w:basedOn w:val="Normal"/>
      <w:next w:val="Normal"/>
      <w:pPr>
        <w:outlineLvl w:val="{outline_level}"/>
        <w:spacing w:before="240" w:after="120"/>
      </w:pPr>
      <w:rPr>
        <w:b/>
        <w:sz w:val="{font_size}"/>
        <w:szCs w:val="{font_size}"/>
      </w:rPr>
    </w:style>"""
    styles_el.append(etree.fromstring(xml))


def _set_outline_level(para, level: int):
    """Gán outlineLvl cho paragraph, tuân thủ XML Schema của Word."""
    pPr = para._p.get_or_add_pPr()
    for existing in pPr.findall(qn("w:outlineLvl")):
        pPr.remove(existing)
    el = OxmlElement("w:outlineLvl")
    el.set(qn("w:val"), str(level))

    # w:outlineLvl PHẢI nằm trước w:rPr để tránh nhao chữ
    rPr = pPr.find(qn('w:rPr'))
    if rPr is not None:
        rPr.addprevious(el)
    else:
        sectPr = pPr.find(qn('w:sectPr'))
        if sectPr is not None:
            sectPr.addprevious(el)
        else:
            pPr.append(el)


def _get_style(doc, primary: str, fallback: str):
    """Lấy style từ document, thử primary trước rồi fallback."""
    try:
        return doc.styles[primary]
    except KeyError:
        try:
            return doc.styles[fallback]
        except KeyError:
            return None


# ───────────────────────────────────────────────
# Core detection logic
# ───────────────────────────────────────────────

def _detect_heading_level(text: str, is_bold: bool) -> int:
    """
    Detect heading level dựa trên nội dung text và bold.
    Returns: 0 = không phải heading, 1 = H1, 2 = H2, 3 = H3
    """
    if not text or len(text) >= 120:
        return 0

    # ── Ưu tiên 1: Regex cấu trúc pháp lý ──
    # H3 check trước H2 vì H3 pattern (1.1.1.) cũng match H2 (1.1.)
    if H3_PATTERNS.match(text):
        return 3

    if H1_PATTERNS.match(text):
        return 1

    if H2_PATTERNS.match(text):
        return 2

    # ── Ưu tiên 2: ALL CAPS = H1 (KHÔNG cần bold) ──
    # Google Drive OCR thường xoá bold nhưng giữ nguyên chữ hoa
    if _is_all_caps(text) and len(text) < 100:
        return 1

    # ── Ưu tiên 3: Bold + ngắn = H2 (fallback) ──
    word_count = len(text.split())
    if is_bold and word_count <= 15:
        return 2

    return 0


# ───────────────────────────────────────────────
# API: format từ bytes (dùng trong pipeline)
# ───────────────────────────────────────────────

def fix_headings_in_memory(docx_bytes: bytes) -> bytes:
    """
    Nhận DOCX bytes (từ Google Drive OCR export),
    phục hồi heading styles, trả về DOCX bytes đã fix.
    """
    doc = Document(BytesIO(docx_bytes))
    count = _apply_heading_fixes(doc)

    if count == 0:
        return docx_bytes

    buf = BytesIO()
    doc.save(buf)
    buf.seek(0)
    return buf.read()


# ───────────────────────────────────────────────
# API: format từ file path
# ───────────────────────────────────────────────

def format_contract(src_path: Path, dst_path: Path) -> int:
    """Format 1 file DOCX trên disk. Trả về số heading đã fix."""
    shutil.copy2(src_path, dst_path)
    doc = Document(str(dst_path))
    count = _apply_heading_fixes(doc)
    doc.save(str(dst_path))
    return count


# ───────────────────────────────────────────────
# Internal: áp dụng fix
# ───────────────────────────────────────────────

def _apply_heading_fixes(doc) -> int:
    """Quét paragraphs và gán heading style. Trả về số heading đã fix."""
    _ensure_heading_style(doc, "Heading1", "heading 1", 0, 28)
    _ensure_heading_style(doc, "Heading2", "heading 2", 1, 24)
    _ensure_heading_style(doc, "Heading3", "heading 3", 2, 22)

    h1_style = _get_style(doc, "Heading 1", "Heading1")
    h2_style = _get_style(doc, "Heading 2", "Heading2")
    h3_style = _get_style(doc, "Heading 3", "Heading3")

    if not h1_style or not h2_style:
        return 0

    style_map = {
        1: (h1_style, 0),
        2: (h2_style, 1),
        3: (h3_style or h2_style, 2),
    }

    count = 0
    for para in doc.paragraphs:
        text = para.text.strip()
        if not text:
            continue

        is_bold = _para_is_bold(para)
        level = _detect_heading_level(text, is_bold)

        if level == 0:
            continue

        # Luôn override heading style + outlineLvl dựa trên nội dung text
        # (Google Drive hay gán sai level, ví dụ tất cả thành Heading 2)
        heading_style, outline_val = style_map[level]
        para.style = heading_style
        _set_outline_level(para, outline_val)
        for run in para.runs:
            run.bold = True
        count += 1

    return count
