#!/usr/bin/env python3
"""
Trích xuất text từ JSON Cloud Vision → Markdown (V7 — Y-bucket + Header/Footer Cleanup)

Kết hợp:
- V6: Y-bucket grouping → layout đẹp, sort X trong mỗi dòng
- V5: Header/Footer detection bằng tọa độ Y → xóa header zone nếu có số trang

Nguyên lý cleanup (từ V5):
1. Quét tất cả block có y_min < 0.15 → header zone
2. Nếu header zone có số trang (Myanmar/Latin) → đánh dấu toàn bộ zone
3. Lọc fragment: bỏ fragment có y_min <= HEADER_Y_MAX và y_min > 0.85
4. Xử lý tự nhiên block merge — chỉ xóa dòng header, giữ body

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

import os
import sys
import json
import glob
import re

# ─── Constants ───────────────────────────────────────────────────────────────

SCANNER_PATTERNS = [
    re.compile(r'scanned\s+with', re.IGNORECASE),
    re.compile(r'cam\s*scanner', re.IGNORECASE),
]

BUCKET_SIZE = 0.022
COLUMN_GAP_THRESHOLD = 0.15
WORD_GAP_THRESHOLD = 0.03

# Cleanup thresholds (tuned for Pali-Taykha)
HEADER_DETECT_Y = 0.15   # detection: y_min <= này → header zone candidate
HEADER_CLEAN_Y = 0.15    # fallback fixed threshold (used when no pure header blocks found)
HEADER_MARGIN = 0.006    # safety margin above pure header block max_y (tight — thà sót header hơn xóa body)
PURE_HEADER_MAX_H = 0.12 # max block height to be considered "pure header" (short blocks)
# Footer: KHÔNG filter — artifact footer dễ phát hiện & xóa ở bước editor

PAGE_NUM_RE = re.compile(r'^[၀-၉ဝ0-9]{1,3}$')


# ─── Text Extraction (from V6) ──────────────────────────────────────────────

def extract_text(obj):
    """Extract text from a paragraph/word, preserving spaces and line breaks."""
    text = ''
    for word in obj.get('words', []):
        for sym in word.get('symbols', []):
            text += sym.get('text', '')
            prop = sym.get('property', {})
            if prop.get('detectedBreak'):
                bt = prop['detectedBreak']['type']
                if bt in ('SPACE', 'SURE_SPACE'):
                    text += ' '
                elif bt in ('EOL_SURE_SPACE', 'LINE_BREAK'):
                    text += '\n'
    return text


def get_bbox(obj):
    """Get (x_min, y_min, x_max, y_max) from normalizedVertices."""
    verts = obj.get('boundingBox', {}).get('normalizedVertices', [])
    if not verts:
        return (0, 0, 0, 0)
    return (
        min(v.get('x', 0) for v in verts),
        min(v.get('y', 0) for v in verts),
        max(v.get('x', 0) for v in verts),
        max(v.get('y', 0) for v in verts),
    )


def is_scanner(text):
    for pat in SCANNER_PATTERNS:
        if pat.search(text):
            return True
    return False


# ─── Header/Footer Detection (from V5) ──────────────────────────────────────

def get_block_bbox(block):
    """Get bbox from a block's normalizedVertices."""
    bbox = block.get("boundingBox", {}).get("normalizedVertices", [{}] * 4)
    ys = [v.get("y", 0) for v in bbox if "y" in v]
    xs = [v.get("x", 0) for v in bbox if "x" in v]
    if not ys or not xs:
        return None
    return min(ys), max(ys), min(xs), max(xs)


def get_block_text(block):
    """Extract concatenated text from a block (for header detection only)."""
    parts = []
    for para in block.get("paragraphs", []):
        for word in para.get("words", []):
            for symbol in word.get("symbols", []):
                parts.append(symbol.get("text", ""))
    return "".join(parts).strip()


def has_page_number(header_blocks):
    """Check if any header block contains a page number."""
    for y_min, y_max, x_min, x_max, text in header_blocks:
        if PAGE_NUM_RE.match(text):
            return True
        # Single-char at edge (might be OCR-damaged page number)
        if len(text) <= 2 and (x_max < 0.30 or x_min > 0.70):
            if re.match(r'^[A-Za-z0-9၀-၉ဝ]+$', text):
                return True
    return False


