#!/usr/bin/env python3
"""
Script trích xuất văn bản thô từ JSON OCR và lưu thành Markdown.
Mỗi 5 trang = 1 file Markdown.
"""

import json
import os
import re
from pathlib import Path

RAW_DIR = Path("/home/tuan-nguyen/.openclaw/workspace/001-vi-dieu-phap-giang-giai/ocr/raw")
OUTPUT_DIR = Path("/home/tuan-nguyen/.openclaw/workspace/001-vi-dieu-phap-giang-giai/extracted/ver-2")

def extract_page_number_from_filename(filename):
    """Trích xuất số trang bắt đầu từ tên file, e.g. output-1-to-5.json -> 1"""
    match = re.search(r'output-(\d+)-to-', filename)
    if match:
        return int(match.group(1))
    return 0

def process_json_file(json_path):
    """Xử lý một file JSON, trả về list (page_number, text) tuples."""
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    results = []
    start_page = extract_page_number_from_filename(json_path.name)
    
    responses = data.get('responses', [])
    for i, response in enumerate(responses):
        page_num = start_page + i
        full_text = response.get('fullTextAnnotation', {}).get('text', '')
        results.append((page_num, full_text))
    
    return results

def main():
    # Tạo thư mục output
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    
    # Lấy tất cả file JSON, sắp xếp theo số trang bắt đầu
    json_files = sorted(RAW_DIR.glob("*.json"), key=lambda p: extract_page_number_from_filename(p.name))
    
    print(f"Tìm thấy {len(json_files)} file JSON")
    
    all_pages = []
    for json_file in json_files:
        print(f"Đang xử lý: {json_file.name}")
        pages = process_json_file(json_file)
        all_pages.extend(pages)
    
    print(f"Tổng cộng {len(all_pages)} trang")
    
    # Gộp mỗi 5 trang thành một file Markdown
    for i in range(0, len(all_pages), 5):
        batch = all_pages[i:i+5]
        start_page = batch[0][0]
        end_page = batch[-1][0]
        
        output_filename = f"pages-{start_page}-to-{end_page}.md"
        output_path = OUTPUT_DIR / output_filename
        
        content = []
        for page_num, text in batch:
            content.append(f"## PAGE {page_num}")
            content.append("")
            content.append(text)
            content.append("")
            content.append("---")
            content.append("")
        
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(content))
        
        print(f"Đã lưu: {output_filename} (trang {start_page}-{end_page})")
    
    print(f"\nHoàn thành! Tổng cộng {len(all_pages) // 5} file Markdown trong {OUTPUT_DIR}")

if __name__ == "__main__":
    main()