#!/usr/bin/env python3
"""
差旅报销自动化脚本

Python 只做两类确定性工作：
1. extract：整理输入附件，提取 PDF 原文到本次归档目录的 TXT
2. fill：读取大模型生成的 JSON，填入 Excel 模板，并完成本次归档目录命名

中间的发票理解、金额分类、日期/事由/人员判断全部由大模型完成。
模板使用 .xlsx，完全依赖 Python 包实现跨平台，不依赖 LibreOffice。
"""

import argparse
import json
import re
import shutil
import zipfile
from datetime import date
from pathlib import Path
from tempfile import NamedTemporaryFile
from xml.etree import ElementTree as ET

import pdfplumber


OUTPUT_DIR_NAME = "报销输出"
RAW_TEXT_FILENAME = "pdf_raw_text.txt"
ANALYZED_RESULT_FILENAME = "analyzed_result.json"
TEMPLATE_FILE = Path("差旅费明细表.xlsx")
ATTACHMENT_SUFFIXES = {".pdf", ".jpg", ".jpeg", ".png"}


SCHEMA_EXAMPLE = {
    "meta": {
        "出差人": "从文本中提取或询问",
        "出差事由": "从文本中提取或询问",
        "出发日期": "从文本中提取或询问",
        "结束日期": "从文本中提取或询问",
        "出发地": "从文本中提取或询问",
        "目的地": "从文本中提取或询问",
        "出行方式": "从文本中提取或询问",
        "出差天数": "从文本中提取或询问",
        "事由简写": "从文本中提取或询问，适合文件名使用的简短事由描述"
    },
    "expenses": {
        "住宿": {
            "金额": 510.0,
            "单据张数": 1,
            "说明": "广州市番禺明乐宾馆住宿费"
        },
        "交通": {
            "金额": 2013.77,
            "单据张数": 5,
            "说明": "机票1690.00 + 滴滴323.77"
        },
        "饮食": {
            "金额": 0.0,
            "单据张数": 0,
            "说明": "未发现餐饮发票"
        },
        "其他": {
            "金额": 0.0,
            "单据张数": 0,
            "说明": "无其他费用"
        }
    },
    "summary": {
        "附件总张数": 6,
        "报销总额": 2523.77,
        "预借差旅费": 0.0,
        "补领金额": 2523.77,
        "退还金额": 0.0
    },
    "evidence": [
        {
            "文件": "住宿.pdf",
            "类别": "住宿",
            "金额": 510.0,
            "依据": "价税合计（小写）¥510.00"
        },
        {
            "文件": "机票预订出行订单尾号87377629-电子普通发票.pdf",
            "类别": "交通",
            "金额": 1690.0,
            "依据": "价税合计（小写）¥1690.00"
        }
    ],
    "questions": []
}


def command_extract(_args):
    archive_dir = prepare_archive_dir()
    output_dir = archive_dir / OUTPUT_DIR_NAME
    output_dir.mkdir(exist_ok=True)
    raw_text_file = output_dir / RAW_TEXT_FILENAME

    pdf_files = sorted(archive_dir.glob("*.pdf"))
    if not pdf_files:
        raise SystemExit(f"{archive_dir} 中未找到 PDF 文件。请把 PDF 放到根目录后重新运行 extract。")

    with open(raw_text_file, "w", encoding="utf-8") as output:
        output.write("=" * 60 + "\n")
        output.write("PDF 原始内容汇总\n")
        output.write("=" * 60 + "\n\n")

        for pdf_path in pdf_files:
            output.write(f"\n{'=' * 60}\n")
            output.write(f"文件: {pdf_path}\n")
            output.write("=" * 60 + "\n")

            try:
                with pdfplumber.open(pdf_path) as pdf:
                    for page_index, page in enumerate(pdf.pages, start=1):
                        text = page.extract_text() or ""
                        output.write(f"\n--- 第 {page_index} 页 ---\n")
                        output.write(text)
                        output.write("\n")
            except Exception as exc:
                output.write(f"\n[读取失败: {exc}]\n")

    print(f"本次附件目录：{archive_dir}")
    print(f"已提取 {len(pdf_files)} 个 PDF 文件到 {raw_text_file}")
    print(f"请将分析结果保存为：{output_dir / ANALYZED_RESULT_FILENAME}")


