from __future__ import annotations

import argparse
import html
import json
import os
import re
import zipfile
from pathlib import Path


P_RE = re.compile(r"<p\b([^>]*)>(.*?)</p>", re.S | re.I)
H_RE = re.compile(r"<h([1-6])\b([^>]*)>(.*?)</h\1>", re.S | re.I)
RUBY_RE = re.compile(r"<ruby\b[^>]*>.*?</ruby>", re.S | re.I)
RB_RE = re.compile(r"<rb\b[^>]*>(.*?)</rb>", re.S | re.I)
RT_RE = re.compile(r"<rt\b[^>]*>(.*?)</rt>", re.S | re.I)
TAG_RE = re.compile(r"<[^>]+>")
NAV_LABEL_TRANSLATIONS = {
    "\u76ee\u6b21": "\u76ee\u5f55",
    "\u5c11\u5e74\u671f\u3000\u5341\u4e8c\u6b73\u306e\u521d\u590f": "\u5c11\u5e74\u671f \u5341\u4e8c\u5c81\u521d\u590f",
    "\u5c11\u5e74\u671f\u3000\u5341\u4e8c\u6b73\u306e\u76db\u590f": "\u5c11\u5e74\u671f \u5341\u4e8c\u5c81\u76db\u590f",
    "\u5c11\u5e74\u671f\u3000\u5341\u4e09\u6b73\u306e\u79cb": "\u5c11\u5e74\u671f \u5341\u4e09\u5c81\u4e4b\u79cb",
    "\u7d42\u7ae0": "\u7ec8\u7ae0",
    "\u30d8\u30f3\u30c0\u30fc\u30bd\u30f3\u30b9\u30b1\u30fc\u30eb1.0\u3000Ver0.3": "\u4ea8\u5fb7\u68ee\u91cf\u88681.0 Ver0.3",
    "\u5965\u4ed8": "\u7248\u6743\u9875",
    "\u3042\u3068\u304c\u304d": "\u540e\u8bb0",
}


def extract_paragraphs(html: str) -> list[str]:
    return [m.group(2) for m in P_RE.finditer(html)]


def extract_headings(text: str) -> list[str]:
    return [m.group(3) for m in H_RE.finditer(text)]


def strip_tags(fragment: str) -> str:
    return html.unescape(TAG_RE.sub("", fragment)).strip()


def translate_known_labels(text: str) -> str:
    result = text
    for jp, cn in NAV_LABEL_TRANSLATIONS.items():
        result = result.replace(jp, cn)
    return result


def clean_term_value(value: str) -> str:
    return str(value).split("#", 1)[0].strip()


def load_terms(novel_root: Path) -> dict[str, str]:
    candidates = [novel_root / "mingcibiao.json", novel_root.parent / "mingcibiao.json"]
    term_path = next((path for path in candidates if path.exists()), None)
    if term_path is None:
        return {}
    raw = term_path.read_text(encoding="utf-8-sig")
    raw = re.sub(r",\s*([}\]])", r"\1", raw)
    return {str(key): clean_term_value(value) for key, value in json.loads(raw).items()}


def split_ruby_term(value: str) -> tuple[str, str] | None:
    match = re.fullmatch(r"(.+?)\(([^()<>]+)\)", clean_term_value(value))
    if not match:
        return None
    rb, rt = match.group(1).strip(), match.group(2).strip()
    if not rb or not rt:
        return None
    return rb, rt


def source_ruby_items(fragment: str) -> list[tuple[str, str]]:
    items = []
    for ruby in RUBY_RE.findall(fragment):
        rb_parts = [strip_tags(part) for part in RB_RE.findall(ruby)]
        rt_parts = [strip_tags(part) for part in RT_RE.findall(ruby)]
        rb = "".join(rb_parts) or strip_tags(re.sub(r"<rt\b[^>]*>.*?</rt>", "", ruby, flags=re.S | re.I))
        rt = "/".join(part for part in rt_parts if part)
        if rb and rt:
            items.append((rb, rt))
    return items


def translated_ruby_for(rb: str, rt: str, terms: dict[str, str]) -> tuple[str, str] | None:
    for key in (f"{rb}({rt})", rb, rt):
        value = terms.get(key)
        if not value:
            continue
        parsed = split_ruby_term(value)
        if parsed:
            return parsed
    return None


def restore_translated_ruby(source_inner: str, translated_inner: str, terms: dict[str, str]) -> str:
    if not terms:
        return translated_inner
    restored = translated_inner
    for rb, rt in source_ruby_items(source_inner):
        parsed = translated_ruby_for(rb, rt, terms)
        if not parsed:
            continue
        cn_rb, cn_rt = parsed
        literal = f"{cn_rb}({cn_rt})"
        if literal not in restored:
            continue
        ruby_html = f"<ruby><rb>{html.escape(cn_rb)}</rb><rt>{html.escape(cn_rt)}</rt></ruby>"
        restored = restored.replace(literal, ruby_html, 1)
    return restored


