#!/usr/bin/env python3
"""
V8: JSON → Markdown với Y-bucket layout + URL-only cleanup
- Giữ toàn bộ body text, KHÔNG header detection dựa trên tọa độ Y
- Chỉ xóa blocks có chứa URL pattern (http://www.dhammadownload.com)

Usage:
  python3 json_to_markdown_v8.py <json_dir> <output_dir> [--dry-run] [--debug]
"""

import json, os, sys, re, glob
from collections import defaultdict

# ── Constants ──
BUCKET_SIZE = 0.022
COLUMN_GAP_THRESHOLD = 0.15
WORD_GAP_THRESHOLD = 0.03

# URL patterns cần xóa
URL_PATTERNS = [
    re.compile(r'http://www\.dhammadownload\.com', re.IGNORECASE),
    re.compile(r'dhammadownload', re.IGNORECASE),
]

# ── Helpers ──
def extract_text(block):
    """Extract full concatenated text from a block."""
    parts = []
    for para in block.get("paragraphs", []):
        for word in para.get("words", []):
            for sym in word.get("symbols", []):
                parts.append(sym.get("text", ""))
                if sym.get("property", {}).get("detectedBreak", {}).get("type") == "SPACE":
                    parts.append(" ")
            parts.append(" ")
    return "".join(parts).strip()


def contains_url(text):
    """Check if text matches any URL pattern."""
    return any(p.search(text) for p in URL_PATTERNS)


def get_bbox(block):
    """Get normalized bbox for a block."""
    bbox = block.get("boundingBox", {}).get("normalizedVertices", [])
    if bbox and len(bbox) >= 4:
        y_min = min(v.get("y", 0) for v in bbox)
        y_max = max(v.get("y", 0) for v in bbox)
        x_min = min(v.get("x", 0) for v in bbox)
        x_max = max(v.get("x", 0) for v in bbox)
        return y_min, y_max, x_min, x_max
    return 0, 0, 0, 0


def extract_fragments(block):
    """Split block into word-level fragments with positions."""
    fragments = []
    for para in block.get("paragraphs", []):
        for word in para.get("words", []):
            symbols = word.get("symbols", [])
            if not symbols:
                continue
            # Word text
            text_parts = []
            for s in symbols:
                text_parts.append(s.get("text", ""))
            text = "".join(text_parts)
            # Word bbox
            word_bbox = word.get("boundingBox", {}).get("normalizedVertices", [])
            if word_bbox and len(word_bbox) >= 4:
                wy = min(v.get("y", 0) for v in word_bbox)
                wx_min = min(v.get("x", 0) for v in word_bbox)
                wx_max = max(v.get("x", 0) for v in word_bbox)
            else:
                wy = 0
                wx_min = 0
                wx_max = 0

            # Check for space after word
            has_space = False
            last_sym = symbols[-1] if symbols else None
            if last_sym and last_sym.get("property", {}).get("detectedBreak", {}).get("type") == "SPACE":
                has_space = True

            fragments.append({
                "text": text,
                "y": wy,
                "x_min": wx_min,
                "x_max": wx_max,
                "space_after": has_space
            })
    return fragments


def page_to_text(page_data, page_num, debug=False):
    """Convert one page's blocks to markdown text (URL blocks removed)."""
    blocks = page_data.get("blocks", [])
    all_fragments = []
    url_removed = 0

    for block in blocks:
        text = extract_text(block)
        if contains_url(text):
            url_removed += 1
            continue  # Skip URL blocks entirely

        frags = extract_fragments(block)
        all_fragments.extend(frags)

    if not all_fragments:
        return f"## PAGE {page_num}\n\n*(trang trống)*\n", {
            "total_frags": 0, "kept": 0, "removed": url_removed, "buckets": 0
        }

    # Y-bucket grouping
    buckets = []
    sorted_frags = sorted(all_fragments, key=lambda f: (f["y"], f["x_min"]))

    for frag in sorted_frags:
        placed = False
        for bucket in buckets:
            # Check if this fragment belongs in this bucket (same Y-line)
            bucket_y = bucket["y_center"]
            if abs(frag["y"] - bucket_y) < BUCKET_SIZE:
                bucket["fragments"].append(frag)
                # Update bucket y_center
                all_y = [f["y"] for f in bucket["fragments"]]
                bucket["y_center"] = sum(all_y) / len(all_y)
                placed = True
                break
        if not placed:
            buckets.append({
                "y_center": frag["y"],
                "fragments": [frag]
            })

    # Sort buckets by Y
    buckets.sort(key=lambda b: b["y_center"])

    # Generate markdown lines with column-aware spacing
    output_lines = [f"## PAGE {page_num}"]

    for bi, bucket in enumerate(buckets):
        # Sort fragments in bucket by X
        bucket["fragments"].sort(key=lambda f: f["x_min"])
        line_parts = []
        last_x_end = 0

        for fi, frag in enumerate(bucket["fragments"]):
            if fi == 0:
                line_parts.append(frag["text"])
            else:
                gap = frag["x_min"] - last_x_end
                if gap >= COLUMN_GAP_THRESHOLD:
                    line_parts.append("    " + frag["text"])
                elif gap >= WORD_GAP_THRESHOLD:
                    line_parts.append("  " + frag["text"])
                else:
                    line_parts.append(" " + frag["text"])

            last_x_end = frag["x_max"]

        line = "".join(line_parts)
        if line.strip():
            output_lines.append(line)

        # Add blank line between visually separated buckets
        if bi < len(buckets) - 1:
            next_y = buckets[bi + 1]["y_center"]
            current_max_y = max(f["y"] for f in bucket["fragments"]) if bucket["fragments"] else bucket["y_center"]
            if next_y - current_max_y > BUCKET_SIZE * 2:
                # Paragraph break
                pass  # blank line already added by \n join

    text = "\n".join(output_lines) + "\n"
    return text, {
        "total_frags": len(all_fragments),
        "kept": len(all_fragments),
        "removed": url_removed,
        "buckets": len(buckets)
    }