def compute_adaptive_clean_y(header_blocks):
    """
    Compute per-page adaptive header clean threshold.
    Header = 2-3 SHORT corner-positioned blocks:
      1. Left corner (x_max < 0.40) — page number or chapter name
      2. Right corner (x_min > 0.60) — chapter name or page number  
      3. URL (http/www) at very top (y_max < 0.03)
    Mixed blocks (tall blocks spanning header+body) are handled separately
    in the fragment filter — their content is always preserved.
    """
    corner_blocks = []
    for y_min, y_max, x_min, x_max, text in header_blocks:
        h = y_max - y_min
        if h >= PURE_HEADER_MAX_H:
            continue  # mixed block — handled in fragment filter
        at_left = x_max < 0.40
        at_right = x_min > 0.60
        is_url = ('http' in text.lower() or 'www.' in text.lower()) and y_max < 0.03
        if is_url:
            corner_blocks.append((y_min, y_max))
        elif at_left:
            # Left corner: page number (≤2 chars matching PAGE_NUM) or ≥3 chars
            if len(text) <= 2:
                if PAGE_NUM_RE.match(text):
                    corner_blocks.append((y_min, y_max))
            else:
                corner_blocks.append((y_min, y_max))
        elif at_right:
            corner_blocks.append((y_min, y_max))
    
    if 1 <= len(corner_blocks) <= 3:
        max_header_y = max(y_max for _, y_max in corner_blocks)
        return min(max_header_y + HEADER_MARGIN, HEADER_DETECT_Y)
    return HEADER_CLEAN_Y  # fallback — let Phase 2 handle


def detect_header_footer(blocks):
    """
    Detect header blocks. Returns:
      should_clean: bool — True if header zone has page numbers → clean
      header_clean_y: float — per-page adaptive clean threshold
      header_bboxes: list of (y_min, y_max) for all header zone blocks
      header_texts: set of header block texts (for logging)
      footer_texts: empty set (footer detection removed — editor handles it)
    """
    header_blocks = []  # (y_min, y_max, x_min, x_max, text)
    footer_texts = set()

    for block in blocks:
        if block.get('blockType') not in (None, 'TEXT'):
            continue
        bbox = get_block_bbox(block)
        if not bbox:
            continue
        y_min, y_max, x_min, x_max = bbox
        text = get_block_text(block)
        if not text:
            continue

        if y_min <= HEADER_DETECT_Y:
            header_blocks.append((y_min, y_max, x_min, x_max, text))

    if not header_blocks:
        return False, HEADER_CLEAN_Y, [], set(), footer_texts

    should_clean = has_page_number(header_blocks)
    header_clean_y = compute_adaptive_clean_y(header_blocks)
    header_texts = {text for _, _, _, _, text in header_blocks}
    header_bboxes = [(y_min, y_max) for y_min, y_max, _, _, _ in header_blocks]

    return should_clean, header_clean_y, header_bboxes, header_texts, footer_texts


# ─── Y-Bucket Processing (from V6) ──────────────────────────────────────────

def split_into_lines(paragraphs):
    """Split multi-line paragraphs into individual line fragments.
    Returns list of (x_min, y_min, x_max, y_max, text_line).
    """
    result = []
    for x_min, y_min, x_max, y_max, text in paragraphs:
        lines = text.split('\n')
        lines = [l.strip() for l in lines if l.strip()]
        if not lines:
            continue

        if len(lines) == 1:
            result.append((x_min, y_min, x_max, y_max, lines[0]))
        else:
            num_lines = len(lines)
            line_h = (y_max - y_min) / num_lines
            for i, line in enumerate(lines):
                ly_min = y_min + i * line_h
                ly_max = ly_min + line_h
                result.append((x_min, ly_min, x_max, ly_max, line))
    return result


def y_bucket_key(y_min, y_max):
    """Map a Y range to a bucket key based on its midpoint."""
    y_mid = (y_min + y_max) / 2
    return round(y_mid / BUCKET_SIZE) * BUCKET_SIZE


def bucket_to_text(fragments):
    """Convert fragments in a single Y-bucket to a text line.
    Fragments are sorted by X (left→right) with column/word spacing.
    """
    if not fragments:
        return ''

    sorted_frags = sorted(fragments, key=lambda f: f[0])
    parts = []
    prev_x_max = 0

    for i, (x_min, y_min, x_max, y_max, text) in enumerate(sorted_frags):
        if i > 0:
            gap = x_min - prev_x_max
            if gap > COLUMN_GAP_THRESHOLD:
                parts.append('    ')
            elif gap > WORD_GAP_THRESHOLD:
                parts.append('  ')
            else:
                parts.append(' ')
        parts.append(text)
        prev_x_max = x_max

    return ''.join(parts)