def set_doc_lang_and_title(text: str, title: str, lang: str = "zh-CN") -> str:
    text = re.sub(r"<title>.*?</title>", f"<title>{title}</title>", text, count=1, flags=re.S)
    text = re.sub(r'xml:lang="[^"]+"', f'xml:lang="{lang}"', text, count=1)
    if 'lang="' in text:
        text = re.sub(r'lang="[^"]+"', f'lang="{lang}"', text, count=1)
    else:
        text = text.replace('<html ', f'<html lang="{lang}" ', 1)
    text = re.sub(r'class="([^"]*)\bvrtl\b([^"]*)"', lambda m: f'class="{(m.group(1) + " hltr " + m.group(2)).strip()}"', text, count=1)
    return text


def add_gray_style(attrs: str) -> str:
    if re.search(r"\bstyle\s*=", attrs, re.I):
        if re.search(r'style="[^"]*"', attrs, re.I):
            def repl(m: re.Match[str]) -> str:
                body = m.group(1)
                body = re.sub(r"\bcolor\s*:\s*[^;]+;?", "", body, flags=re.I).strip()
                if body and not body.endswith(";"):
                    body += ";"
                body = (body + " color:#808080;").strip()
                return f'style="{body}"'

            return re.sub(r'style="([^"]*)"', repl, attrs, count=1, flags=re.I)
        if re.search(r"style='[^']*'", attrs, re.I):
            def repl2(m: re.Match[str]) -> str:
                body = m.group(1)
                body = re.sub(r"\bcolor\s*:\s*[^;]+;?", "", body, flags=re.I).strip()
                if body and not body.endswith(";"):
                    body += ";"
                body = (body + " color:#808080;").strip()
                return f"style='{body}'"

            return re.sub(r"style='([^']*)'", repl2, attrs, count=1, flags=re.I)
    return attrs + ' style="color:#808080;"'


def transform_xhtml(
    source_text: str,
    translated_text: str | None,
    bilingual: bool,
    title: str,
    terms: dict[str, str],
) -> str:
    source_text = set_doc_lang_and_title(source_text, title)
    translated_headings = extract_headings(translated_text or "")
    if translated_headings:
        heading_iter = iter(translated_headings)

        def replace_heading(match: re.Match[str]) -> str:
            try:
                inner = next(heading_iter)
            except StopIteration:
                inner = translate_known_labels(match.group(3))
            return f"<h{match.group(1)}{match.group(2)}>{inner}</h{match.group(1)}>"

        source_text = H_RE.sub(replace_heading, source_text)
    else:
        source_text = translate_known_labels(source_text)
    translated_paragraphs = extract_paragraphs(translated_text or "")
    matches = list(P_RE.finditer(source_text))

    pieces: list[str] = []
    last = 0
    for index, match in enumerate(matches):
        pieces.append(source_text[last:match.start()])
        attrs = match.group(1)
        source_inner = match.group(2)
        translated_inner = translated_paragraphs[index] if index < len(translated_paragraphs) else source_inner
        translated_inner = restore_translated_ruby(source_inner, translated_inner, terms)

        if bilingual:
            gray_attrs = add_gray_style(attrs)
            pieces.append(f"<p{gray_attrs}>{source_inner}</p>")
            if strip_tags(translated_inner):
                pieces.append(f"<p{attrs}>{translated_inner}</p>")
        else:
            inner = translated_inner if translated_inner.strip() else source_inner
            pieces.append(f"<p{attrs}>{inner}</p>")

        last = match.end()

    pieces.append(source_text[last:])
    return "".join(pieces)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Build translated and bilingual EPUB variants for any novel volume.")
    parser.add_argument(
        "--mode",
        choices=("translated", "bilingual", "both"),
        default="both",
        help="Which deliverable(s) to build.",
    )
    parser.add_argument(
        "--source-epub",
        default=os.environ.get("SOURCE_EPUB"),
        help="Path to the source EPUB.",
    )
    parser.add_argument(
        "--novel-root",
        default=os.environ.get("NOVEL_ROOT"),
        help="Path to the novel project root.",
    )
    parser.add_argument(
        "--out-dir",
        default=os.environ.get("OUT_DIR"),
        help="Directory where deliverables are written.",
    )
    parser.add_argument(
        "--title",
        default=os.environ.get("EPUB_TITLE"),
        help="Optional display title. If omitted, the source EPUB title is used.",
    )
    return parser.parse_args()


def load_manifest(novel_root: Path) -> dict:
    manifest_path = novel_root / "text" / "_epub_import_manifest.json"
    if manifest_path.exists():
        return json.loads(manifest_path.read_text(encoding="utf-8"))
    final_dir = novel_root / "final_text"
    files = [
        {
            "source_epub_path": f"text/{path.name}",
            "output": path.name,
        }
        for path in sorted(final_dir.glob("part*.html"))
    ]
    return {
        "metadata": {"title": novel_root.name},
        "files": files,
    }


