from __future__ import annotations

from pathlib import Path

from docx import Document
from docx.enum.text import WD_COLOR_INDEX
from docx.shared import RGBColor

from review_agent.regulatory_info_package.schemas import MergedField


def write_docx_from_template(
    source_path: str | Path,
    output_path: str | Path,
    merged_fields: dict[str, MergedField],
) -> tuple[int, int, int]:
    source = Path(source_path)
    output = Path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    if source.exists():
        document = Document(source)
    else:
        document = Document()
    replacements = {f"{{{{{key}}}}}": field for key, field in merged_fields.items()}
    highlight_count = 0
    missing_count = 0
    llm_only_count = 0
    for paragraph in document.paragraphs:
        for placeholder, field in replacements.items():
            if placeholder in paragraph.text:
                _replace_paragraph_text(paragraph, paragraph.text.replace(placeholder, field.value), field)
                if field.highlight_reason != "none":
                    highlight_count += 1
                if field.highlight_reason == "missing":
                    missing_count += 1
                if field.highlight_reason == "llm_only":
                    llm_only_count += 1
    document.add_page_break()
    heading = document.add_paragraph()
    heading_run = heading.add_run("预生成字段")
    heading_run.bold = True
    table = document.add_table(rows=1, cols=4)
    table.rows[0].cells[0].text = "字段"
    table.rows[0].cells[1].text = "值"
    table.rows[0].cells[2].text = "来源"
    table.rows[0].cells[3].text = "待确认"
    for field in merged_fields.values():
        cells = table.add_row().cells
        cells[0].text = field.label
        cells[1].text = field.value
        cells[2].text = field.source
        cells[3].text = "是" if field.needs_review else "否"
        if field.highlight_reason != "none":
            highlight_count += 1
        if field.highlight_reason == "missing":
            missing_count += 1
        if field.highlight_reason == "llm_only":
            llm_only_count += 1
    document.save(output)
    return highlight_count, missing_count, llm_only_count


def _replace_paragraph_text(paragraph, text: str, field: MergedField) -> None:
    for run in paragraph.runs:
        run.text = ""
    run = paragraph.add_run(text)
    if field.highlight_reason != "none":
        run.font.highlight_color = WD_COLOR_INDEX.YELLOW
    if field.highlight_reason == "conflict":
        run.font.color.rgb = RGBColor(255, 0, 0)