# ── Main ──
def json_to_markdown(json_path, output_dir, start_page, debug=False):
    """Process one JSON file → Markdown."""
    with open(json_path) as f:
        data = json.load(f)

    pages = data.get("responses", [])
    page_texts = []
    total_stats = {"total_frags": 0, "kept": 0, "removed": 0, "buckets": 0}

    for i, page in enumerate(pages):
        annotation = page.get("fullTextAnnotation", {})
        pages_info = annotation.get("pages", [])

        if not pages_info:
            page_num = start_page + i
            page_texts.append(f"## PAGE {page_num}\n\n*(trang trống)*\n")
            continue

        page_num = start_page + i
        text, stats = page_to_text(pages_info[0], page_num, debug=debug)

        for k in total_stats:
            total_stats[k] += stats[k]

        # Add annotation comment
        comment = f"<!-- {stats['total_frags']}→{stats['kept']} fragments, {stats['buckets']} buckets, {stats['removed']} url removed | v8_url_only -->"
        text = text.replace(f"## PAGE {page_num}\n", f"## PAGE {page_num}\n{comment}\n", 1)

        page_texts.append(text)

        if debug:
            print(f"   Page {page_num:3d}  {stats['total_frags']}→{stats['kept']} frags  "
                  f"{stats['buckets']} buckets  url_removed={stats['removed']}")

    full_text = "\n".join(page_texts)

    # Write output file
    basename = os.path.splitext(os.path.basename(json_path))[0]
    out_path = os.path.join(output_dir, f"{basename}.md")
    with open(out_path, "w") as f:
        f.write(full_text)

    if debug:
        print(f"   ✅ {basename}.md ({len(full_text):,} bytes)")

    return out_path, total_stats


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python3 json_to_markdown_v8.py <json_dir> <output_dir> [--dry-run] [--debug]")
        sys.exit(1)

    json_dir = sys.argv[1]
    output_dir = sys.argv[2]
    dry_run = "--dry-run" in sys.argv
    debug = "--debug" in sys.argv

    os.makedirs(output_dir, exist_ok=True)

    # Natural sort by start page number (extract from filename like output-X-to-Y.json)
    def _start_page(fname):
        m = re.search(r'output-(\d+)-to-\d+', os.path.basename(fname))
        return int(m.group(1)) if m else 0
    json_files = sorted(glob.glob(os.path.join(json_dir, "*.json")), key=_start_page)
    if not json_files:
        print(f"❌ Không tìm thấy file JSON nào trong {json_dir}")
        sys.exit(1)

    print(f"🔧 V8: Y-bucket layout + URL-only cleanup (no header detection)")
    print(f"📄 {len(json_files)} file JSON tìm thấy\n")

    total_pages = 0
    grand_stats = {"total_frags": 0, "kept": 0, "removed": 0, "buckets": 0}
    merged_texts = []

    for jf in json_files:
        if dry_run:
            print(f"   [DRY-RUN] {os.path.basename(jf)}")
            # Quick page count
            with open(jf) as f:
                data = json.load(f)
            n_pages = len(data.get("responses", []))
            total_pages += n_pages
            continue

        out_path, stats = json_to_markdown(jf, output_dir, total_pages + 1, debug=debug)
        n_pages = len(json.load(open(jf)).get("responses", []))
        total_pages += n_pages

        for k in grand_stats:
            grand_stats[k] += stats[k]

        merged_texts.append(open(out_path).read())

    # Write merged full file
    if not dry_run and merged_texts:
        merged_path = os.path.join(output_dir, "Pali-Taykha-full.md")
        with open(merged_path, "w") as f:
            f.write("\n".join(merged_texts))
        print(f"\n📄 Merged: {merged_path} ({sum(len(t) for t in merged_texts):,} bytes)")

    print(f"\n{'='*60}")
    print(f"📊 {total_pages} trang | {grand_stats['removed']} URL blocks removed")
    if not dry_run:
        print(f"🎉 Done! → {output_dir}")
    else:
        print(f"🔍 DRY-RUN: {len(json_files)} files, ~{total_pages} pages")
    print(f"{'='*60}")