# ─── Page Processing ────────────────────────────────────────────────────────

def page_to_text(page_data, page_num, debug=False):
    """Convert a single Cloud Vision page to clean Markdown text."""
    blocks = page_data.get('blocks', [])

    # ── Phase 0: Header/Footer Detection ──
    should_clean, header_clean_y, header_bboxes, header_texts, footer_texts = detect_header_footer(blocks)
    
    # Collect mixed block Y-ranges (blocks that span header+body → don't clean their content)
    mixed_block_ranges = []
    for block in blocks:
        if block.get('blockType') not in (None, 'TEXT'):
            continue
        bbox = get_block_bbox(block)
        if not bbox:
            continue
        y_min, y_max, x_min, x_max = bbox
        if y_min <= HEADER_DETECT_Y and (y_max - y_min) >= PURE_HEADER_MAX_H:
            mixed_block_ranges.append((y_min, y_max))

    # ── Phase 1: Extract all paragraphs ──
    all_paragraphs = []
    for block in blocks:
        if block.get('blockType') not in (None, 'TEXT'):
            continue
        for para in block.get('paragraphs', []):
            text = extract_text(para)
            text = text.strip()
            if not text:
                continue

            x_min, y_min, x_max, y_max = get_bbox(para)

            if is_scanner(text):
                continue

            all_paragraphs.append((x_min, y_min, x_max, y_max, text))

    if not all_paragraphs:
        return f"## PAGE {page_num}\n\n*(trang trống)*\n", \
               should_clean, header_clean_y, header_texts, footer_texts, 0

    # ── Phase 2: Split into line fragments ──
    fragments = split_into_lines(all_paragraphs)
    total_frags = len(fragments)

    # ── Phase 3: Filter header/footer fragments ──
    removed_count = 0
    if should_clean:
        kept = []
        for frag in fragments:
            y_min, y_max = frag[1], frag[3]
            # Check if fragment belongs to a mixed block → always keep
            in_mixed = any(my_min <= y_min <= my_max for my_min, my_max in mixed_block_ranges)
            if in_mixed:
                kept.append(frag)
                continue
            # Header: remove if fragment starts <= per-page adaptive header_clean_y
            if y_min > header_clean_y:
                kept.append(frag)
            else:
                removed_count += 1
        fragments = kept

    if not fragments:
        return f"## PAGE {page_num}\n\n*(trang trống — header/footer removed)*\n", \
               should_clean, header_clean_y, header_texts, footer_texts, removed_count

    # ── Phase 4: Group into Y-buckets ──
    buckets = {}
    for frag in fragments:
        key = y_bucket_key(frag[1], frag[3])
        if key not in buckets:
            buckets[key] = []
        buckets[key].append(frag)

    # ── Phase 5: Build output ──
    output_lines = [f"## PAGE {page_num}"]
    if debug:
        clean_info = f"adaptive_y={header_clean_y:.3f}" if should_clean else "no_clean"
        output_lines.append(
            f"<!-- {total_frags}→{len(fragments)} fragments, "
            f"{len(buckets)} buckets, {removed_count} removed | {clean_info} -->"
        )
    output_lines.append("")

    for bk in sorted(buckets.keys()):
        line = bucket_to_text(buckets[bk])
        if line:
            output_lines.append(line)

    output_lines.append("")
    return '\n'.join(output_lines), should_clean, header_clean_y, header_texts, footer_texts, removed_count


