#!/usr/bin/env python3
"""
Tách các file Markdown 5-trang thành batch 3 trang.
Đọc từ thư mục extracted/, ghi ra thư mục mới (cùng extracted/ hoặc extracted-3pp/).

Usage:
  python3 scripts/split_to_3pp.py <input_dir> [output_dir]
"""
import os, sys, re, glob

def parse_page_content(filepath):
    """Parse a markdown file into a list of (page_num, content) tuples."""
    with open(filepath, 'r', encoding='utf-8') as f:
        text = f.read()
    
    pages = []
    # Split by ## PAGE N markers
    parts = re.split(r'(## PAGE \d+\n)', text)
    
    current_header = None
    for part in parts:
        if part.startswith('## PAGE '):
            current_header = part.strip()
        elif current_header and part.strip():
            page_num = int(re.search(r'## PAGE (\d+)', current_header).group(1))
            pages.append((page_num, current_header + '\n\n' + part.strip() + '\n'))
    
    return pages


def main():
    if len(sys.argv) < 2:
        print("Usage: python3 split_to_3pp.py <input_dir> [output_dir]")
        sys.exit(1)
    
    input_dir = sys.argv[1]
    output_dir = sys.argv[2] if len(sys.argv) > 2 else input_dir
    
    # Collect all pages from all files
    md_files = sorted(
        glob.glob(os.path.join(input_dir, 'output-*.md')),
        key=lambda f: int(re.search(r'output-(\d+)-to', os.path.basename(f)).group(1))
    )
    
    # Filter: only process files starting from page 51
    all_pages = []
    for mf in md_files:
        base = int(re.search(r'output-(\d+)-to', os.path.basename(mf)).group(1))
        if base < 51:
            continue  # Skip already-processed batches 1-50
        pages = parse_page_content(mf)
        all_pages.extend(pages)
    
    all_pages.sort(key=lambda x: x[0])
    
    if not all_pages:
        print("❌ Không tìm thấy trang nào từ 51 trở đi")
        sys.exit(1)
    
    print(f"📄 Tổng: {len(all_pages)} trang (từ trang {all_pages[0][0]} đến {all_pages[-1][0]})")
    
    # Group into batches of 3 pages
    BATCH_SIZE = 3
    batch_num = 0
    
    for i in range(0, len(all_pages), BATCH_SIZE):
        batch_pages = all_pages[i:i + BATCH_SIZE]
        start_page = batch_pages[0][0]
        end_page = batch_pages[-1][0]
        
        fname = f"output-{start_page}-to-{end_page}.md"
        outpath = os.path.join(output_dir, fname)
        
        content = '\n'.join(p[1] for p in batch_pages)
        
        os.makedirs(output_dir, exist_ok=True)
        with open(outpath, 'w', encoding='utf-8') as f:
            f.write(content)
        
        batch_num += 1
        print(f"   ✅ {fname} ({len(batch_pages)} trang, {len(content):,} bytes)")
    
    print(f"\n🎉 Done! {batch_num} file 3-trang → {os.path.abspath(output_dir)}")


if __name__ == '__main__':
    main()