def output_title(manifest: dict, override: str | None) -> str:
    return override or manifest["metadata"]["title"]


def output_filename(title: str, suffix: str) -> str:
    safe_title = re.sub(r'[\\/:*?"<>|]+', "_", title).strip()
    return f"{safe_title}_{suffix}.epub"


def main() -> int:
    args = parse_args()
    if not args.source_epub or not args.novel_root:
        raise SystemExit("SOURCE_EPUB and NOVEL_ROOT must be provided via args or environment.")

    source_epub = Path(args.source_epub)
    novel_root = Path(args.novel_root)
    output_dir = Path(args.out_dir) if args.out_dir else novel_root / "deliverables"
    output_dir.mkdir(parents=True, exist_ok=True)

    manifest = load_manifest(novel_root)
    terms = load_terms(novel_root)
    book_title = output_title(manifest, args.title)
    source_to_output = {item["source_epub_path"]: item["output"] for item in manifest["files"]}
    source_output_to_xhtml = {item["output"]: item["source_epub_path"] for item in manifest["files"]}

    final_dir = novel_root / "final_text"
    translated_by_output = {p.name: p for p in final_dir.glob("part*.html")}

    variants = [
        {
            "name": output_filename(book_title, "简中译本"),
            "title": f"{book_title}（简中译本）",
            "bilingual": False,
            "mode": "translated",
        },
        {
            "name": output_filename(book_title, "中日双语版"),
            "title": f"{book_title}（中日双语版）",
            "bilingual": True,
            "mode": "bilingual",
        },
    ]
    if args.mode != "both":
        variants = [item for item in variants if item["mode"] == args.mode]

    with zipfile.ZipFile(source_epub, "r") as src_zip:
        entries = src_zip.infolist()
        source_bytes = {entry.filename: src_zip.read(entry.filename) for entry in entries}

    for variant in variants:
        out_path = output_dir / variant["name"]
        if out_path.exists():
            out_path.unlink()

        with zipfile.ZipFile(out_path, "w") as out_zip:
            out_zip.writestr("mimetype", b"application/epub+zip", compress_type=zipfile.ZIP_STORED)
            for entry in entries:
                name = entry.filename
                if name == "mimetype":
                    continue

                data = source_bytes[name]
                if name.lower().endswith(".opf"):
                    text = data.decode("utf-8")
                    text = re.sub(
                        r"<dc:title>.*?</dc:title>",
                        f"<dc:title>{variant['title']}</dc:title>",
                        text,
                        count=1,
                        flags=re.S,
                    )
                    text = re.sub(
                        r"<dc:language>.*?</dc:language>",
                        "<dc:language>zh-CN</dc:language>",
                        text,
                        count=1,
                        flags=re.S,
                    )
                    text = re.sub(
                        r'<meta\s+content="vertical-rl"\s+name="primary-writing-mode"\s*/>',
                        '<meta content="horizontal-lr" name="primary-writing-mode"/>',
                        text,
                        count=1,
                    )
                    text = re.sub(
                        r'page-progression-direction="rtl"',
                        'page-progression-direction="ltr"',
                        text,
                        count=1,
                    )
                    data = text.encode("utf-8")
                elif name.lower().endswith(".ncx"):
                    text = data.decode("utf-8")
                    text = re.sub(r'xml:lang="[^"]+"', 'xml:lang="zh-CN"', text, count=1)
                    text = re.sub(
                        r"<docTitle>\s*<text>.*?</text>\s*</docTitle>",
                        f"<docTitle><text>{variant['title']}</text></docTitle>",
                        text,
                        count=1,
                        flags=re.S,
                    )
                    text = translate_known_labels(text)
                    data = text.encode("utf-8")
                elif (
                    (name.startswith("item/xhtml/") and name.endswith(".xhtml"))
                    or name in source_to_output
                ):
                    manifest_source = source_output_to_xhtml.get("part" + name[len("item/xhtml/p-") : -len(".xhtml")] + ".html")
                    output_name = source_to_output.get(name) if name in source_to_output else None
                    if output_name is None and manifest_source is not None:
                        output_name = source_to_output.get(manifest_source)
                    if output_name:
                        translated_file = translated_by_output.get(output_name)
                        if translated_file and translated_file.exists():
                            source_text = data.decode("utf-8")
                            translated_text = translated_file.read_text(encoding="utf-8")
                            data = transform_xhtml(
                                source_text,
                                translated_text,
                                variant["bilingual"],
                                variant["title"],
                                terms,
                            ).encode("utf-8")

                compress_type = zipfile.ZIP_STORED if name == "mimetype" else zipfile.ZIP_DEFLATED
                out_zip.writestr(name, data, compress_type=compress_type)

        print(f"Wrote {out_path}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())