def json_to_markdown(json_path, output_dir, start_page, debug=False):
    """Process one JSON file → clean Markdown file."""
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    basename = os.path.splitext(os.path.basename(json_path))[0]
    out_path = os.path.join(output_dir, f"{basename}.md")

    page_texts = []
    for i, response in enumerate(data.get('responses', [])):
        page_num = start_page + i
        fta = response.get('fullTextAnnotation', {})

        if 'pages' not in fta:
            raw_text = fta.get('text', '').strip()
            page_texts.append(
                f"## PAGE {page_num}\n\n{raw_text if raw_text else '*(trang trống)*'}\n"
            )
            continue

        pages = fta.get('pages', [])
        if pages:
            page_text, should_clean, clean_y, h_texts, f_texts, removed = \
                page_to_text(pages[0], page_num, debug=debug)

            # Log
            nonlocal_log = {'cleaned': False, 'skipped': False, 'removed': 0}
            if h_texts or f_texts:
                if should_clean:
                    status = '✅'
                    nonlocal_log['cleaned'] = True
                else:
                    status = '⏭️ '
                    nonlocal_log['skipped'] = True
                log_texts = sorted(h_texts | f_texts)
                if not should_clean:
                    log_texts.append('(no page number → skip)')
                else:
                    log_texts.insert(0, f'[clean_y={clean_y:.3f}]')
                print(f"   Page {page_num:2d}  {status}  xóa: {log_texts}")
            nonlocal_log['removed'] = removed

            page_texts.append((page_text, nonlocal_log))
        else:
            page_texts.append(f"## PAGE {page_num}\n\n*(trang trống)*\n")

    os.makedirs(output_dir, exist_ok=True)
    with open(out_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(pt[0] if isinstance(pt, tuple) else pt for pt in page_texts))

    stats = {
        'pages': len(page_texts),
        'cleaned': sum(1 for pt in page_texts if isinstance(pt, tuple) and pt[1]['cleaned']),
        'skipped': sum(1 for pt in page_texts if isinstance(pt, tuple) and pt[1]['skipped']),
        'removed': sum(pt[1]['removed'] for pt in page_texts if isinstance(pt, tuple)),
    }
    return out_path, stats


# ─── Main ────────────────────────────────────────────────────────────────────

def main():
    if len(sys.argv) < 3:
        print("Usage: python3 json_to_markdown_v7.py <json_dir> <output_dir> [--dry-run] [--debug]")
        print()
        print("  json_dir    — thư mục chứa các file JSON từ Cloud Vision")
        print("  output_dir  — thư mục xuất file Markdown")
        print("  --dry-run   — preview 25 dòng đầu mỗi file, không lưu thật")
        print("  --debug     — hiển thị fragment/bucket count và log cleanup")
        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

    json_files = sorted(
        glob.glob(os.path.join(json_dir, '*.json')),
        key=lambda f: int(re.search(r'output-(\d+)-to', os.path.basename(f)).group(1))
    )

    if not json_files:
        print(f"❌ Không tìm thấy file JSON nào trong {json_dir}")
        sys.exit(1)

    print(f"🔧 V7: Y-bucket sorting + Header/Footer cleanup")
    print(f"📄 {len(json_files)} file JSON tìm thấy\n")

    total_pages = total_removed = pages_cleaned = pages_skipped = 0
    all_md = []

    for jf in json_files:
        out_path, stats = json_to_markdown(jf, output_dir, total_pages + 1, debug=debug)
        total_pages += stats['pages']
        total_removed += stats['removed']
        pages_cleaned += stats['cleaned']
        pages_skipped += stats['skipped']

        if dry_run:
            with open(out_path, 'r') as f:
                preview = ''.join(f.readline() for _ in range(30))
            print(f"\n{'='*60}")
            print(f"📄 {os.path.basename(out_path)} ({stats['pages']} trang) — PREVIEW:")
            print(f"{'='*60}")
            print(preview)
            print(f"... (saved at {out_path})")
        else:
            print(f"   ✅ {os.path.basename(out_path)} ({stats['pages']} trang)")

        # Gather merged content
        with open(out_path, 'r') as f:
            all_md.append(f.read())

    # Write merged full file
    if not dry_run:
        merged_path = os.path.join(output_dir, "Pali-Taykha-full.md")
        with open(merged_path, 'w', encoding='utf-8') as f:
            f.write("# သောဠသမကျမ်း — OCR Cleaned (V7: Y-bucket + Header/Footer)\n\n")
            f.write('\n'.join(all_md))
            f.write(f"\n\n---\n**Tổng:** {total_pages} trang | ")
            f.write(f"**Header/Footer:** {pages_cleaned} cleaned | ")
            f.write(f"**Fragments removed:** {total_removed}\n")
        print(f"\n📄 Merged: {merged_path} ({os.path.getsize(merged_path):,} bytes)")

    if pages_skipped:
        print(f"   ⏭️  {pages_skipped} trang có header zone nhưng không có số trang → bỏ qua")

    print(f"\n{'='*60}")
    print(f"📊 {total_pages} trang | {pages_cleaned} có header → xóa {total_removed} fragments")
    print(f"🎉 Done! → {os.path.abspath(output_dir)}")
    print(f"{'='*60}")


if __name__ == '__main__':
    main()