#!/usr/bin/env python3
"""
pdf-exporter: 将 HTML 报告导出为支持中文的 PDF 文件
用法:
  python3 export_pdf.py --input <html_file_path> --output <pdf_file_path> [--title <标题>] [--analysis <分析文字>]

参数:
  --input     输入 HTML 文件路径
  --output    输出 PDF 文件路径（建议保存到 /workspace/Document/PersonalDocuments/）
  --title     PDF 标题（可选）
  --analysis  AI 分析内容（可选），支持 Markdown 格式，## 开头表示小标题，- 开头表示列表项

依赖:
  - reportlab (已安装于 /workspace/.local/lib/python3.11/site-packages)
  - 中文字体: /workspace/skills/pdf-exporter/fonts/wqy-microhei.ttc
"""

import sys
import os
import re
import json
import argparse
import html as html_module
from datetime import datetime

# 将 reportlab 安装路径加入 sys.path
sys.path.insert(0, '/workspace/.local/lib/python3.11/site-packages')

try:
    from reportlab.lib.pagesizes import A4
    from reportlab.lib.units import mm
    from reportlab.lib import colors
    from reportlab.pdfbase import pdfmetrics
    from reportlab.pdfbase.ttfonts import TTFont
    from reportlab.platypus import (
        SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle,
        HRFlowable, KeepTogether
    )
    from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
    from reportlab.lib.enums import TA_LEFT, TA_CENTER, TA_RIGHT, TA_JUSTIFY
except ImportError as e:
    print(json.dumps({"status": "error", "message": f"reportlab 未安装: {e}"}))
    sys.exit(1)

# ── 字体注册 ──────────────────────────────────────────────────────────────────
SKILL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
FONT_PATH = os.path.join(SKILL_DIR, 'fonts', 'wqy-microhei.ttc')
FONT_NAME = 'WQYMicroHei'


def register_chinese_font():
    """注册中文字体，失败时回退到 Helvetica 并打印警告"""
    if not os.path.exists(FONT_PATH):
        sys.stderr.write(f"[WARNING] 中文字体文件不存在: {FONT_PATH}\n")
        return False
    try:
        pdfmetrics.registerFont(TTFont(FONT_NAME, FONT_PATH))
        return True
    except Exception as e:
        sys.stderr.write(f"[WARNING] 字体注册失败: {e}\n")
        return False


# ── HTML 文本清理 ─────────────────────────────────────────────────────────────
def strip_html_tags(text):
    text = re.sub(r'<style[^>]*>.*?</style>', '', text, flags=re.DOTALL | re.IGNORECASE)
    text = re.sub(r'<script[^>]*>.*?</script>', '', text, flags=re.DOTALL | re.IGNORECASE)
    text = re.sub(r'<[^>]+>', '', text)
    text = html_module.unescape(text)
    return text.strip()


def extract_html_sections(html_content):
    """从 HTML 中提取标题、段落、表格等结构化内容"""
    sections = []

    m = re.search(r'<title[^>]*>(.*?)</title>', html_content, re.IGNORECASE | re.DOTALL)
    if m:
        sections.append({'type': 'title', 'text': strip_html_tags(m.group(1))})

    body_m = re.search(r'<body[^>]*>(.*?)</body>', html_content, re.IGNORECASE | re.DOTALL)
    body = body_m.group(1) if body_m else html_content

    pattern = re.compile(
        r'(<h[1-4][^>]*>.*?</h[1-4]>|<p[^>]*>.*?</p>|<table[^>]*>.*?</table>)',
        re.IGNORECASE | re.DOTALL
    )
    for match in pattern.finditer(body):
        tag_html = match.group(1)
        tag_name = re.match(r'<(\w+)', tag_html).group(1).lower()

        if tag_name in ('h1', 'h2', 'h3', 'h4'):
            level = int(tag_name[1])
            text = strip_html_tags(tag_html).strip()
            if text:
                sections.append({'type': 'heading', 'level': level, 'text': text})

        elif tag_name == 'p':
            text = strip_html_tags(tag_html).strip()
            if text and len(text) > 1:
                sections.append({'type': 'paragraph', 'text': text})

        elif tag_name == 'table':
            rows_data = []
            for row_m in re.finditer(r'<tr[^>]*>(.*?)</tr>', tag_html, re.IGNORECASE | re.DOTALL):
                row_html = row_m.group(1)
                cells = [strip_html_tags(c).strip()
                         for c in re.findall(r'<t[hd][^>]*>(.*?)</t[hd]>', row_html,
                                             re.IGNORECASE | re.DOTALL)]
                if cells:
                    rows_data.append(cells)
            if rows_data:
                sections.append({'type': 'table', 'rows': rows_data})

    return sections


