#!/usr/bin/env python3
"""
Docs as Code — 构建入口

读取 project.yaml → 逐章节处理 @import → Markdown → HTML → 图片内嵌
→ 组装 → 注入锚点 → 提取目录 → 模板渲染 → 自包含 HTML 报告。

用法:
    python doc_builder/build.py

输出:
    output/<标题>.html  （自包含 HTML，可离线分发）
"""

import base64
import csv
import importlib.util
import io
import json
import re
import sys
from datetime import date
from pathlib import Path

import markdown
import yaml
from jinja2 import Environment, FileSystemLoader, select_autoescape

# ---------- 路径配置 ----------
PROJECT_ROOT = Path(__file__).resolve().parent.parent
CHAPTERS_DIR = PROJECT_ROOT / "chapters"
ASSETS_DIR = PROJECT_ROOT / "assets"
OUTPUT_DIR = PROJECT_ROOT / "output"
BUILDER_DIR = PROJECT_ROOT / "doc_builder"
TEMPLATES_DIR = BUILDER_DIR / "templates"
THEMES_DIR = BUILDER_DIR / "themes"
RENDERERS_DIR = BUILDER_DIR / "renderers"
CHECKS_DIR = BUILDER_DIR / "checks"

# ---------- 配置加载 ----------

def load_config():
    path = PROJECT_ROOT / "project.yaml"
    if not path.exists():
        raise FileNotFoundError("找不到 project.yaml")
    with open(path, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f) or {}

    config.setdefault("title", "未命名文档")
    config.setdefault("subtitle", "")
    config.setdefault("author", "")
    config.setdefault("version", "")
    config.setdefault("doc_type", "技术文档")
    config.setdefault("logo", "")
    config.setdefault("lang", "zh-CN")
    config.setdefault("date", date.today().isoformat())
    return config


# ---------- 插件发现 ----------

def load_plugin_modules(directory):
    """扫描目录下的 *.py 文件并导入为模块字典。"""
    modules = {}
    if not directory.exists():
        return modules
    for py_file in sorted(directory.glob("*.py")):
        if py_file.name.startswith("_"):
            continue
        try:
            spec = importlib.util.spec_from_file_location(py_file.stem, py_file)
            if spec is None or spec.loader is None:
                continue
            mod = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(mod)
            modules[py_file.stem] = mod
            print(f"  已加载: {py_file.stem}")
        except Exception as e:
            print(f"  警告: 加载 {py_file.name} 失败: {e}")
    return modules


# ---------- @import 处理 ----------

IMPORT_RE = re.compile(r'^@import\s+"([^"]+)"(?:\s+using\s+(\S+))?\s*$', re.MULTILINE)


def resolve_import_path(raw, base_dir):
    path = Path(raw)
    if path.is_absolute():
        return path
    return (base_dir / path).resolve()


def default_data_renderer(filepath):
    """默认数据文件渲染：CSV → HTML 表格，YAML/JSON → 代码块。"""
    ext = filepath.suffix.lower()
    if ext == ".csv":
        with open(filepath, newline="", encoding="utf-8") as f:
            reader = csv.reader(f)
            rows = list(reader)
        if not rows:
            return ""
        lines = ["<table>"]
        lines.append("<tr>" + "".join(f"<th>{c}</th>" for c in rows[0]) + "</tr>")
        for row in rows[1:]:
            lines.append("<tr>" + "".join(f"<td>{c}</td>" for c in row) + "</tr>")
        lines.append("</table>")
        return "".join(lines)
    elif ext in (".yaml", ".yml"):
        with open(filepath, "r", encoding="utf-8") as f:
            data = yaml.safe_load(f) or {}
        return f"<pre><code>{yaml.dump(data, allow_unicode=True)}</code></pre>"
    elif ext == ".json":
        with open(filepath, "r", encoding="utf-8") as f:
            data = json.load(f)
        return f"<pre><code>{json.dumps(data, ensure_ascii=False, indent=2)}</code></pre>"
    else:
        content = filepath.read_text(encoding="utf-8")
        return f"<pre><code>{content}</code></pre>"


def render_with_module(renderers, renderer_name, filepath):
    """调用指定渲染器。"""
    if renderer_name not in renderers:
        raise ValueError(f"找不到指定渲染器：{renderer_name}")
    func = getattr(renderers[renderer_name], "render", None)
    if not callable(func):
        raise ValueError(f"{renderer_name} 没有 render(content: str) -> str 函数")
    return func(str(filepath))


