"""
hf_dataset_loader.py
────────────────────
Loads Romanian legal documents from the HuggingFace multi_eurlex dataset
and converts them into the same articles.jsonl format that pipeline.py produces.

This is a temporary substitute while legislatie.just.ro is rate-limiting us.

Install: pip install datasets
"""

import json
import re
from pathlib import Path
from datasets import load_dataset

DATA_DIR = Path("data")
DATA_DIR.mkdir(exist_ok=True)
ARTICLES_FILE = DATA_DIR / "articles.jsonl"


def split_into_chunks(text: str, law_id: int, title: str) -> list[dict]:
    """
    MultiEURLEX documents don't have article boundaries marked
    the same way as legislatie.just.ro, so we do a best-effort split.

    Strategy:
      1. Try to split on "Article X" / "Articolul X" patterns first.
      2. If none found, fall back to splitting on paragraph boundaries
         (every ~500 characters), so the chunks aren't too big for the embedder.
    """
    # Try article-style splitting first
    article_pattern = re.compile(
        r"(Articolul\s+\d+|Article\s+\d+)",
        re.IGNORECASE
    )
    matches = list(article_pattern.finditer(text))

    if matches:
        chunks = []
        for i, match in enumerate(matches):
            start = match.start()
            end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
            chunk_text = text[start:end].strip()
            if len(chunk_text) > 80:  # skip tiny fragments
                chunks.append({
                    "law_id":         law_id,
                    "law_title":      title,
                    "article_number": match.group(0).strip(),
                    "text":           chunk_text,
                    "chunk":          f"{title}\n{match.group(0).strip()}\n\n{chunk_text}",
                })
        return chunks

    # Fallback: paragraph-based chunking
    paragraphs = [p.strip() for p in text.split("\n\n") if len(p.strip()) > 80]
    chunks = []
    for i, para in enumerate(paragraphs):
        article_num = f"Paragraf {i + 1}"
        chunks.append({
            "law_id":         law_id,
            "law_title":      title,
            "article_number": article_num,
            "text":           para,
            "chunk":          f"{title}\n{article_num}\n\n{para}",
        })
    return chunks


def load_and_convert(max_docs: int = 500):
    """
    Load Romanian documents from multi_eurlex and write them to articles.jsonl.

    max_docs: how many documents to load (500 is a good starting point).
              The full dataset is ~30k documents — start small.
    """
    print("Downloading multi_eurlex Romanian split from HuggingFace...")
    print("(This downloads ~200MB on first run, then it's cached locally.)\n")

    # 'all_languages' config contains Romanian under the 'ro' key
    dataset = load_dataset(
        "multi_eurlex",
        "ro",            # Romanian subset
        split="train",
        trust_remote_code=True,
    )

    print(f"Dataset loaded. Total documents available: {len(dataset)}")
    print(f"Processing first {max_docs} documents...\n")

    total_articles = 0

    with open(ARTICLES_FILE, "w", encoding="utf-8") as f:
        for i, doc in enumerate(dataset):
            if i >= max_docs:
                break

            law_id = i + 1  # synthetic ID since HF doesn't have real law IDs
            title  = doc.get("text", "")[:80].split("\n")[0].strip() or f"Document {law_id}"
            text   = doc.get("text", "")

            if not text or len(text) < 100:
                continue

            chunks = split_into_chunks(text, law_id, title)
            for chunk in chunks:
                f.write(json.dumps(chunk, ensure_ascii=False) + "\n")
                total_articles += 1

            if (i + 1) % 50 == 0:
                print(f"  Processed {i + 1}/{max_docs} documents "
                      f"({total_articles} articles so far)...")

    print(f"\n✓ Done! Wrote {total_articles} articles to {ARTICLES_FILE}")
    print(f"  File size: {ARTICLES_FILE.stat().st_size // 1024} KB")
    print(f"\nNext step: run  python indexer.py  to build the FAISS index.")


if __name__ == "__main__":
    load_and_convert(max_docs=500)