# ── AI 分析文字解析 ───────────────────────────────────────────────────────────
def parse_analysis_text(analysis_text):
    """
    将 Agent 传入的分析文字解析为结构化 sections。
    支持：
      - ## 小标题（绿色加粗）
      - # 大标题（同 ## 处理）
      - - / • / * 开头的列表项
      - 1. 2. 3. 数字列表
      - 普通段落（空行分隔）
    """
    sections = []
    if not analysis_text or not analysis_text.strip():
        return sections

    lines = analysis_text.strip().split('\n')
    current_para = []

    def flush_para():
        if current_para:
            text = ' '.join(current_para).strip()
            if text:
                sections.append({'type': 'analysis_paragraph', 'text': text})
            current_para.clear()

    for line in lines:
        line_stripped = line.strip()

        if not line_stripped:
            flush_para()
            continue

        if line_stripped.startswith('## '):
            flush_para()
            sections.append({'type': 'analysis_heading', 'text': line_stripped[3:].strip()})
        elif line_stripped.startswith('# '):
            flush_para()
            sections.append({'type': 'analysis_heading', 'text': line_stripped[2:].strip()})
        elif line_stripped.startswith(('- ', '• ', '* ', '· ')):
            flush_para()
            sections.append({'type': 'analysis_bullet', 'text': line_stripped[2:].strip()})
        elif re.match(r'^\d+[\.\、]\s+', line_stripped):
            flush_para()
            text = re.sub(r'^\d+[\.\、]\s+', '', line_stripped)
            sections.append({'type': 'analysis_bullet', 'text': text})
        else:
            current_para.append(line_stripped)

    flush_para()
    return sections


# ── PDF 样式 ──────────────────────────────────────────────────────────────────
def build_styles(has_chinese_font):
    base_font = FONT_NAME if has_chinese_font else 'Helvetica'

    return {
        'doc_title': ParagraphStyle(
            'DocTitle', fontName=base_font, fontSize=20, leading=28,
            spaceAfter=8, alignment=TA_CENTER,
            textColor=colors.HexColor('#1a1a2e'),
        ),
        'h1': ParagraphStyle(
            'H1', fontName=base_font, fontSize=16, leading=22,
            spaceBefore=14, spaceAfter=6,
            textColor=colors.HexColor('#1677FF'),
        ),
        'h2': ParagraphStyle(
            'H2', fontName=base_font, fontSize=13, leading=18,
            spaceBefore=10, spaceAfter=4,
            textColor=colors.HexColor('#333333'),
        ),
        'h3': ParagraphStyle(
            'H3', fontName=base_font, fontSize=11, leading=16,
            spaceBefore=8, spaceAfter=3,
            textColor=colors.HexColor('#555555'),
        ),
        'h4': ParagraphStyle(
            'H4', fontName=base_font, fontSize=10, leading=14,
            spaceBefore=6, spaceAfter=2,
            textColor=colors.HexColor('#777777'),
        ),
        'body': ParagraphStyle(
            'Body', fontName=base_font, fontSize=10, leading=16,
            spaceAfter=6, textColor=colors.HexColor('#333333'),
        ),
        'table_header': ParagraphStyle(
            'TableHeader', fontName=base_font, fontSize=9, leading=13,
            alignment=TA_CENTER, textColor=colors.white,
        ),
        'table_cell': ParagraphStyle(
            'TableCell', fontName=base_font, fontSize=9, leading=13,
            alignment=TA_LEFT, textColor=colors.HexColor('#333333'),
        ),
        # ── AI 分析专用样式 ──
        'ai_section_title': ParagraphStyle(
            'AISectionTitle', fontName=base_font, fontSize=15, leading=22,
            spaceBefore=16, spaceAfter=8,
            textColor=colors.HexColor('#0d6e4f'),
        ),
        'ai_heading': ParagraphStyle(
            'AIHeading', fontName=base_font, fontSize=11, leading=16,
            spaceBefore=10, spaceAfter=4,
            textColor=colors.HexColor('#0d6e4f'),
        ),
        'ai_paragraph': ParagraphStyle(
            'AIParagraph', fontName=base_font, fontSize=10, leading=17,
            spaceAfter=6, alignment=TA_JUSTIFY,
            textColor=colors.HexColor('#1a1a1a'),
            leftIndent=4,
        ),
        'ai_bullet': ParagraphStyle(
            'AIBullet', fontName=base_font, fontSize=10, leading=16,
            spaceAfter=4, textColor=colors.HexColor('#1a1a1a'),
            leftIndent=16, bulletIndent=4,
        ),
        'ai_footer': ParagraphStyle(
            'AIFooter', fontName=base_font, fontSize=8, leading=12,
            spaceAfter=4, textColor=colors.HexColor('#888888'),
            alignment=TA_RIGHT,
        ),
    }


