#!/usr/bin/env python3
"""
Post-release dataset processing script for CaDaR

This script:
1. Loads the arbml/darija dataset from HuggingFace
2. Applies bizi2ara transformation on the 'darija' column
3. Creates a new 'darija_ar' column with Arabic script results
4. Pushes the transformed dataset to OiTe/CaDaR with version tagging

Usage:
    python scripts/process_darija_dataset.py --version <version>

Environment variables:
    HF_TOKEN: HuggingFace API token (required)
    GH_TOKEN: GitHub token (optional, for repo info)
"""

import argparse
import logging
import os
import sys
from typing import Optional

import pandas as pd
from datasets import Dataset, DatasetDict, load_dataset
from huggingface_hub import HfApi, create_repo, whoami
from huggingface_hub.errors import HfHubHTTPError

import cadar

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler("dataset_processing.log"),
    ],
)
logger = logging.getLogger(__name__)


class DatasetProcessor:
    """Process Darija dataset with CaDaR transformations"""

    def __init__(self, version: str):
        self.version = version.strip("v")  # Remove 'v' prefix if present
        self.source_dataset = "arbml/darija"
        self.target_dataset = "OiTe/CaDaR"
        self.hf_token = os.getenv("HF_TOKEN")

        if not self.hf_token:
            raise ValueError("HF_TOKEN environment variable is required")

        # Initialize CaDaR processor
        self.cadar_processor = cadar.CaDaR(darija="Ma")
        logger.info(f"Initialized CaDaR processor for Moroccan Darija")
        logger.info(f"Processing version: {self.version}")

    def load_source_dataset(self) -> DatasetDict:
        """Load the source dataset from HuggingFace"""
        try:
            logger.info(f"Loading dataset: {self.source_dataset}")
            dataset_dict = load_dataset(self.source_dataset, token=self.hf_token)
            logger.info(f"Dataset loaded successfully")

            # Log dataset info
            for split_name, dataset in dataset_dict.items():
                logger.info(f"Split '{split_name}': {len(dataset)} rows")

            return dataset_dict

        except Exception as e:
            logger.error(f"Failed to load source dataset: {e}")
            raise

    def process_text_batch(self, texts: list[str]) -> list[str]:
        """Apply bizi2ara transformation to a batch of texts"""
        results = []
        total = len(texts)

        for i, text in enumerate(texts, 1):
            try:
                if pd.isna(text) or text is None:
                    results.append("")
                    continue

                # Apply bizi2ara transformation
                arabic_text = self.cadar_processor.bizi2ara(str(text))
                results.append(arabic_text)

                if i % 100 == 0 or i == total:
                    logger.info(f"Processed {i}/{total} texts in batch")

            except Exception as e:
                logger.warning(f"Failed to process text '{text[:50]}...': {e}")
                results.append("")  # Empty string for failed conversions

        return results

    def process_dataset(self, dataset_dict: DatasetDict) -> DatasetDict:
        """Process the entire dataset"""
        processed_dict = DatasetDict()

        for split_name, dataset in dataset_dict.items():
            logger.info(f"Processing split: {split_name}")

            # Convert to pandas for easier manipulation
            df = dataset.to_pandas()
            logger.info(f"Split '{split_name}' has {len(df)} rows")

            # Check if 'darija' column exists
            if "darija" not in df.columns:
                raise ValueError(f"'darija' column not found in split '{split_name}'")

            # Log some examples
            sample_texts = df["darija"].dropna().head(3).tolist()
            logger.info(f"Sample texts from '{split_name}':")
            for i, text in enumerate(sample_texts, 1):
                logger.info(f"  {i}. {text}")

            # Process the darija column in batches
            batch_size = 1000
            total_rows = len(df)
            arabic_texts = []

            for start_idx in range(0, total_rows, batch_size):
                end_idx = min(start_idx + batch_size, total_rows)
                batch_texts = df["darija"].iloc[start_idx:end_idx].tolist()

                logger.info(f"Processing batch {start_idx}-{end_idx}/{total_rows}")
                batch_arabic = self.process_text_batch(batch_texts)
                arabic_texts.extend(batch_arabic)

            # Add new column
            df["darija_ar"] = arabic_texts

            # Log some results
            sample_results = df[["darija", "darija_ar"]].dropna().head(3)
            logger.info(f"Sample results from '{split_name}':")
            for _, row in sample_results.iterrows():
                logger.info(f"  {row['darija']} → {row['darija_ar']}")

            # Convert back to Dataset
            processed_dataset = Dataset.from_pandas(df)
            processed_dict[split_name] = processed_dataset

        return processed_dict

    def push_to_huggingface(self, dataset_dict: DatasetDict):
        """Push processed dataset to HuggingFace"""
        try:
            # Create repository if it doesn't exist
            try:
                create_repo(
                    repo_id=self.target_dataset,
                    token=self.hf_token,
                    repo_type="dataset",
                    private=False,
                    exist_ok=True,
                )
                logger.info(f"Repository {self.target_dataset} ready")
            except HfHubHTTPError as e:
                if "already exists" in str(e):
                    logger.info(f"Repository {self.target_dataset} already exists")
                else:
                    raise

            # Get user info
            user_info = whoami(token=self.hf_token)
            logger.info(
                f"Pushing as user: {user_info['name']} ({user_info['username']})"
            )

            # Create dataset card
            dataset_card = self._create_dataset_card()

            # Push dataset
            logger.info(f"Pushing dataset to: {self.target_dataset}")

            dataset_dict.push_to_hub(
                repo_id=self.target_dataset,
                token=self.hf_token,
                commit_message=f"Update dataset with CaDaR v{self.version} transformations",
            )

            # Create README.md with dataset card content
            api = HfApi(token=self.hf_token)
            api.upload_file(
                path_or_fileobj=dataset_card.encode("utf-8"),
                path_in_repo="README.md",
                repo_id=self.target_dataset,
                repo_type="dataset",
                commit_message=f"Add dataset card for CaDaR v{self.version}",
            )

            logger.info(f"Dataset successfully pushed to: {self.target_dataset}")

        except Exception as e:
            logger.error(f"Failed to push dataset: {e}")
            raise

    def _create_dataset_card(self) -> str:
        """Create dataset card content"""
        card_content = f"""---
language: 
- ar
- pcm
license: apache-2.0
tags:
- darija
- moroccan-arabic
- arabic
- transliteration
- cadar
size_categories:
- 10K<n<100K
---

# CaDaR Transformed Darija Dataset

This dataset contains Darija (Moroccan Arabic) text transformed using **CaDaR v{self.version}**.

## Dataset Description

The dataset contains parallel text in:
- `darija`: Original Latin script (Bizi) text from the arbml/darija dataset
- `darija_ar`: Arabic script text generated using CaDaR's bizi2ara transformation

## Transformation Process

Each entry in the dataset was processed using:

```python
import cadar
processor = cadar.CaDaR(darija="Ma")
arabic_text = processor.bizi2ara(latin_text)
```

## Version Information

- **CaDaR Version**: v{self.version}
- **Transformation**: bizi2ara (Latin to Arabic)
- **Dialect**: Moroccan Darija (Ma)
- **Processing Date**: {pd.Timestamp.now().strftime("%Y-%m-%d %H:%M:%S")}

## Dataset Structure

The dataset contains the following splits:

{{ dataset_splits }}

## Usage Example

```python
from datasets import load_dataset

# Load the dataset
dataset = load_dataset("OiTe/CaDaR")

# Access the data
train_data = dataset['train']
example = train_data[0]

print(f"Latin: {{example['darija']}}")
print(f"Arabic: {{example['darija_ar']}}")
```

## Citation

If you use this dataset, please cite:

```bibtex
@dataset{{cadar_darija_{self.version.replace(".", "_")},
  title = {{CaDaR Transformed Darija Dataset}},
  author = {{CaDaR Team}},
  year = {pd.Timestamp.now().year},
  version = {{v{self.version}}},
  publisher = {{Hugging Face}},
  howpublished = {{\\url{{https://huggingface.co/datasets/OiTe/CaDaR}}}}
}}
```
"""
        return card_content

    def run(self):
        """Main processing pipeline"""
        try:
            logger.info("Starting dataset processing pipeline")

            # Load source dataset
            dataset_dict = self.load_source_dataset()

            # Process dataset
            processed_dict = self.process_dataset(dataset_dict)

            # Push to HuggingFace
            self.push_to_huggingface(processed_dict)

            logger.info("Dataset processing completed successfully")

        except Exception as e:
            logger.error(f"Dataset processing failed: {e}")
            sys.exit(1)


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(description="Process Darija dataset with CaDaR")
    parser.add_argument("--version", required=True, help="CaDaR version (e.g., 0.1.11)")

    args = parser.parse_args()

    # Check environment variables
    if not os.getenv("HF_TOKEN"):
        logger.error("HF_TOKEN environment variable is required")
        sys.exit(1)

    # Run processor
    processor = DatasetProcessor(args.version)
    processor.run()


if __name__ == "__main__":
    main()
