#!/usr/bin/env python3
"""
Data Streaming and Parsing Utility for the HackIndia Challenge.
Author: Team Ascended
Description: Memory-efficient streaming of financial QA datasets (specifically
             mishface123/adaption-econ-finance-qa-pairs) with batch generation,
             parsing, and formatting capabilities.
"""

import argparse
import sys
from typing import Dict, Generator, List, Optional
try:
    from datasets import load_dataset
except ImportError:
    print("Warning: 'datasets' library not found. Please install it via 'pip install datasets'.")


def stream_dataset(
    dataset_name: str,
    split: str = "train",
    batch_size: int = 1,
    limit: Optional[int] = None
) -> Generator[List[Dict], None, None]:
    """
    Streams a dataset from Hugging Face and yields parsed data in batches.

    Args:
        dataset_name: Hugging Face dataset identifier.
        split: The dataset split to stream (e.g., 'train', 'test').
        batch_size: Number of records to yield at a time.
        limit: Max number of records to process before stopping.

    Yields:
        A list of parsed dictionaries containing formatted instruction-response pairs.
    """
    try:
        # Load the dataset in streaming mode to minimize RAM overhead
        dataset = load_dataset(dataset_name, split=split, streaming=True)
    except Exception as e:
        print(f"Error loading dataset {dataset_name}: {e}", file=sys.stderr)
        return

    batch = []
    count = 0

    for record in dataset:
        parsed_record = parse_record(record)
        batch.append(parsed_record)
        count += 1

        if len(batch) == batch_size:
            yield batch
            batch = []

        if limit is not None and count >= limit:
            break

    if batch:
        yield batch


def parse_record(record: Dict) -> Dict:
    """
    Parses a single record from the dataset and standardizes its keys.
    Designed specifically for 'mishface123/adaption-econ-finance-qa-pairs'.

    Args:
        record: The raw record dict from the dataset.

    Returns:
        A standardized dictionary containing instruction, context, and response.
    """
    # Standardize keys depending on schema
    instruction = record.get("instruction") or record.get("question") or ""
    context = record.get("context") or ""
    response = record.get("output") or record.get("response") or record.get("answer") or ""
    metadata = record.get("metadata") or {}

    # Extract target market metadata (e.g., Indian localization indicator)
    localized = metadata.get("localized", False) or "india" in str(metadata).lower()

    return {
        "instruction": instruction.strip(),
        "context": context.strip(),
        "response": response.strip(),
        "localized": localized,
        "raw_keys": list(record.keys())
    }


def main():
    parser = argparse.ArgumentParser(
        description="Stream and parse datasets for the AutoScientist pipeline."
    )
    parser.add_argument(
        "--dataset",
        type=str,
        default="mishface123/adaption-econ-finance-qa-pairs",
        help="Hugging Face dataset name to stream."
    )
    parser.add_argument(
        "--split",
        type=str,
        default="train",
        help="Dataset split (train, validation, test)."
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=4,
        help="Batch size for yield/printing."
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=10,
        help="Maximum records to stream/parse."
    )
    args = parser.parse_args()

    print(f"Streaming dataset: {args.dataset} (split: {args.split})...")
    batch_generator = stream_dataset(
        dataset_name=args.dataset,
        split=args.split,
        batch_size=args.batch_size,
        limit=args.limit
    )

    for i, batch in enumerate(batch_generator):
        print(f"\n--- Batch {i+1} (Size: {len(batch)}) ---")
        for record_idx, record in enumerate(batch):
            print(f"\nRecord {record_idx+1}:")
            print(f"  Instruction : {record['instruction'][:120]}...")
            print(f"  Context     : {record['context'][:120]}...")
            print(f"  Response    : {record['response'][:120]}...")
            print(f"  Localized   : {record['localized']}")


if __name__ == "__main__":
    main()