"""
InferRoute Streaming Benchmark Loader (Hugging Face streaming=True).

Supports zero-disk-download streaming evaluation from:
1. allenai/WildChat-4.8M (3.2M real ChatGPT conversations up to Aug 2025, ideal for production traffic simulation)
2. HuggingFaceH4/no_robots (10,000 category-labeled instructions, ideal for task-aware routing correctness)
3. OpenAssistant/oasst1 (Multi-turn conversation trees)
4. lmsys/lmsys-chat-1m (Gated multi-model arena dataset)
"""

import os
import sys
import json
import time
import asyncio
from typing import Generator, Dict, Any, Optional

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Lazy import datasets
try:
    from datasets import load_dataset
    HAS_HF_DATASETS = True
except ImportError:
    HAS_HF_DATASETS = False


def stream_wildchat_prompts(limit: int = 1000) -> Generator[Dict[str, Any], None, None]:
    """
    Streams real user prompts from allenai/WildChat-4.8M with streaming=True.
    Zero 15GB repo download required.
    """
    if not HAS_HF_DATASETS:
        print("[WARN] datasets library not installed. Cannot stream WildChat-4.8M.")
        return

    print(f"[STREAM] Connecting to allenai/WildChat-4.8M (streaming=True)...")
    try:
        ds = load_dataset("allenai/WildChat-4.8M", split="train", streaming=True)
        count = 0
        for row in ds:
            conversation = row.get("conversation", [])
            user_messages = [
                msg.get("content", "")
                for msg in conversation
                if isinstance(msg, dict) and msg.get("role") == "user" and msg.get("content")
            ]

            if user_messages:
                prompt = user_messages[0]
                count += 1
                yield {
                    "id": f"wildchat_{count:06d}",
                    "source": "allenai/WildChat-4.8M",
                    "category": "wildchat_real_user",
                    "prompt": prompt.strip(),
                    "requires_json": ("json" in prompt.lower() or "schema" in prompt.lower())
                }
                if count >= limit:
                    break
        print(f"[OK] Streamed {count:,} real prompts from allenai/WildChat-4.8M")
    except Exception as e:
        print(f"[ERR] WildChat streaming exception: {e}")


def stream_no_robots_prompts(limit: int = 1000) -> Generator[Dict[str, Any], None, None]:
    """
    Streams category-labeled instructions from HuggingFaceH4/no_robots with streaming=True.
    Used for task-aware routing correctness benchmarking.
    """
    if not HAS_HF_DATASETS:
        print("[WARN] datasets library not installed. Cannot stream no_robots.")
        return

    print(f"[STREAM] Connecting to HuggingFaceH4/no_robots (streaming=True)...")
    try:
        ds = load_dataset("HuggingFaceH4/no_robots", split="train", streaming=True)
        count = 0
        for row in ds:
            prompt = row.get("prompt", "")
            category = row.get("category", "general")
            if prompt:
                count += 1
                yield {
                    "id": f"norobots_{count:06d}",
                    "source": "HuggingFaceH4/no_robots",
                    "category": category,
                    "prompt": prompt.strip(),
                    "requires_json": category in ["coding", "summarization", "classification"]
                }
                if count >= limit:
                    break
        print(f"[OK] Streamed {count:,} labeled prompts from HuggingFaceH4/no_robots")
    except Exception as e:
        print(f"[ERR] NoRobots streaming exception: {e}")


def run_streaming_benchmark(limit: int = 500, dataset_choice: str = "wildchat"):
    """
    Executes a real-time streaming benchmark directly feeding prompts into InferRoute.
    """
    if dataset_choice == "wildchat":
        generator = stream_wildchat_prompts(limit=limit)
    else:
        generator = stream_no_robots_prompts(limit=limit)

    prompts = list(generator)
    print(f"[BENCHMARK] Loaded {len(prompts)} prompts from HF stream. Ready for gateway evaluation.")
    return prompts


if __name__ == "__main__":
    prompts = run_streaming_benchmark(limit=10, dataset_choice="norobots")
    for p in prompts[:3]:
        print("--- Sample Prompt ---")
        print(f"ID: {p['id']} | Category: {p['category']}")
        print(f"Content: {p['prompt'][:100]}...\n")