def prepare_archive_dir():
    root_attachments = find_root_attachments()
    today_dir = Path(date.today().strftime("%Y%m%d"))

    if root_attachments:
        archive_dir = choose_archive_dir(today_dir)
        archive_dir.mkdir(exist_ok=True)
        for attachment in root_attachments:
            target = unique_file_path(archive_dir / attachment.name)
            shutil.move(str(attachment), target)
        print(f"已移动 {len(root_attachments)} 个输入文件到 {archive_dir}")
        return archive_dir

    existing_today_dir = find_existing_today_dir(today_dir)
    if existing_today_dir:
        return existing_today_dir

    raise SystemExit("当前根目录未找到 PDF/JPG/JPEG/PNG 文件。请把本次发票、行程单或截图放到根目录后再运行 extract。")


def find_root_attachments():
    return sorted(
        path for path in Path(".").iterdir()
        if path.is_file() and path.suffix.lower() in ATTACHMENT_SUFFIXES
    )


def choose_archive_dir(today_dir):
    if not today_dir.exists():
        return today_dir

    if has_attachments(today_dir):
        answer = input(f"{today_dir} 已存在且包含附件，是否合并本次根目录附件？[y/N] ").strip().lower()
        if answer in {"y", "yes"}:
            return today_dir
        return next_numbered_dir(today_dir)

    return today_dir


def find_existing_today_dir(today_dir):
    candidates = [today_dir] + sorted(Path(".").glob(f"{today_dir.name}-*"))
    candidates = [path for path in candidates if path.is_dir() and list(path.glob("*.pdf"))]
    if not candidates:
        return None
    return candidates[-1]


def has_attachments(directory):
    return any(
        path.is_file() and path.suffix.lower() in ATTACHMENT_SUFFIXES
        for path in directory.iterdir()
    )


def next_numbered_dir(base_dir):
    index = 1
    while True:
        candidate = Path(f"{base_dir.name}-{index}")
        if not candidate.exists():
            return candidate
        index += 1


def unique_file_path(path):
    if not path.exists():
        return path

    index = 1
    while True:
        candidate = path.with_name(f"{path.stem}-{index}{path.suffix}")
        if not candidate.exists():
            return candidate
        index += 1


def command_schema(args):
    output_dir = Path(OUTPUT_DIR_NAME)
    output_dir.mkdir(exist_ok=True)
    text = json.dumps(SCHEMA_EXAMPLE, ensure_ascii=False, indent=2)
    if args.output:
        schema_path = output_dir / "schema_example.json"
        schema_path.write_text(text + "\n", encoding="utf-8")
        print(f"已写入 {schema_path}")
    else:
        print(text)


def command_fill(args):
    result = load_json(args.json_path)
    validate_result(result)
    ensure_fill_dependencies()

    json_path = Path(args.json_path)
    output_dir = json_path.parent
    archive_dir = output_dir.parent if output_dir.name == OUTPUT_DIR_NAME else Path(".")

    output_file = fill_template(result, archive_dir)
    verify_output(output_file, result)
    final_archive_dir = finalize_archive_dir(archive_dir, result)
    final_output_file = final_archive_dir / output_file.name
    cleanup_output_dir(final_archive_dir / OUTPUT_DIR_NAME)
    cleanup_pycache()

    print(f"完成：{final_output_file}")
    print(f"本次归档目录：{final_archive_dir}")


def load_json(json_path):
    path = Path(json_path)
    if not path.exists():
        raise SystemExit(f"找不到 JSON 文件：{json_path}")

    with open(path, encoding="utf-8") as f:
        try:
            return json.load(f)
        except json.JSONDecodeError as exc:
            raise SystemExit(f"JSON 格式错误：{exc}")