def rewrite_image_paths(text, source_dir):
    """把被导入 Markdown 中的相对图片路径改为绝对路径，供后续 base64 嵌入。"""
    def repl(match):
        alt = match.group(1)
        src = match.group(2)
        if src.startswith(("http://", "https://", "data:")) or Path(src).is_absolute():
            return match.group(0)
        abs_path = (source_dir / src).resolve()
        return f'![{alt}]({abs_path})'
    return re.sub(r'!\[([^\]]*)\]\(([^)]+)\)', repl, text)


def process_imports(text, base_dir, renderers, _imported=None):
    """扫描并替换 @import 指令。支持 .md 注入和数据文件导入。"""
    if _imported is None:
        _imported = set()

    def repl(match):
        raw = match.group(1)
        renderer_name = match.group(2)
        target = resolve_import_path(raw, base_dir)

        # 外部 Markdown 导入
        if raw.endswith(".md"):
            if target in _imported:
                raise RuntimeError(f"检测到循环 @import：{target}")
            _imported.add(target)
            if not target.exists():
                raise FileNotFoundError(f"找不到要导入的 Markdown 文件：{target}")
            md = target.read_text(encoding="utf-8")
            md = rewrite_image_paths(md, target.parent)
            md = process_imports(md, target.parent, renderers, _imported)
            return md

        # 数据导入
        if not target.exists():
            raise FileNotFoundError(f"找不到要导入的数据文件：{target}")
        if renderer_name:
            renderer_name = renderer_name.removesuffix(".py")
            return render_with_module(renderers, renderer_name, target)
        return default_data_renderer(target)

    return IMPORT_RE.sub(repl, text)


# ---------- 图片属性 {w=50%} ----------

IMAGE_ATTR_RE = re.compile(r'!\[([^\]]*)\]\(([^)]+)\)\{([^}]*)\}')


def parse_attrs(attr_str):
    attrs = {}
    for part in attr_str.split(","):
        part = part.strip()
        if "=" in part:
            k, v = part.split("=", 1)
            attrs[k.strip()] = v.strip()
    return attrs


def process_image_attrs(text):
    """处理 ![](path){w=50%} 图片属性语法，转换为 HTML img 标签。"""
    def repl(match):
        alt = match.group(1)
        src = match.group(2)
        attr_str = match.group(3)
        attrs = parse_attrs(attr_str)
        style = ""
        if "w" in attrs:
            style = f'width:{attrs["w"]};'
        cls = attrs.get("class", "")
        cls_attr = f' class="{cls}"' if cls else ""
        style_attr = f' style="{style}"' if style else ""
        return f'<img src="{src}" alt="{alt}"{cls_attr}{style_attr} />'
    return IMAGE_ATTR_RE.sub(repl, text)


# ---------- KaTeX 预处理 ----------

def preprocess_katex(text):
    """
    将 Markdown 中的 LaTeX 公式包装为原始 HTML，
    防止 Markdown 解析器错误解释公式中的 _ * 等字符。
    """
    # 块级公式 $$...$$
    def protect_display(match):
        latex = match.group(1)
        return f'<div class="math-display">$${latex}$$</div>'
    text = re.sub(r'\$\$\s*(.+?)\s*\$\$', protect_display, text, flags=re.DOTALL)

    # 行内公式 $...$
    def protect_inline(match):
        latex = match.group(1)
        return f'<span class="math-inline">${latex}$</span>'
    text = re.sub(r'(?<!\d)\$([^$\s].*?[^$\s])\$(?!\d)', protect_inline, text)

    return text


# ---------- 自定义代码块渲染 ----------

FENCED_CODEBLOCK_OPEN_RE = re.compile(r'^```(\w+)(?:\s+[^\n]*)?$')


def render_custom_codeblocks_md(text, renderers):
    """把有对应渲染器的 fenced code block 替换为 HTML。"""
    lines = text.splitlines(keepends=True)
    out_lines = []
    i = 0
    while i < len(lines):
        line = lines[i]
        m = FENCED_CODEBLOCK_OPEN_RE.match(line)
        if not m:
            out_lines.append(line)
            i += 1
            continue

        lang = m.group(1)
        renderer_name = f"render_{lang.replace('-', '_')}"
        if renderer_name not in renderers:
            out_lines.append(line)
            i += 1
            continue

        func = getattr(renderers[renderer_name], "render", None)
        if not callable(func):
            out_lines.append(line)
            i += 1
            continue

        # 收集到闭合 fence
        start = i + 1
        j = start
        while j < len(lines) and lines[j].strip() != '```':
            j += 1
        if j >= len(lines):
            out_lines.append(line)
            i += 1
            continue

        content = "".join(lines[start:j]).rstrip("\n")
        rendered = func(content)
        if not rendered.endswith("\n"):
            rendered += "\n"
        out_lines.append(rendered)
        # 跳过 fence 后的换行
        if j + 1 < len(lines) and lines[j + 1] == "\n":
            i = j + 2
        else:
            i = j + 1

    return "".join(out_lines)


