import os
import json
import pymupdf
from pathlib import Path
from typing import Dict, Any, List

def extract_text_from_pdf(pdf_path: str) -> Dict[str, Any]:
    """Extract text and metadata from PDF"""
    try:
        doc = pymupdf.open(pdf_path)
        
        text = ""
        for page_num in range(len(doc)):
            page = doc[page_num]
            text += page.get_text()
        
        metadata = doc.metadata
        doc.close()
        
        # Basic cleaning
        text = clean_text(text)
        
        return {
            "title": metadata.get('title', Path(pdf_path).stem),
            "text": text,
            "author": metadata.get('author', ''),
            "subject": metadata.get('subject', ''),
            "num_pages": len(doc)
        }
    except Exception as e:
        print(f"Error processing {pdf_path}: {e}")
        return None

def process_pdf_directory(input_dir: str, output_dir: str) -> int:
    """Process all PDFs in a directory"""
    os.makedirs(output_dir, exist_ok=True)
    
    count = 0
    for pdf_file in Path(input_dir).glob("*.pdf"):
        result = extract_text_from_pdf(str(pdf_file))
        
        if result and result.get('text'):
            output_file = os.path.join(output_dir, f"{pdf_file.stem}.json")
            
            doc = {
                "id": f"pdf_{pdf_file.stem}",
                "title": result['title'],
                "text": result['text'],
                "source": "pdf",
                "metadata": {
                    "author": result.get('author'),
                    "subject": result.get('subject'),
                    "pages": result.get('num_pages')
                }
            }
            
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(doc, f, ensure_ascii=False, indent=2)
            
            count += 1
            if count % 100 == 0:
                print(f"Processed {count} PDFs")
    
    return count

def clean_text(text: str) -> str:
    """Clean extracted PDF text"""
    import re
    
    # Remove multiple spaces
    text = re.sub(r'  +', ' ', text)
    # Remove multiple newlines
    text = re.sub(r'\n\n+', '\n', text)
    # Remove special characters
    text = re.sub(r'[\x00-\x08\x0b-\x0c\x0e-\x1f]', '', text)
    
    return text.strip()

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 3:
        print("Usage: python parse_pdf.py <input_dir> <output_dir>")
        sys.exit(1)
    
    input_dir = sys.argv[1]
    output_dir = sys.argv[2]
    
    count = process_pdf_directory(input_dir, output_dir)
    print(f"Successfully processed {count} PDFs")