# ── AI 分析章节构建 ───────────────────────────────────────────────────────────
def build_analysis_block(analysis_sections, styles):
    """构建 AI 分析章节的 flowables，插入在数据内容之后"""
    flowables = []

    flowables.append(Spacer(1, 10))
    flowables.append(HRFlowable(width='100%', thickness=1.5,
                                color=colors.HexColor('#0d6e4f'), spaceAfter=6))
    flowables.append(Paragraph('AI 智能分析', styles['ai_section_title']))
    flowables.append(HRFlowable(width='100%', thickness=0.5,
                                color=colors.HexColor('#b2dfdb'), spaceAfter=8))

    for sec in analysis_sections:
        if sec['type'] == 'analysis_heading':
            flowables.append(Paragraph(sec['text'], styles['ai_heading']))
        elif sec['type'] == 'analysis_paragraph':
            flowables.append(Paragraph(sec['text'], styles['ai_paragraph']))
        elif sec['type'] == 'analysis_bullet':
            flowables.append(Paragraph('- ' + sec['text'], styles['ai_bullet']))

    gen_time = datetime.now().strftime('%Y-%m-%d %H:%M')
    flowables.append(Spacer(1, 6))
    flowables.append(Paragraph(
        f'以上分析由 AI 自动生成，仅供参考 · 生成时间：{gen_time}',
        styles['ai_footer']
    ))
    flowables.append(HRFlowable(width='100%', thickness=0.5,
                                color=colors.HexColor('#b2dfdb'), spaceBefore=4, spaceAfter=8))

    return flowables