# ---------- Markdown → HTML ----------

def markdown_to_html(text):
    md = markdown.Markdown(extensions=[
        "extra",
        "tables",
        "fenced_code",
        "toc",
    ])
    return md.convert(text)


# ---------- 图片 base64 内嵌（HTML 级别） ----------

IMG_SRC_RE = re.compile(r'<img([^>]*?)src=["\']([^"\']+)["\']([^>]*)>', re.IGNORECASE)

MIME_TABLE = {
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".gif": "image/gif",
    ".svg": "image/svg+xml",
    ".webp": "image/webp",
}


def guess_mime(ext):
    return MIME_TABLE.get(ext.lower(), "application/octet-stream")


def embed_images(html, base_dirs):
    """扫描 HTML 中的 img 标签，将本地图片替换为 base64 data URI。"""
    def resolve_src(src):
        if Path(src).is_absolute():
            path = Path(src)
            if path.exists():
                return path
            return None
        for base in base_dirs:
            path = base / src
            if path.exists():
                return path
        return None

    def repl(match):
        prefix = match.group(1)
        src = match.group(2)
        suffix = match.group(3)

        if src.startswith(("http://", "https://", "data:")):
            return match.group(0)

        path = resolve_src(src)
        if path is None:
            print(f"  [警告] 找不到图片，保留原路径：{src}")
            return match.group(0)
        try:
            mime = guess_mime(path.suffix)
            data = path.read_bytes()
            b64 = base64.b64encode(data).decode("ascii")
            return f'<img{prefix}src="data:{mime};base64,{b64}"{suffix}>'
        except Exception as e:
            print(f"  [警告] 图片 base64 编码失败（{src}）：{e}")
            return match.group(0)

    return IMG_SRC_RE.sub(repl, html)


# ---------- 目录与锚点 ----------

def slugify(text):
    """生成 HTML 锚点 ID。"""
    anchor = re.sub(r'[^\w\s一-鿿-]', '', text)
    return anchor.strip().replace(" ", "-")[:50]


def generate_toc(html):
    """从 HTML 中提取 H1-H4 标题，生成目录。"""
    toc = []
    for m in re.finditer(r'<h([1-4])[^>]*>(.*?)</h\1>', html, re.DOTALL):
        level = int(m.group(1))
        text = re.sub(r'<.*?>', '', m.group(2)).strip()
        toc.append({"level": level, "text": text, "anchor": slugify(text)})
    return toc


def inject_anchors(html):
    """为所有 H1-H4 标签注入 id 属性，使侧边栏目录可跳转。"""
    def repl(match):
        level = match.group(1)
        attrs = match.group(2)
        inner = match.group(3)
        anchor = slugify(re.sub(r'<.*?>', '', inner).strip())
        return f'<h{level} id="{anchor}"{attrs}>{inner}</h{level}>'
    return re.sub(
        r'<h([1-4])([^>]*)>(.*?)</h\1>',
        repl,
        html,
        flags=re.DOTALL,
    )


# ---------- Logo 处理 ----------

def find_logo(logo_path_str):
    """定位 logo 文件；需在 project.yaml 明确指定 logo 字段。"""
    if not logo_path_str:
        return None
    path = Path(logo_path_str)
    if path.is_absolute():
        return path
    return PROJECT_ROOT / path


def embed_logo(logo_path, max_width=400):
    """压缩 logo 并返回 base64 data URI。"""
    if logo_path is None:
        return None

    try:
        from PIL import Image
    except ImportError:
        print("  提示: 未安装 Pillow，跳过 logo 处理")
        return None

    try:
        img = Image.open(logo_path)
        w, h = img.size
        if w > max_width:
            ratio = max_width / w
            img = img.resize((max_width, int(h * ratio)), Image.Resampling.LANCZOS)

        ext = logo_path.suffix.lower()
        if ext == ".svg":
            data = logo_path.read_bytes()
            b64 = base64.b64encode(data).decode("ascii")
            return f"data:image/svg+xml;base64,{b64}"

        buf = io.BytesIO()
        if img.mode in ("RGBA", "P"):
            img.save(buf, format="PNG", optimize=True)
            mime = "image/png"
        else:
            img = img.convert("RGB")
            img.save(buf, format="JPEG", optimize=True, quality=90)
            mime = "image/jpeg"
        b64 = base64.b64encode(buf.getvalue()).decode("ascii")
        return f"data:{mime};base64,{b64}"
    except Exception as e:
        print(f"  警告: logo 处理失败: {e}")
        return None


