import bz2
import xml.etree.ElementTree as ET
import json
import os
from pathlib import Path
from typing import Generator, Dict, Any

def extract_text_from_wikipedia_dump(bz2_path: str, output_dir: str, limit: int = None) -> int:
    """Extract text from Wikipedia XML dump"""
    os.makedirs(output_dir, exist_ok=True)
    
    count = 0
    with bz2.BZ2File(bz2_path, 'r') as f:
        context = ET.iterparse(f, events=('start', 'end'))
        context = iter(context)
        
        event, root = next(context)
        
        for event, elem in context:
            if event == 'end' and elem.tag.endswith('page'):
                doc = parse_page(elem)
                if doc and doc.get('text'):
                    filename = f"{output_dir}/wiki_{doc['id']}.json"
                    with open(filename, 'w', encoding='utf-8') as out:
                        json.dump(doc, out, ensure_ascii=False)
                    
                    count += 1
                    if count % 1000 == 0:
                        print(f"Processed {count} documents")
                    
                    if limit and count >= limit:
                        break
                
                root.clear()
    
    return count

def parse_page(page_elem) -> Dict[str, Any]:
    """Parse a single Wikipedia page"""
    ns = {'wiki': 'http://www.mediawiki.org/xml/export-0.10/'}
    
    title_elem = page_elem.find('wiki:title', ns) or page_elem.find('.//title')
    page_id_elem = page_elem.find('wiki:id', ns) or page_elem.find('.//id')
    revision_elem = page_elem.find('.//revision')
    
    if revision_elem is None:
        return None
    
    text_elem = revision_elem.find('wiki:text', ns) or revision_elem.find('.//text')
    
    if title_elem is None or text_elem is None:
        return None
    
    title = title_elem.text or ""
    page_id = page_id_elem.text if page_id_elem is not None else ""
    text = text_elem.text or ""
    
    # Skip redirects and special pages
    if text.startswith('#REDIRECT') or title.startswith('Wikipedia:'):
        return None
    
    # Basic text cleaning
    text = clean_wikipedia_markup(text)
    
    if len(text.split()) < 50:
        return None
    
    return {
        "id": f"wiki_{page_id}",
        "title": title,
        "text": text,
        "source": "wikipedia",
        "lang": "ru"  # Change for English
    }

def clean_wikipedia_markup(text: str) -> str:
    """Remove Wikipedia markup"""
    import re
    
    # Remove wiki links
    text = re.sub(r'\[\[([^\|]*\|)?([^\]]*)\]\]', r'\2', text)
    # Remove external links
    text = re.sub(r'\[http[^\s]+ ([^\]]*)\]', r'\1', text)
    # Remove templates
    text = re.sub(r'\{\{[^}]*\}\}', '', text)
    # Remove categories
    text = re.sub(r'\[\[Category:[^\]]*\]\]', '', text)
    # Remove refs
    text = re.sub(r'<ref[^>]*>[^<]*</ref>', '', text)
    # Remove HTML tags
    text = re.sub(r'<[^>]*>', '', text)
    # Clean whitespace
    text = re.sub(r'\n\n+', '\n', text)
    text = re.sub(r' +', ' ', text)
    
    return text.strip()

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 3:
        print("Usage: python parse_wikipedia.py <input_bz2> <output_dir> [limit]")
        sys.exit(1)
    
    input_file = sys.argv[1]
    output_dir = sys.argv[2]
    limit = int(sys.argv[3]) if len(sys.argv) > 3 else None
    
    count = extract_text_from_wikipedia_dump(input_file, output_dir, limit)
    print(f"Successfully processed {count} documents")