#!/usr/bin/env python3
"""Extract images from docx with context (surrounding paragraphs).

Parses document.xml to find <w:drawing> elements and extracts:
- Image file path (from word/media/)
- Relationship ID
- Context: text from paragraphs before and after the image

Output: JSON with image list and context info.
"""

import json
import sys
import zipfile
import xml.etree.ElementTree as ET
from pathlib import Path


# Namespaces
W_NS = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
R_NS = "http://schemas.openxmlformats.org/officeDocument/2006/relationships"
A_NS = "http://schemas.openxmlformats.org/drawingml/2006/main"

NS = {
    "w": W_NS,
    "wp": "http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing",
    "a": A_NS,
    "pic": "http://schemas.openxmlformats.org/drawingml/2006/picture",
    "r": R_NS,
}


def extract_docx_images(docx_path: str, output_dir: str) -> dict:
    """Extract images from docx with context information.
    
    Args:
        docx_path: Path to the .docx file
        output_dir: Directory to extract images to
    
    Returns:
        Dict with 'images' list containing path, relId, contextBefore, contextAfter
    """
    docx = Path(docx_path)
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    
    if not docx.exists():
        return {"error": f"File not found: {docx_path}", "images": []}
    
    try:
        with zipfile.ZipFile(docx, "r") as zf:
            # Extract all media files
            media_files = [f for f in zf.namelist() if f.startswith("word/media/")]
            for mf in media_files:
                zf.extract(mf, out)
            
            # Parse document.xml for image positions
            with zf.open("word/document.xml") as f:
                doc_xml = f.read()
            
            # Parse relationships
            with zf.open("word/_rels/document.xml.rels") as f:
                rels_xml = f.read()
    except Exception as e:
        return {"error": f"Failed to read docx: {e}", "images": []}
    
    # Build relId -> media file mapping
    rels_root = ET.fromstring(rels_xml)
    rel_map = {}
    for rel in rels_root:
        rel_id = rel.get("Id")
        target = rel.get("Target")  # e.g., "media/image1.png"
        if rel_id and target and target.startswith("media/"):
            rel_map[rel_id] = target
    
    # Parse document to find images with context
    doc_root = ET.fromstring(doc_xml)
    body = doc_root.find(".//w:body", NS)
    if body is None:
        return {"error": "No body found in document", "images": []}
    
    paragraphs = list(body.findall("w:p", NS))
    
    images = []
    for idx, para in enumerate(paragraphs):
        # Find all drawings in this paragraph
        drawings = para.findall(".//w:drawing", NS)
        if not drawings:
            continue
        
        # Get context: text from previous and next non-empty paragraphs
        context_before = ""
        for i in range(idx - 1, -1, -1):
            text = "".join(t.text or "" for t in paragraphs[i].findall(".//w:t", NS))
            if text.strip():
                context_before = text.strip()
                break
        
        context_after = ""
        for i in range(idx + 1, len(paragraphs)):
            text = "".join(t.text or "" for t in paragraphs[i].findall(".//w:t", NS))
            if text.strip():
                context_after = text.strip()
                break
        
        # Extract image info from each drawing
        for drawing in drawings:
            # Find the blip reference
            blip = drawing.find(".//a:blip", NS)
            if blip is None:
                continue
            
            rel_id = blip.get(f"{{{R_NS}}}embed")
            if not rel_id:
                continue
            
            media_target = rel_map.get(rel_id)
            if not media_target:
                continue
            
            # Copy image to output directory
            # zip extracts to word/media/, rels target is media/
            src_path = out / "word" / media_target
            if src_path.exists():
                filename = Path(media_target).name
                dest_path = out / filename
                if src_path != dest_path:
                    import shutil
                    shutil.copy2(str(src_path), str(dest_path))
                
                images.append({
                    "filename": filename,
                    "path": str(dest_path),
                    "relId": rel_id,
                    "contextBefore": context_before,
                    "contextAfter": context_after,
                })
    
    return {"images": images}


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: extract-docx-images.py <docx_path> <output_dir>")
        sys.exit(1)
    
    result = extract_docx_images(sys.argv[1], sys.argv[2])
    print(json.dumps(result, ensure_ascii=False, indent=2))