# ---------- 模板渲染 ----------

def inline_css():
    """读取主题 CSS。"""
    css_path = THEMES_DIR / "report.css"
    if css_path.exists():
        return css_path.read_text(encoding="utf-8")
    return ""


def render_template(config, content, toc, logo_data_uri=None):
    env = Environment(
        loader=FileSystemLoader(TEMPLATES_DIR),
        autoescape=select_autoescape(["html", "xml"]),
    )
    template = env.get_template("report.html")
    css = inline_css()
    return template.render(
        title=config["title"],
        subtitle=config["subtitle"],
        author=config["author"],
        version=config["version"],
        doc_type=config["doc_type"],
        logo=logo_data_uri,
        date=config["date"],
        content=content,
        toc=toc,
        css=css,
    )


# ---------- 检查脚本 ----------

def run_checks(checks):
    issues = []
    for name, module in checks.items():
        func = getattr(module, "check", None)
        if not callable(func):
            continue
        try:
            result = func(str(PROJECT_ROOT))
            if result:
                issues.extend(result)
        except Exception as e:
            print(f"  警告: 检查脚本 {name} 执行失败: {e}")
    if issues:
        print("检查发现问题：")
        for item in issues:
            print(f"  - {item}")


# ---------- 主入口 ----------

def main():
    print("=== Docs as Code 构建 ===")
    print()

    # 0. 加载配置
    config = load_config()
    print(f"[配置] {config['title']} — {config['version']}")
    print(f"  章节数: {len(config.get('chapters', []))}")

    # 0a. 加载插件
    print("[插件] 加载扩展模块 ...")
    renderers = load_plugin_modules(RENDERERS_DIR)
    checks = load_plugin_modules(CHECKS_DIR)

    # 0b. 运行检查
    if config.get("checks", True) and checks:
        print("[检查] 运行检查脚本 ...")
        run_checks(checks)

    # 1. 逐章节处理
    chapters = config.get("chapters", [])
    if not chapters:
        sys.exit("错误: project.yaml 中未定义 chapters 列表")

    html_parts = []
    for ch_file in chapters:
        ch_path = CHAPTERS_DIR / ch_file
        if not ch_path.exists():
            print(f"  警告: 章节文件不存在，跳过: {ch_file}")
            continue

        text = ch_path.read_text(encoding="utf-8")

        # @import 处理
        text = process_imports(text, ch_path.parent, renderers)

        # 图片属性处理
        text = process_image_attrs(text)

        # KaTeX 预处理
        text = preprocess_katex(text)

        # 自定义代码块渲染
        text = render_custom_codeblocks_md(text, renderers)

        # Markdown → HTML
        chapter_html = markdown_to_html(text)

        # 图片 base64 内嵌
        base_dirs = [ch_path.parent, ASSETS_DIR, PROJECT_ROOT]
        chapter_html = embed_images(chapter_html, base_dirs)

        html_parts.append(chapter_html)

    full_html = "\n\n".join(html_parts)
    print(f"  HTML 总字符数: {len(full_html)}")

    # 2. 注入锚点 + 提取目录
    full_html = inject_anchors(full_html)
    toc = generate_toc(full_html)
    print(f"  目录条目数: {len(toc)}")

    # 3. Logo 处理
    logo_path = find_logo(config.get("logo", ""))
    logo_data_uri = embed_logo(logo_path) if logo_path else None
    if logo_data_uri:
        print("  logo: 已内嵌")
    else:
        print("  logo: 未配置，封面将不显示 logo")

    # 4. 模板渲染 + 输出
    output_filename = re.sub(r'[^\w\-.]', '_', config["title"]) + ".html"
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    output_path = OUTPUT_DIR / output_filename

    final_html = render_template(config, full_html, toc, logo_data_uri)
    output_path.write_text(final_html, encoding="utf-8")
    print(f"\n构建完成：{output_path}")
    print(f"文件大小: {output_path.stat().st_size:,} 字节")


if __name__ == "__main__":
    main()