def validate_result(result):
    require_keys(result, ["meta", "expenses", "summary"], "根对象")

    meta = result["meta"]
    expenses = result["expenses"]
    summary = result["summary"]

    require_keys(meta, [
        "出差人", "出差事由", "出发日期", "结束日期", "出发地", "目的地",
        "出行方式", "出差天数", "事由简写"
    ], "meta")

    require_keys(expenses, ["住宿", "交通", "饮食", "其他"], "expenses")
    for category in ["住宿", "交通", "饮食", "其他"]:
        require_keys(expenses[category], ["金额", "单据张数", "说明"], f"expenses.{category}")
        require_number(expenses[category]["金额"], f"expenses.{category}.金额")
        require_integer(expenses[category]["单据张数"], f"expenses.{category}.单据张数")

    require_keys(summary, ["附件总张数", "报销总额", "预借差旅费", "补领金额", "退还金额"], "summary")

    expense_total = sum(float(expenses[c]["金额"]) for c in ["住宿", "交通", "饮食", "其他"])
    if abs(float(summary["报销总额"]) - expense_total) > 0.01:
        raise SystemExit(
            f"报销总额不一致：summary.报销总额={summary['报销总额']}，"
            f"费用合计={expense_total:.2f}"
        )

    receipt_total = sum(int(expenses[c]["单据张数"]) for c in ["住宿", "交通", "饮食", "其他"])
    if int(summary["附件总张数"]) != receipt_total:
        raise SystemExit(
            f"附件总张数不一致：summary.附件总张数={summary['附件总张数']}，"
            f"分类单据合计={receipt_total}"
        )


def require_keys(obj, keys, label):
    missing = [key for key in keys if key not in obj]
    if missing:
        raise SystemExit(f"{label} 缺少字段：{', '.join(missing)}")


def require_number(value, label):
    if not isinstance(value, (int, float)):
        raise SystemExit(f"{label} 必须是数字，当前为：{value!r}")


def require_integer(value, label):
    if not isinstance(value, int):
        raise SystemExit(f"{label} 必须是整数，当前为：{value!r}")


def ensure_fill_dependencies():
    if not TEMPLATE_FILE.exists():
        raise SystemExit(
            "找不到模板文件：差旅费明细表.xlsx\n"
            "请将 Excel 模板另存为 .xlsx 后放到当前目录。"
        )

    try:
        import openpyxl  # noqa: F401
    except ImportError:
        raise SystemExit("缺少 openpyxl，请运行：python -m pip install openpyxl")


def fill_template(result, output_dir):
    output_dir.mkdir(exist_ok=True)

    meta = result["meta"]
    output_file = output_dir / build_output_filename(meta)

    shutil.copy(TEMPLATE_FILE, output_file)
    update_xlsx_cells(output_file, build_cell_values(result), build_formula_cached_values(result))
    return output_file


def build_cell_values(result):
    meta = result["meta"]
    expenses = result["expenses"]
    return {
        "E3": meta["出差事由"],
        "A6": f'{meta["出发日期"]}-{meta["结束日期"]}',
        "B6": f'{meta["出发地"]}-{meta["目的地"]}',
        "C6": f'{meta["出差天数"]}天',
        "D6": meta["出行方式"],
        "F6": meta["出差天数"],
        "H6": expenses["住宿"]["单据张数"],
        "I6": expenses["住宿"]["金额"],
        "H7": expenses["饮食"]["单据张数"],
        "I7": expenses["饮食"]["金额"],
        "H8": expenses["交通"]["单据张数"],
        "I8": expenses["交通"]["金额"],
        "H9": expenses["其他"]["单据张数"],
        "I9": expenses["其他"]["金额"],
    }


def build_formula_cached_values(result):
    summary = result["summary"]
    advance = summary["预借差旅费"]
    total = summary["报销总额"]
    return {
        "H2": summary["附件总张数"],
        "F9": result["meta"]["出差天数"],
        "C10": total,
        "H10": max(total - advance, 0),
        "H11": max(advance - total, 0),
    }


