export_hub.py

"""Export simulation traces to local files and Hugging Face Hub."""

import json
import os
import re
import shutil
import tempfile
from datetime import datetime
from pathlib import Path

from schemas import SimulationResult, Town, BroadcastEvent


# PII patterns to sanitize from exports
PII_PATTERNS = [
    (r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', '[EMAIL_REMOVED]'),
    (r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b', '[PHONE_REMOVED]'),
    (r'\b\d{1,5}\s+[\w\s]+(?:Street|St|Avenue|Ave|Road|Rd|Boulevard|Blvd|Drive|Dr|Lane|Ln|Way|Court|Ct)\b', '[ADDRESS_REMOVED]'),
]

DATASET_CARD_TEMPLATE = """# Analog Town Trace Dataset

## Description

This dataset contains fictional agent state-transition traces generated by the
Analog Town simulator.

**Town:** {town_name}
**Event:** {event_title}
**Agents:** {agent_count}
**Generated:** {created_at}

## Intended Use

This dataset is intended for demonstrating small-model perspective rehearsal,
fictional character simulation, and inspectable state-evolution prompting.

## Limitations

The traces are fictional. They should not be interpreted as predictions of real
people.

## Ethical Considerations

Users should avoid uploading real private personal information. Real individuals
should be fictionalized or represented only with consent.

## Data Fields

- `town_id` — identifier for the simulated town
- `event_id` — identifier for the broadcast event
- `agent_id` — identifier for the responding agent
- `frequency` — the agent's radio frequency
- `previous_state` — agent state before the event
- `transition` — the state transition details
- `updated_state` — agent state after the event
- `created_at` — timestamp of generation

## Source

Generated by [Analog Town](https://github.com/analog-town) — a shortwave simulator for fictional minds.
"""


class ExportManager:
    """Manages export of simulation traces to local files and HF Hub."""

    def _sanitize_text(self, text: str) -> str:
        """Remove PII patterns from text."""
        for pattern, replacement in PII_PATTERNS:
            text = re.sub(pattern, replacement, text, flags=re.IGNORECASE)
        return text

    def _sanitize_dict(self, data: dict) -> dict:
        """Recursively sanitize PII from a dictionary."""
        sanitized = {}
        for key, value in data.items():
            if isinstance(value, str):
                sanitized[key] = self._sanitize_text(value)
            elif isinstance(value, dict):
                sanitized[key] = self._sanitize_dict(value)
            elif isinstance(value, list):
                sanitized[key] = [
                    self._sanitize_dict(item) if isinstance(item, dict)
                    else self._sanitize_text(item) if isinstance(item, str)
                    else item
                    for item in value
                ]
            else:
                sanitized[key] = value
        return sanitized

    def _make_event_id(self, event: BroadcastEvent) -> str:
        """Generate a simple event ID from the title."""
        return re.sub(r'[^a-z0-9]+', '_', event.title.lower()).strip('_')[:50]

    def export_local(
        self,
        result: SimulationResult,
        town: Town,
        output_dir: str = "outputs",
    ) -> str:
        """Export simulation traces to local directory.

        Creates:
        - town.json
        - broadcast_event.json
        - agent_traces.jsonl
        - README.md

        Returns the output directory path.
        """
        # Create timestamped output directory
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        export_dir = Path(output_dir) / f"export_{result.town_id}_{timestamp}"
        export_dir.mkdir(parents=True, exist_ok=True)

        event_id = self._make_event_id(result.event)

        # 1. town.json
        town_data = self._sanitize_dict(town.model_dump())
        with open(export_dir / "town.json", "w") as f:
            json.dump(town_data, f, indent=2, ensure_ascii=False)

        # 2. broadcast_event.json
        event_data = self._sanitize_dict(result.event.model_dump())
        with open(export_dir / "broadcast_event.json", "w") as f:
            json.dump(event_data, f, indent=2, ensure_ascii=False)

        # 3. agent_traces.jsonl
        # Find agent frequencies from town
        freq_map = {agent.id: agent.frequency for agent in town.agents}

        with open(export_dir / "agent_traces.jsonl", "w") as f:
            for transition in result.transitions:
                trace_line = {
                    "town_id": result.town_id,
                    "event_id": event_id,
                    "agent_id": transition.agent_id,
                    "frequency": freq_map.get(transition.agent_id, 0.0),
                    "previous_state": {},  # We don't store previous state separately
                    "transition": self._sanitize_dict(
                        transition.model_dump(
                            exclude={"updated_state"}
                        )
                    ),
                    "updated_state": self._sanitize_dict(
                        transition.updated_state.model_dump()
                    ),
                    "created_at": result.created_at,
                }
                f.write(json.dumps(trace_line, ensure_ascii=False) + "\n")

        # 4. README.md (dataset card)
        readme_content = DATASET_CARD_TEMPLATE.format(
            town_name=town.name,
            event_title=result.event.title,
            agent_count=len(result.transitions),
            created_at=result.created_at,
        )
        with open(export_dir / "README.md", "w") as f:
            f.write(readme_content)

        return str(export_dir)

    def export_zip(
        self,
        result: SimulationResult,
        town: Town,
        output_dir: str = "outputs",
    ) -> str:
        """Export simulation traces as a ZIP file.

        Returns the path to the ZIP file.
        """
        # First export to temp directory
        with tempfile.TemporaryDirectory() as tmp_dir:
            export_dir = self.export_local(result, town, tmp_dir)

            # Create ZIP
            zip_dir = Path(output_dir)
            zip_dir.mkdir(parents=True, exist_ok=True)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            zip_name = f"analog_town_{result.town_id}_{timestamp}"
            zip_path = shutil.make_archive(
                str(zip_dir / zip_name), "zip", export_dir
            )

        return zip_path

    def export_to_hub(
        self,
        result: SimulationResult,
        town: Town,
        repo_id: str,
        token: str | None = None,
    ) -> str:
        """Upload simulation traces to Hugging Face Hub as a dataset.

        Args:
            result: The simulation result
            town: The town configuration
            repo_id: HF repo ID (e.g., "username/analog-town-graybridge")
            token: HF token (uses env HF_TOKEN if not provided)

        Returns:
            The URL of the created dataset
        """
        from huggingface_hub import HfApi

        token = token or os.getenv("HF_TOKEN")
        if not token:
            raise ValueError("HF_TOKEN required for Hub upload")

        api = HfApi(token=token)

        # Export to temp directory first
        with tempfile.TemporaryDirectory() as tmp_dir:
            export_dir = self.export_local(result, town, tmp_dir)

            # Create or get repo
            try:
                api.create_repo(
                    repo_id=repo_id,
                    repo_type="dataset",
                    exist_ok=True,
                )
            except Exception as e:
                print(f"Note: repo creation issue (may already exist): {e}")

            # Upload all files
            api.upload_folder(
                folder_path=export_dir,
                repo_id=repo_id,
                repo_type="dataset",
            )

        return f"https://huggingface.co/datasets/{repo_id}"

    def generate_dataset_card(self, town: Town, event: BroadcastEvent) -> str:
        """Generate a dataset card markdown string."""
        return DATASET_CARD_TEMPLATE.format(
            town_name=town.name,
            event_title=event.title,
            agent_count=len(town.agents),
            created_at=datetime.now().isoformat(),
        )