src/senate_parser.py

"""
Senate Data Parser - Extract and parse Akoma Ntoso documents from Senate

Handles:
- Extraction from senate.tar.zst (compressed archive)
- Parsing Akoma Ntoso XML format
- Metadata extraction (URN, title, dates, etc.)
- Vectorization and FAISS integration

Data source: https://github.com/SenatoDellaRepubblica/AkomaNtosoBulkData
"""

import json
import logging
import tarfile
import tempfile
from pathlib import Path
from typing import List, Dict, Any, Optional
from xml.etree import ElementTree as ET
from datetime import datetime

logger = logging.getLogger(__name__)


class AkomaNtosoBulkParser:
    """Parse Akoma Ntoso XML documents from Senate bulk data."""
    
    NAMESPACES = {
        'akn': 'http://docs.oasis-open.org/legaldocml/ns/akn/3.0',
        'html': 'http://www.w3.org/1999/xhtml',
    }
    
    def __init__(self, work_dir: Optional[Path] = None):
        """
        Initialize parser.
        
        Args:
            work_dir: Working directory for extraction (default: temp)
        """
        self.work_dir = work_dir or Path(tempfile.gettempdir()) / "senate_parser"
        self.work_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"Senate parser initialized: {self.work_dir}")
    
    def extract_archive(self, archive_path: Path) -> Path:
        """
        Extract senate.tar.zst archive.
        
        Args:
            archive_path: Path to senate.tar.zst
        
        Returns:
            Path to extraction directory
        """
        extract_dir = self.work_dir / "extracted"
        extract_dir.mkdir(exist_ok=True)
        
        try:
            logger.info(f"Extracting {archive_path.name}...")
            with tarfile.open(archive_path, "r:zst") as tar:
                tar.extractall(path=extract_dir)
            
            xml_count = len(list(extract_dir.rglob("*.xml")))
            logger.info(f"✓ Extracted {xml_count:,} XML files")
            return extract_dir
        except Exception as e:
            logger.error(f"Extraction failed: {e}")
            raise
    
    def parse_xml(self, xml_file: Path) -> Optional[Dict[str, Any]]:
        """
        Parse single Akoma Ntoso XML document.
        
        Args:
            xml_file: Path to .xml file
        
        Returns:
            Metadata dict or None if parse fails
        """
        try:
            tree = ET.parse(xml_file)
            root = tree.getroot()
            
            # Extract metadata from preface/preamble
            meta = {
                "source": "Senate",
                "source_file": str(xml_file.name),
                "type": "Senate Document",
                "file_path": str(xml_file),
            }
            
            # Try common Akoma Ntoso paths
            for tag in ['meta', 'identification', 'FRBRIdentification']:
                elem = root.find(f".//akn:{tag}", self.NAMESPACES)
                if elem is not None:
                    meta.update(self._extract_from_element(elem))
            
            # Extract title from preface
            preface = root.find(".//akn:preface", self.NAMESPACES)
            if preface is not None:
                p = preface.find(".//akn:p", self.NAMESPACES)
                if p is not None and p.text:
                    meta["source_title"] = p.text.strip()
            
            # Extract main text
            body = root.find(".//akn:body", self.NAMESPACES)
            if body is not None:
                text_parts = []
                for p in body.findall(".//akn:p", self.NAMESPACES):
                    if p.text:
                        text_parts.append(p.text.strip())
                if text_parts:
                    meta["text_content"] = "\n".join(text_parts)
            
            # Generate URN if not present
            if "urn" not in meta or not meta["urn"]:
                meta["urn"] = self._generate_urn_from_filename(xml_file.name)
            
            logger.debug(f"Parsed: {meta.get('source_title', 'Unknown')[:50]}")
            return meta
        except Exception as e:
            logger.warning(f"Failed to parse {xml_file.name}: {e}")
            return None
    
    def _extract_from_element(self, elem: ET.Element) -> Dict[str, str]:
        """Extract metadata from XML element."""
        meta = {}
        
        for key in ['uri', 'urn', 'name', 'date', 'author', 'publisher']:
            child = elem.find(f".//akn:{key}", self.NAMESPACES)
            if child is not None and child.get("value"):
                meta[key] = child.get("value")
        
        return meta
    
    def _generate_urn_from_filename(self, filename: str) -> str:
        """Generate URN from filename if not in XML."""
        # Akoma Ntoso files typically named like: atto_dd_YYYYMMDD_XX_001.xml
        # Generate: urn:nir:senato:documento:YYYY:XX
        parts = filename.replace(".xml", "").split("_")
        if len(parts) >= 3:
            date_part = parts[2]
            year = date_part[:4] if len(date_part) >= 4 else "2000"
            doc_id = parts[3] if len(parts) > 3 else "001"
            return f"urn:nir:senato:documento:{year}:{doc_id}"
        return f"urn:nir:senato:documento:unknown"
    
    def parse_directory(self, dir_path: Path, limit: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        Parse all XML files in directory.
        
        Args:
            dir_path: Directory containing XML files
            limit: Max files to parse (optional)
        
        Returns:
            List of parsed metadata dicts
        """
        documents = []
        xml_files = sorted(dir_path.rglob("*.xml"))
        
        logger.info(f"Parsing {len(xml_files):,} XML files...")
        for i, xml_file in enumerate(xml_files):
            if limit and i >= limit:
                logger.info(f"Stopped at limit: {limit} files")
                break
            
            meta = self.parse_xml(xml_file)
            if meta:
                documents.append(meta)
            
            if (i + 1) % 100 == 0:
                logger.info(f"  Processed {i+1:,} files...")
        
        logger.info(f"✓ Parsed {len(documents):,} documents")
        return documents
    
    def export_to_jsonl(self, documents: List[Dict], output_path: Path) -> int:
        """
        Export parsed documents to JSONL for FAISS indexing.
        
        Args:
            documents: List of metadata dicts
            output_path: Output JSONL file
        
        Returns:
            Number of records exported
        """
        try:
            with open(output_path, "w", encoding="utf-8") as f:
                for doc in documents:
                    # Ensure required fields for FAISS integration
                    doc.setdefault("id", doc.get("urn", "unknown"))
                    doc.setdefault("validity_status", "in_corso")  # Senate docs are current
                    doc.setdefault("domain", "Legislativo")
                    doc.setdefault("jurisdiction", "Italiana")
                    doc.setdefault("is_vigente", True)
                    doc.setdefault("is_zombie", False)
                    doc.setdefault("is_sunset", False)
                    
                    f.write(json.dumps(doc, ensure_ascii=False) + "\n")
            
            logger.info(f"✓ Exported {len(documents)} documents to {output_path.name}")
            return len(documents)
        except Exception as e:
            logger.error(f"Export failed: {e}")
            return 0
    
    def prepare_for_faiss(
        self,
        archive_path: Path,
        output_jsonl: Path,
        limit: Optional[int] = None
    ) -> Dict[str, Any]:
        """
        End-to-end pipeline: extract → parse → export for FAISS integration.
        
        Args:
            archive_path: Path to senate.tar.zst
            output_jsonl: Output JSONL path for FAISS
            limit: Max documents to parse
        
        Returns:
            Summary dict with counts and paths
        """
        try:
            # Extract
            extract_dir = self.extract_archive(archive_path)
            
            # Parse
            documents = self.parse_directory(extract_dir, limit=limit)
            
            # Export
            exported = self.export_to_jsonl(documents, output_jsonl)
            
            return {
                "success": True,
                "total_documents": len(documents),
                "exported": exported,
                "output_file": str(output_jsonl),
                "output_size_mb": output_jsonl.stat().st_size / (1024 * 1024) if output_jsonl.exists() else 0,
            }
        except Exception as e:
            logger.error(f"Pipeline failed: {e}")
            return {
                "success": False,
                "error": str(e),
            }


def merge_senate_with_faiss(
    senate_jsonl: Path,
    faiss_metadata: Path,
    output_path: Path,
    keep_duplicates: bool = False
) -> int:
    """
    Merge Senate data with existing FAISS metadata.
    
    Args:
        senate_jsonl: Senate documents in JSONL
        faiss_metadata: Existing FAISS metadata JSONL
        output_path: Output merged JSONL
        keep_duplicates: If False, skip Senate docs already in FAISS (by URN)
    
    Returns:
        Total records in merged file
    """
    try:
        # Load existing URNs
        existing_urns = set()
        with open(faiss_metadata, "r", encoding="utf-8") as f:
            for line in f:
                try:
                    doc = json.loads(line.strip())
                    if "urn" in doc:
                        existing_urns.add(doc["urn"])
                except json.JSONDecodeError:
                    pass
        
        logger.info(f"Existing FAISS has {len(existing_urns):,} unique URNs")
        
        # Merge
        total = 0
        skipped = 0
        with open(output_path, "w", encoding="utf-8") as out:
            # Write existing
            with open(faiss_metadata, "r", encoding="utf-8") as f:
                for line in f:
                    out.write(line)
                    total += 1
            
            # Write Senate (if not duplicate)
            with open(senate_jsonl, "r", encoding="utf-8") as f:
                for line in f:
                    try:
                        doc = json.loads(line.strip())
                        urn = doc.get("urn")
                        
                        if not keep_duplicates and urn in existing_urns:
                            skipped += 1
                        else:
                            out.write(line)
                            total += 1
                    except json.JSONDecodeError:
                        pass
        
        logger.info(f"✓ Merged: {total:,} total, {skipped:,} duplicates skipped")
        return total
    except Exception as e:
        logger.error(f"Merge failed: {e}")
        return 0


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # Example usage
    parser = AkomaNtosoBulkParser()
    
    # Parse senate.tar.zst from HF dataset
    archive = Path("/tmp/senate.tar.zst")
    output = Path("/tmp/senate_metadata.jsonl")
    
    if archive.exists():
        result = parser.prepare_for_faiss(archive, output, limit=1000)
        print(f"Result: {result}")