def update_xlsx_cells(xlsx_path, cell_values, formula_cached_values):
    temp_file = NamedTemporaryFile(delete=False, suffix=".xlsx", dir=xlsx_path.parent)
    temp_path = Path(temp_file.name)
    temp_file.close()

    try:
        with zipfile.ZipFile(xlsx_path, "r") as source:
            with zipfile.ZipFile(temp_path, "w", compression=zipfile.ZIP_DEFLATED) as target:
                for item in source.infolist():
                    data = source.read(item.filename)
                    if item.filename == "xl/worksheets/sheet1.xml":
                        data = update_sheet_xml(data, cell_values, formula_cached_values)
                    target.writestr(item, data)
        temp_path.replace(xlsx_path)
    except Exception:
        temp_path.unlink(missing_ok=True)
        raise


def update_sheet_xml(sheet_xml, cell_values, formula_cached_values):
    namespace = "http://schemas.openxmlformats.org/spreadsheetml/2006/main"
    ET.register_namespace("", namespace)
    root = ET.fromstring(sheet_xml)
    for cell_ref, value in cell_values.items():
        cell = root.find(f".//{{{namespace}}}c[@r='{cell_ref}']")
        if cell is None:
            raise SystemExit(f"模板缺少单元格：{cell_ref}")
        set_cell_xml_value(cell, value, namespace)
    for cell_ref, value in formula_cached_values.items():
        cell = root.find(f".//{{{namespace}}}c[@r='{cell_ref}']")
        if cell is None:
            raise SystemExit(f"模板缺少单元格：{cell_ref}")
        set_formula_cached_value(cell, value, namespace)
    return ET.tostring(root, encoding="utf-8", xml_declaration=True)


def set_cell_xml_value(cell, value, namespace):
    for child in list(cell):
        if child.tag in {f"{{{namespace}}}v", f"{{{namespace}}}f", f"{{{namespace}}}is"}:
            cell.remove(child)

    if isinstance(value, (int, float)):
        cell.attrib.pop("t", None)
        value_node = ET.SubElement(cell, f"{{{namespace}}}v")
        value_node.text = format_number_for_excel(value)
    else:
        cell.set("t", "inlineStr")
        inline_string = ET.SubElement(cell, f"{{{namespace}}}is")
        text_node = ET.SubElement(inline_string, f"{{{namespace}}}t")
        text_node.text = str(value)


def set_formula_cached_value(cell, value, namespace):
    cell.attrib.pop("t", None)
    value_node = None
    for child in cell:
        if child.tag == f"{{{namespace}}}v":
            value_node = child
            break
    if value_node is None:
        value_node = ET.SubElement(cell, f"{{{namespace}}}v")
    value_node.text = format_number_for_excel(value)


def format_number_for_excel(value):
    if isinstance(value, int):
        return str(value)
    return f"{float(value):.10f}".rstrip("0").rstrip(".")


def finalize_archive_dir(archive_dir, result):
    if archive_dir == Path("."):
        return archive_dir

    target_dir = unique_archive_dir(archive_dir.parent / build_archive_dir_name(result["meta"]))
    if archive_dir.resolve() == target_dir.resolve():
        return archive_dir

    archive_dir.rename(target_dir)
    return target_dir


def unique_archive_dir(path):
    if not path.exists():
        return path

    index = 1
    while True:
        candidate = path.with_name(f"{path.name}-{index}")
        if not candidate.exists():
            return candidate
        index += 1


def cleanup_output_dir(output_dir):
    if output_dir.exists():
        shutil.rmtree(output_dir)


def cleanup_pycache():
    pycache_dir = Path("__pycache__")
    if pycache_dir.exists():
        shutil.rmtree(pycache_dir)


def build_archive_dir_name(meta):
    start_date = normalize_date_for_filename(meta["出发日期"])
    end_date = normalize_date_for_filename(meta["结束日期"])
    route = sanitize_filename(f'{meta["出发地"]}-{meta["目的地"]}')
    reason = sanitize_filename(meta["事由简写"])
    return f"{start_date}-{end_date}_{route}_{reason}"


def build_output_filename(meta):
    date_str = normalize_date_for_filename(meta["出发日期"])
    safe_parts = [
        "差旅费明细表",
        sanitize_filename(meta["出差人"]),
        date_str,
        sanitize_filename(meta["目的地"]),
        sanitize_filename(meta["事由简写"]),
    ]
    return "_".join(safe_parts) + ".xlsx"


