#!/usr/bin/env python3
"""Парсер WI-файлов ElsaWin (OLE2 → XML → JSONL)"""
import os, re, json, sys
import olefile

def parse_wi_file(path):
    """Извлечь XML-потоки из WI (OLE2) файла."""
    docs = []
    
    with olefile.OleFileIO(path) as ole:
        for stream in ole.listdir():
            stream_name = "/".join(stream)
            if not stream_name.endswith(".xml"):
                continue
            
            data = ole.openstream(stream).read()
            text = data.decode("utf-8", errors="replace")
            
            # Extract key elements
            title_parts = []
            
            m = re.search(r"<marke>(.*?)</marke>", text, re.IGNORECASE)
            marke = m.group(1) if m else ""
            
            m = re.search(r"<modell>(.*?)</modell>", text, re.IGNORECASE)
            modell = m.group(1) if m else ""
            
            m = re.search(r"<obergrup>(.*?)</obergrup>", text, re.IGNORECASE)
            obergrup = m.group(1) if m else ""
            
            m = re.search(r"<baugrup>(.*?)</baugrup>", text, re.IGNORECASE)
            baugrup = m.group(1) if m else ""
            
            # Название документа
            m = re.search(r"<h-kap[^>]*?typen-bez=\"([^\"]+)\"", text)
            typen = m.group(1) if m else ""
            
            # Title from vorspann
            title = " ".join(p for p in [obergrup, baugrup] if p)
            if marke:
                title = f"{marke} {modell}: {title}" if modell else f"{marke}: {title}"
            
            # Clean body
            body = re.sub(r"<[^>]+>", " ", text)
            body = re.sub(r"\s+", " ", body).strip()
            
            docs.append({
                "title": title,
                "marke": marke,
                "modell": modell,
                "obergrup": obergrup,
                "baugrup": baugrup,
                "text": body,
            })
    
    return docs


def main():
    if len(sys.argv) < 3:
        print(f"Usage: {sys.argv[0]} <elsa_docs_dir> <output_dir>")
        sys.exit(1)
    
    elsa_docs = os.path.abspath(sys.argv[1])
    out_dir = os.path.abspath(sys.argv[2])
    os.makedirs(out_dir, exist_ok=True)
    
    # Collect en-GB WI files
    wi_files = []
    for root, dirs, files in os.walk(elsa_docs):
        for f in files:
            if f.endswith(".wi") and "en-GB" in f:
                wi_files.append(os.path.join(root, f))
    
    print(f"Found {len(wi_files)} en-GB WI files")
    
    # By category
    by_cat = {}
    for f in wi_files:
        if "/rl/" in f:
            cat = "rl"
        elif "/igg/" in f:
            cat = "igg"
        elif "/ki/" in f:
            cat = "ki"
        else:
            cat = "other"
        by_cat.setdefault(cat, []).append(f)
    
    for cat, files in by_cat.items():
        print(f"  {cat}: {len(files)} files")
    
    # Parse
    total_docs = 0
    part_size = 5000
    part_idx = 0
    buf = []
    
    for cat, files in by_cat.items():
        for fpath in sorted(files):
            try:
                docs = parse_wi_file(fpath)
                for doc in docs:
                    doc["source"] = f"elsa_wi_{cat}"
                    doc["file"] = os.path.relpath(fpath, elsa_docs)
                    buf.append(doc)
                    total_docs += 1
            except Exception as e:
                print(f"  ERROR {fpath}: {e}", file=sys.stderr)
            
            if len(buf) >= part_size:
                out_file = os.path.join(out_dir, f"part_{part_idx:04d}.jsonl")
                with open(out_file, "w", encoding="utf-8") as f:
                    for r in buf:
                        f.write(json.dumps(r, ensure_ascii=False) + "\n")
                print(f"  Part {part_idx:04d}: {len(buf)} → {out_file}")
                buf = []
                part_idx += 1
    
    if buf:
        out_file = os.path.join(out_dir, f"part_{part_idx:04d}.jsonl")
        with open(out_file, "w", encoding="utf-8") as f:
            for r in buf:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
        print(f"  Part {part_idx:04d}: {len(buf)} → {out_file}")
    
    print(f"\nTotal: {total_docs} XML documents from {len(wi_files)} WI files → {out_dir}")


if __name__ == "__main__":
    main()