#!/usr/bin/env python3
"""
Extract raw text from Cloud Vision JSON output and save as Markdown files.
Each JSON file (5 pages) → one .md file, with ## PAGE X headers.
"""
import os
import json
import glob
import sys

def extract_text_from_json(json_path):
    """Extract raw text from a Cloud Vision JSON output file."""
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    pages_text = []

    # Cloud Vision async output structure: { "responses": [ { "fullTextAnnotation": { "text": "..." } }, ... ] }
    responses = data.get('responses', [])
    for resp in responses:
        annotation = resp.get('fullTextAnnotation', {})
        text = annotation.get('text', '').strip()
        pages_text.append(text)

    return pages_text


def main():
    raw_dir = sys.argv[1] if len(sys.argv) > 1 else '001-vi-dieu-phap-giang-giai/ocr/raw/'
    out_dir = sys.argv[2] if len(sys.argv) > 2 else '001-vi-dieu-phap-giang-giai/ocr/markdown/'

    os.makedirs(out_dir, exist_ok=True)

    # Find and sort JSON files
    json_files = sorted(
        glob.glob(os.path.join(raw_dir, 'output-*.json')),
        key=lambda f: int(os.path.basename(f).replace('output-', '').split('-to-')[0])
    )

    if not json_files:
        print(f"❌ No output-*.json files found in {raw_dir}")
        sys.exit(1)

    print(f"Found {len(json_files)} JSON files in {raw_dir}")

    total_pages = 0
    for json_file in json_files:
        basename = os.path.basename(json_file)
        # Parse page range from filename: e.g., "output-1-to-5.json" → (1, 5)
        parts = basename.replace('output-', '').replace('.json', '').split('-to-')
        start_page = int(parts[0])
        end_page = int(parts[1])

        # Map to original PDF page numbers (start at 181)
        pdf_start = 180 + start_page
        pdf_end = 180 + end_page

        # Extract raw text
        pages_text = extract_text_from_json(json_file)
        num_pages = len(pages_text)

        # Build markdown
        md_name = f"vdp-{pdf_start}-{pdf_end}.md"
        md_path = os.path.join(out_dir, md_name)

        lines = []
        for i, text in enumerate(pages_text):
            page_num = pdf_start + i
            lines.append(f"## PAGE {page_num}")
            lines.append("")
            lines.append(text)
            lines.append("")

        with open(md_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(lines))

        total_pages += num_pages
        print(f"   ✅ {basename} → {md_name} ({num_pages} pages)")

    print(f"\n🎉 Done! {total_pages} pages → {len(json_files)} markdown files in {os.path.abspath(out_dir)}")


if __name__ == '__main__':
    main()