def normalize_date_for_filename(date_text):
    parts = re.split(r"[/-]", str(date_text))
    if len(parts) != 3:
        return sanitize_filename(str(date_text))
    return f"{int(parts[0])}{int(parts[1]):02d}{int(parts[2]):02d}"


def sanitize_filename(text):
    text = str(text).strip()
    text = re.sub(r"[\\/:*?\"<>|\s]+", "", text)
    return text or "未填写"


def verify_output(output_file, result):
    import openpyxl

    wb = openpyxl.load_workbook(output_file, data_only=False)
    ws = wb.active
    meta = result["meta"]

    expected_route = f'{meta["出发地"]}-{meta["目的地"]}'
    expected_date = f'{meta["出发日期"]}-{meta["结束日期"]}'

    checks = [
        (ws["E3"].value, meta["出差事由"], "出差事由"),
        (ws["A6"].value, expected_date, "日期"),
        (ws["B6"].value, expected_route, "起止地点"),
    ]
    for actual, expected, label in checks:
        if actual != expected:
            raise SystemExit(f"输出校验失败：{label} 实际={actual!r} 预期={expected!r}")

    formula_checks = {
        "H2": "附件张数合计",
        "F9": "天数合计",
        "C10": "报销总额",
        "H10": "补领金额",
    }
    for cell, label in formula_checks.items():
        value = ws[cell].value
        if not isinstance(value, str) or not value.startswith("="):
            raise SystemExit(f"输出校验失败：{label} 单元格 {cell} 公式丢失，当前值={value!r}")

    verify_xlsx_styles_preserved(output_file)


def verify_xlsx_styles_preserved(output_file):
    with zipfile.ZipFile(TEMPLATE_FILE, "r") as template, zipfile.ZipFile(output_file, "r") as output:
        if template.read("xl/styles.xml") != output.read("xl/styles.xml"):
            raise SystemExit("输出校验失败：样式表 xl/styles.xml 与模板不一致")

        template_styles = read_sheet_cell_style_ids(template.read("xl/worksheets/sheet1.xml"))
        output_styles = read_sheet_cell_style_ids(output.read("xl/worksheets/sheet1.xml"))
        for cell_ref in filled_cell_refs():
            if template_styles.get(cell_ref) != output_styles.get(cell_ref):
                raise SystemExit(
                    f"输出校验失败：单元格 {cell_ref} 样式编号变化，"
                    f"模板={template_styles.get(cell_ref)!r} 输出={output_styles.get(cell_ref)!r}"
                )


def read_sheet_cell_style_ids(sheet_xml):
    namespace = "http://schemas.openxmlformats.org/spreadsheetml/2006/main"
    root = ET.fromstring(sheet_xml)
    return {
        cell.attrib["r"]: cell.attrib.get("s")
        for cell in root.findall(f".//{{{namespace}}}c")
        if "r" in cell.attrib
    }


def filled_cell_refs():
    return ["E3", "A6", "B6", "C6", "D6", "H6", "I6", "H7", "I7", "H8", "I8", "H9", "I9"]


def build_parser():
    parser = argparse.ArgumentParser(description="差旅报销自动化脚本")
    subparsers = parser.add_subparsers(dest="command", required=True)

    extract_parser = subparsers.add_parser("extract", help="整理根目录附件并提取 PDF 原文到本次归档目录")
    extract_parser.set_defaults(func=command_extract)

    schema_parser = subparsers.add_parser("schema", help="输出大模型需要生成的 JSON schema 示例")
    schema_parser.add_argument("--output", action="store_true", help="写入 报销输出/schema_example.json，而不是打印到终端")
    schema_parser.set_defaults(func=command_schema)

    fill_parser = subparsers.add_parser("fill", help="读取 JSON 并填入差旅费明细表模板，再归档本次目录")
    fill_parser.add_argument("json_path", help="大模型生成的 analyzed_result.json 路径")
    fill_parser.set_defaults(func=command_fill)

    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()