"""
天目 OCR 布局质量集成测试。

生成模拟真实文档的图片（含标题、正文、列表、编号项、分隔线等），
验证：
  1. PP-OCRv6 正确识别文字内容
  2. 布局后处理生成结构化 Markdown（标题级别、段落分组）
  3. 双栏排版识别

使用方法：
  1. 先启动天目服务：uv run tianmu
  2. 运行本脚本：python tests/test_integration_layout.py
"""

from __future__ import annotations

import base64
import io
import json
import os
import sys
import urllib.request

from PIL import Image, ImageDraw, ImageFont

API_BASE = os.environ.get("TIANMU_API_BASE", "http://localhost:8101")


# ── 字体加载 ──────────────────────────────────────────────

def _load_font(size: int) -> ImageFont.FreeTypeFont | None:
    paths = [
        "/System/Library/Fonts/PingFang.ttc",
        "/System/Library/Fonts/STHeiti Light.ttc",
        "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc",
        "C:/Windows/Fonts/msyh.ttc",
    ]
    for fp in paths:
        if os.path.exists(fp):
            try:
                return ImageFont.truetype(fp, size)
            except Exception:
                continue
    return None


# ── 测试场景 1：标准文档布局 ────────────────────────────

def create_standard_doc() -> bytes:
    """生成一个含标题、正文、列表的标准文档图片。"""
    font_title = _load_font(28)
    font_body = _load_font(16)

    img = Image.new("RGB", (600, 400), color="white")
    draw = ImageDraw.Draw(img)

    y = 20
    lines = []
    title = "项目实施方案"
    lines.append((title, font_title, (20, y)))
    y += 45

    lines.append(("", None, None))  # 空行

    h2 = "一、项目背景"
    lines.append((h2, font_body, (20, y)))
    y += 30

    body_texts = [
        "本项目旨在构建一套基于视觉感知的智能文档处理系统。",
        "系统通过本地OCR引擎实现文字提取，支持多种文档类型。",
    ]
    for t in body_texts:
        lines.append((t, font_body, (20, y)))
        y += 25

    lines.append(("", None, None))
    y += 5

    items = [
        "• 核心技术栈：Python + FastAPI + ONNX Runtime",
        "• 模型支持：PP-OCRv6 + PaddleOCR-VL",
        "• 部署方式：Docker 容器化微服务",
    ]
    for item in items:
        lines.append((item, font_body, (20, y)))
        y += 25

    for text, font, pos in lines:
        if text and font and pos:
            draw.text(pos, text, fill="black", font=font)

    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


# ── 测试场景 2：双栏布局 ────────────────────────────────

def create_two_column_doc() -> bytes:
    """生成一个模拟双栏排版的图片。"""
    font_body = _load_font(14)

    img = Image.new("RGB", (700, 300), color="white")
    draw = ImageDraw.Draw(img)

    # 左栏 (x=20~300)
    left_items = [
        "左栏项目一：需求分析",
        "左栏项目二：方案设计",
        "左栏项目三：开发实施",
    ]
    y = 20
    for item in left_items:
        draw.text((20, y), item, fill="black", font=font_body)
        y += 30

    # 右栏 (x=380~680)
    right_items = [
        "右栏概述：本栏说明验收标准",
        "右栏条目1：功能完整性",
        "右栏条目2：性能达标",
    ]
    y = 20
    for item in right_items:
        draw.text((380, y), item, fill="black", font=font_body)
        y += 30

    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


# ── 测试场景 3：编号 + TOC 目录 ─────────────────────────

def create_toc_doc() -> bytes:
    """生成含编号和 TOC 目录的图片。"""
    font = _load_font(16)

    img = Image.new("RGB", (500, 250), color="white")
    draw = ImageDraw.Draw(img)

    items = [
        ("1. 引言", 20),
        ("2. 相关工作", 50),
        ("  2.1 传统方法", 75),
        ("  2.2 深度学习方法", 100),
        ("3. 实验与分析", 130),
        ("  3.1 数据集", 155),
        ("  3.2 实验结果", 180),
    ]
    for text, y in items:
        draw.text((20, y), text, fill="black", font=font)

    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


# ── 运行测试 ──────────────────────────────────────────────

def run_ocr(image_bytes: bytes, engine: str = "ppocrv6") -> dict | None:
    b64 = base64.b64encode(image_bytes).decode("ascii")
    body = json.dumps({"image": b64, "engine": engine}).encode("utf-8")
    try:
        req = urllib.request.Request(
            f"{API_BASE}/v1/ocr",
            data=body,
            headers={"Content-Type": "application/json"},
        )
        with urllib.request.urlopen(req, timeout=30) as resp:
            return json.loads(resp.read())
    except Exception as e:
        print(f"  ✗ API 请求失败: {e}")
        return None


def main():
    print("=" * 60)
    print("  天目 OCR 布局质量集成测试")
    print(f"  API: {API_BASE}")
    print("=" * 60)

    # 1. 标准文档布局
    print("\n▶ 1/3 标准文档（标题 + 正文 + 列表）")
    img1 = create_standard_doc()
    result = run_ocr(img1)
    if result and result.get("text"):
        text = result["text"]
        print(f"  识别文本 ({len(text)} 字符):")
        for line in text.split("\n")[:10]:
            if line.strip():
                print(f"    {line.strip()[:80]}")
        if "项目实施方案" in text:
            print("  ✓ 识别到标题")
        if "•" in text or "- 核心" in text:
            print("  ✓ 识别到列表项")
        print(f"  耗时: {result.get('duration_ms', 0):.0f} ms")
    else:
        err = result.get("error") if result else "无响应"
        print(f"  ✗ 识别失败: {err}")

    # 2. 双栏布局
    print("\n▶ 2/3 双栏排版")
    img2 = create_two_column_doc()
    result = run_ocr(img2)
    if result and result.get("text"):
        text = result["text"]
        print(f"  识别文本 ({len(text)} 字符):")
        for line in text.split("\n")[:10]:
            if line.strip():
                print(f"    {line.strip()[:80]}")
        if "左栏项目" in text and "右栏概述" in text:
            print("  ✓ 双栏内容均被识别")
        print(f"  耗时: {result.get('duration_ms', 0):.0f} ms")
    else:
        err = result.get("error") if result else "无响应"
        print(f"  ✗ 识别失败: {err}")

    # 3. 编号 + 目录
    print("\n▶ 3/3 编号项 + TOC 目录")
    img3 = create_toc_doc()
    result = run_ocr(img3)
    if result and result.get("text"):
        text = result["text"]
        print(f"  识别文本 ({len(text)} 字符):")
        for line in text.split("\n")[:10]:
            if line.strip():
                print(f"    {line.strip()[:80]}")
        checks = [
            ("1. 引言" in text, "识别编号 1."),
            ("2. 相关" in text, "识别编号 2."),
            ("3. 实验" in text, "识别编号 3."),
        ]
        for ok, desc in checks:
            print(f"  {'✓' if ok else '✗'} {desc}")
        print(f"  耗时: {result.get('duration_ms', 0):.0f} ms")
    else:
        err = result.get("error") if result else "无响应"
        print(f"  ✗ 识别失败: {err}")

    print("\n" + "=" * 60)
    print("  测试完成")
    print("=" * 60)


if __name__ == "__main__":
    main()