# ── 内容转 flowables ──────────────────────────────────────────────────────────
def sections_to_flowables(sections, styles, page_width):
    flowables = []
    usable_width = page_width - 40 * mm
    heading_styles = {1: styles['h1'], 2: styles['h2'], 3: styles['h3'], 4: styles['h4']}

    for sec in sections:
        if sec['type'] == 'title':
            flowables.append(Paragraph(sec['text'], styles['doc_title']))
            flowables.append(HRFlowable(width='100%', thickness=2,
                                        color=colors.HexColor('#1677FF'), spaceAfter=12))

        elif sec['type'] == 'heading':
            level = sec.get('level', 2)
            style = heading_styles.get(level, styles['h2'])
            flowables.append(Paragraph(sec['text'], style))
            if level == 1:
                flowables.append(HRFlowable(width='100%', thickness=1,
                                            color=colors.HexColor('#e0e0e0'), spaceAfter=4))

        elif sec['type'] == 'paragraph':
            text = sec['text']
            if len(text.strip()) > 2:
                flowables.append(Paragraph(text, styles['body']))

        elif sec['type'] == 'table':
            rows = sec['rows']
            if not rows:
                continue
            col_count = max(len(r) for r in rows)
            col_width = usable_width / col_count
            table_data = []
            for i, row in enumerate(rows):
                while len(row) < col_count:
                    row.append('')
                cell_style = styles['table_header'] if i == 0 else styles['table_cell']
                table_data.append([Paragraph(str(cell), cell_style) for cell in row])

            t = Table(table_data, colWidths=[col_width] * col_count, repeatRows=1)
            t.setStyle(TableStyle([
                ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#1677FF')),
                ('TEXTCOLOR',  (0, 0), (-1, 0), colors.white),
                ('FONTSIZE',   (0, 0), (-1, 0), 9),
                ('ALIGN',      (0, 0), (-1, 0), 'CENTER'),
                ('VALIGN',     (0, 0), (-1, -1), 'MIDDLE'),
                ('ROWBACKGROUNDS', (0, 1), (-1, -1),
                 [colors.white, colors.HexColor('#f5f8ff')]),
                ('GRID',       (0, 0), (-1, -1), 0.5, colors.HexColor('#d0d0d0')),
                ('LINEBELOW',  (0, 0), (-1, 0), 1.5, colors.HexColor('#1677FF')),
                ('TOPPADDING',    (0, 0), (-1, -1), 5),
                ('BOTTOMPADDING', (0, 0), (-1, -1), 5),
                ('LEFTPADDING',   (0, 0), (-1, -1), 6),
                ('RIGHTPADDING',  (0, 0), (-1, -1), 6),
            ]))
            flowables.append(t)
            flowables.append(Spacer(1, 8))

    return flowables


# ── 主函数 ────────────────────────────────────────────────────────────────────
def generate_pdf(input_path, output_path, title=None, analysis_text=None):
    """读取 HTML，生成 PDF，可选插入 AI 分析章节"""
    has_chinese_font = register_chinese_font()

    try:
        with open(input_path, 'r', encoding='utf-8') as f:
            html_content = f.read()
    except Exception as e:
        return {"status": "error", "message": f"读取 HTML 文件失败: {e}"}

    sections = extract_html_sections(html_content)

    if title:
        sections.insert(0, {'type': 'title', 'text': title})

    styles = build_styles(has_chinese_font)

    os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else '.', exist_ok=True)
    doc = SimpleDocTemplate(
        output_path,
        pagesize=A4,
        leftMargin=20 * mm, rightMargin=20 * mm,
        topMargin=20 * mm,  bottomMargin=20 * mm,
        title=title or '数据分析报告',
        author='OpenClaw Agent',
    )

    page_width = A4[0]
    flowables = sections_to_flowables(sections, styles, page_width)

    # 插入 AI 分析章节（在数据内容之后）
    has_analysis = False
    if analysis_text and analysis_text.strip():
        analysis_sections = parse_analysis_text(analysis_text)
        if analysis_sections:
            flowables.extend(build_analysis_block(analysis_sections, styles))
            has_analysis = True

    if not flowables:
        return {"status": "error", "message": "HTML 内容为空，无法生成 PDF"}

    try:
        doc.build(flowables)
    except Exception as e:
        return {"status": "error", "message": f"PDF 生成失败: {e}"}

    file_size = os.path.getsize(output_path)
    return {
        "status": "success",
        "message": "PDF 导出成功" + ("（含 AI 分析）" if has_analysis else ""),
        "output_path": output_path,
        "file_size_kb": round(file_size / 1024, 1),
        "chinese_font": has_chinese_font,
        "sections_count": len(sections),
        "has_analysis": has_analysis,
    }


def main():
    parser = argparse.ArgumentParser(description='将 HTML 报告导出为支持中文的 PDF')
    parser.add_argument('--input',    required=True,  help='输入 HTML 文件路径')
    parser.add_argument('--output',   required=True,  help='输出 PDF 文件路径')
    parser.add_argument('--title',    default=None,   help='PDF 标题（可选）')
    parser.add_argument('--analysis', default=None,
                        help='AI 分析内容（可选），支持 Markdown 格式：## 小标题，- 列表项，普通段落')
    args = parser.parse_args()

    result = generate_pdf(args.input, args.output, args.title, args.analysis)
    print(json.dumps(result, ensure_ascii=False))
    sys.exit(0 if result['status'] == 'success' else 1)


if __name__ == '__main__':
    main()