"""
TD Dataset Finder v3 — HARD DATA EDITION

v3 changes (March 2026):
- HARD datasets first, easy ones last (model at 95% needs challenge, not basics)
- Added: competition math (MATH/AMC/AIME), PhD-level knowledge (GPQA),
  competitive programming (code_contests, TACO), evolved instructions (WizardLM)
- Removed easy-first ordering that caused training loss = 0.0000
- Fixed: MMLU-Pro split, HellaSwag split

v2 changes:
- Known-good fallback datasets tried FIRST (reliable, parseable)
- HuggingFace search only used if fallbacks don't have enough
- More fallback datasets added per category (4-5 each)
- Retention floor increased to 20% (prevents catastrophic forgetting)
- Stall detection in fallback downloads (skip after 500 unparseable items)

Takes the weakness report from td_weakness.py and finds training data
to fix those weaknesses. Searches multiple sources in priority order.

This is Step 3 of the self-improvement loop.

Sources (in order of priority):
  1. Known-good fallback datasets (curated, tested, parseable)
  2. HuggingFace Datasets API search (if fallbacks aren't enough)
  3. GitHub — repos with training data, benchmarks, problem sets
  4. Kaggle — competitions and datasets
  5. (Future) Web scraping — Stack Overflow, math forums, etc.

Input:
  - WeaknessReport from td_weakness.py

Output:
  - Formatted training data in chat template, ready for training
  - Data allocation follows the weakness report (worse = more data)
"""

import json
import os
import random
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field


# ============================================================
# DATA STRUCTURES
# ============================================================

@dataclass
class DatasetResult:
    """A dataset found from searching."""
    name: str
    source: str  # "huggingface", "github", "kaggle", "web"
    url: str
    description: str
    size_estimate: int  # approximate number of samples
    relevance_score: float  # 0-1, how well it matches the weakness
    category: str  # which weakness category it's for


@dataclass
class FormattedSample:
    """A single training sample in chat template format."""
    system: str
    user: str
    assistant: str
    category: str
    source: str


@dataclass
class TrainingData:
    """Complete training data package ready for the trainer."""
    samples: List[FormattedSample]
    sources_used: List[DatasetResult]
    allocation: Dict[str, int]  # category -> actual samples collected
    total_samples: int
    categories_covered: List[str]
    categories_missing: List[str]  # categories we couldn't find data for


# ============================================================
# CHAT TEMPLATE FORMATTER
# ============================================================

def format_chat_template(
    system: str,
    user: str,
    assistant: str,
) -> str:
    """
    Format a sample into Qwen3-VL's chat template.
    This is the <|im_start|> format that selfimprove.py already uses.
    """
    parts = []
    if system:
        parts.append(f"<|im_start|>system\n{system}<|im_end|>")
    parts.append(f"<|im_start|>user\n{user}<|im_end|>")
    parts.append(f"<|im_start|>assistant\n{assistant}<|im_end|>")
    return "\n".join(parts)


def format_sample_to_dict(sample: FormattedSample) -> Dict:
    """Convert a FormattedSample to the dict format selfimprove.py expects."""
    return {
        "text": format_chat_template(sample.system, sample.user, sample.assistant),
        "category": sample.category,
        "source": sample.source,
    }


# ============================================================
# HUGGINGFACE SEARCH
# ============================================================

def search_huggingface(
    keywords: List[str],
    tags: List[str],
    category: str,
    max_datasets: int = 5,
) -> List[DatasetResult]:
    """
    Search HuggingFace Datasets API for matching datasets.

    Uses the huggingface_hub library to search by keyword and tags.
    Returns metadata about found datasets (doesn't download yet).
    """
    results = []

    try:
        from huggingface_hub import HfApi
        api = HfApi()

        # Search by each keyword
        seen = set()
        for keyword in keywords[:3]:  # Top 3 keywords
            try:
                datasets = api.list_datasets(
                    search=keyword,
                    sort="downloads",
                    direction=-1,
                    limit=max_datasets,
                )

                for ds in datasets:
                    if ds.id in seen:
                        continue
                    seen.add(ds.id)

                    # Estimate relevance based on tag overlap
                    ds_tags = set(ds.tags or [])
                    tag_overlap = len(ds_tags.intersection(set(tags)))
                    relevance = min(0.5 + (tag_overlap * 0.15), 1.0)

                    # Boost if keyword appears in name
                    if any(kw.lower() in ds.id.lower() for kw in keyword.split()):
                        relevance = min(relevance + 0.2, 1.0)

                    results.append(DatasetResult(
                        name=ds.id,
                        source="huggingface",
                        url=f"https://huggingface.co/datasets/{ds.id}",
                        description=ds.id,  # HF API doesn't always return description
                        size_estimate=ds.downloads or 0,  # Use downloads as proxy
                        relevance_score=relevance,
                        category=category,
                    ))
            except Exception as e:
                print(f"    Warning: HF search failed for '{keyword}': {e}")
                continue

    except ImportError:
        print("    Warning: huggingface_hub not installed. Install with: pip install huggingface_hub")
        return []

    # Sort by relevance
    results.sort(key=lambda r: r.relevance_score, reverse=True)
    return results[:max_datasets]


def download_huggingface_dataset(
    dataset_name: str,
    category: str,
    num_samples: int,
    split: str = "train",
) -> List[FormattedSample]:
    """
    Download and format samples from a HuggingFace dataset.

    Handles common dataset formats:
    - QA pairs (question/answer columns)
    - Instruction pairs (instruction/output columns)
    - Chat format (conversations column)
    - Text only (text column) — wrap in QA format
    """
    samples = []

    try:
        from datasets import load_dataset

        print(f"    Downloading {dataset_name} (split={split})...")
        try:
            ds = load_dataset(dataset_name, split=split, streaming=True)
        except Exception:
            # Try without split specification
            try:
                ds = load_dataset(dataset_name, streaming=True)
                # Get first available split
                if hasattr(ds, 'keys'):
                    split_name = list(ds.keys())[0]
                    ds = ds[split_name]
                else:
                    return []
            except Exception as e:
                print(f"    Failed to load {dataset_name}: {e}")
                return []

        # Collect samples (streaming so we don't download everything)
        # Timeout: 5 min per dataset to prevent hanging on slow/broken sources
        import time as _time
        download_start = _time.time()
        download_timeout = 300  # 5 minutes max per dataset
        count = 0
        stall_count = 0
        for item in ds:
            if count >= num_samples:
                break
            if _time.time() - download_start > download_timeout:
                print(f"    Timeout: {dataset_name} took >{download_timeout}s — moving on with {count} samples")
                break

            sample = _format_hf_item(item, category, dataset_name)
            if sample:
                samples.append(sample)
                count += 1
                stall_count = 0
            else:
                stall_count += 1
                # If 500 consecutive items produce nothing, dataset format is wrong
                if stall_count > 500:
                    print(f"    {dataset_name}: 500 consecutive unparseable items — skipping")
                    break

        print(f"    Got {len(samples)} samples from {dataset_name}")

    except ImportError:
        print("    Warning: datasets library not installed. Install with: pip install datasets")
    except Exception as e:
        print(f"    Error downloading {dataset_name}: {e}")

    return samples


def _format_hf_item(item: Dict, category: str, source: str) -> Optional[FormattedSample]:
    """
    Try to format a HuggingFace dataset item into a chat sample.
    Handles many common column naming conventions.
    """
    # Guard: item must be a dict
    if not isinstance(item, dict):
        return None

    # Common column name patterns for user/assistant content
    user_keys = ["question", "instruction", "input", "prompt", "query", "problem", "text"]
    assistant_keys = ["answer", "output", "response", "solution", "target", "completion"]
    system_keys = ["system", "system_prompt", "context"]

    user_text = None
    assistant_text = None
    system_text = ""

    # Try to find user content (guard: value might be non-string like int, list, etc.)
    for key in user_keys:
        if key in item and item[key]:
            val = item[key]
            if isinstance(val, (str, int, float)):
                user_text = str(val).strip()
                if user_text:
                    break
            # Skip lists, dicts, None-like values

    # Try to find assistant content
    for key in assistant_keys:
        if key in item and item[key]:
            val = item[key]
            if isinstance(val, (str, int, float)):
                assistant_text = str(val).strip()
                if assistant_text:
                    break

    # Try to find system content
    for key in system_keys:
        if key in item and item[key]:
            val = item[key]
            if isinstance(val, (str, int, float)):
                system_text = str(val).strip()
                break

    # Handle conversation format
    if not user_text and "conversations" in item:
        convs = item["conversations"]
        if isinstance(convs, list) and len(convs) >= 2:
            for c in convs:
                if not isinstance(c, dict):
                    continue  # Skip malformed conversation entries
                role = str(c.get("role", c.get("from", ""))).strip()
                content = c.get("content", c.get("value", ""))
                if content is None:
                    continue
                content = str(content).strip()
                if role in ["user", "human"] and not user_text:
                    user_text = content
                elif role in ["assistant", "gpt", "bot"] and not assistant_text:
                    assistant_text = content
                elif role == "system" and not system_text:
                    system_text = content

    # Handle messages format (OpenAI-style)
    if not user_text and "messages" in item:
        msgs = item["messages"]
        if isinstance(msgs, list):
            for m in msgs:
                if not isinstance(m, dict):
                    continue  # Skip malformed message entries
                role = str(m.get("role", "")).strip()
                content = m.get("content", "")
                if content is None:
                    continue
                content = str(content).strip()
                if role == "user" and not user_text:
                    user_text = content
                elif role == "assistant" and not assistant_text:
                    assistant_text = content
                elif role == "system" and not system_text:
                    system_text = content

    # Handle MMLU-style multiple choice: question + choices + answer_index
    if not assistant_text and "choices" in item and "answer" in item:
        choices = item["choices"]
        answer_idx = item["answer"]
        if isinstance(choices, list) and len(choices) > 0 and len(choices) <= 26:
            # Format as multiple choice question (cap at 26 = A-Z)
            if user_text:
                choice_strs = [f"{chr(65+i)}. {c}" for i, c in enumerate(choices)]
                user_text = user_text + "\n" + "\n".join(choice_strs)
            if isinstance(answer_idx, int) and 0 <= answer_idx < len(choices):
                assistant_text = f"{chr(65+answer_idx)}. {choices[answer_idx]}"
            elif isinstance(answer_idx, str) and len(answer_idx) == 1 and answer_idx.isalpha():
                idx = ord(answer_idx.upper()) - 65
                if 0 <= idx < len(choices):
                    assistant_text = f"{answer_idx.upper()}. {choices[idx]}"

    # Handle GSM8K-style: answer contains "#### final_number"
    if assistant_text and "####" in assistant_text:
        # Keep the full chain-of-thought — it's good training data for reasoning
        pass

    # Need at least user + assistant
    if not user_text or not assistant_text:
        return None

    # Skip very short or very long samples
    if len(user_text) < 10 or len(assistant_text) < 5:  # Allow short answers (like "42")
        return None
    if len(user_text) > 5000 or len(assistant_text) > 10000:
        return None

    return FormattedSample(
        system=system_text,
        user=user_text,
        assistant=assistant_text,
        category=category,
        source=source,
    )


# ============================================================
# WELL-KNOWN DATASETS PER CATEGORY
# ============================================================

# If the search finds nothing, fall back to these known-good datasets
# Verified dataset names on HuggingFace (March 2026)
# Each tuple: (dataset_name, config_or_subset)
# If a dataset fails to load, it moves to the next one silently
# Each tuple: (dataset_name, config_or_subset, split_override)
# split_override: if None, uses "train". Set to specific split name if needed.
# If a dataset fails to load, it moves to the next one silently.
#
# IMPORTANT: HARD datasets come FIRST. The model is at 95% — it already knows
# basic stuff like GSM8K and alpaca. Training loss = 0 on easy data.
# We need competition math, PhD-level knowledge, complex code, etc.
FALLBACK_DATASETS = {
    "math": [
        # HARD FIRST — competition math, not grade school
        ("lighteval/MATH", "all", None),                     # 12.5K, AMC/AIME competition math
        ("AI-MO/NuminaMath-CoT", None, None),               # 860K chain-of-thought math (diverse difficulty)
        ("TIGER-Lab/MathInstruct", None, None),             # 262K math instruction pairs (mixed difficulty)
        ("microsoft/orca-math-word-problems-200k", None, None),  # 200K word problems (harder than GSM8K)
        ("openai/gsm8k", "main", None),                     # 8.8K grade school — LAST (easiest)
    ],
    "code": [
        # HARD FIRST — competitive programming, real-world code
        ("deepmind/code_contests", None, None),             # Google Code Jam / Codeforces level
        ("BAAI/TACO", None, None),                          # Hard algorithmic problems
        ("m-a-p/CodeFeedback-Filtered-Instruction", None, None),  # 157K code instructions
        ("sahil2801/CodeAlpaca-20k", None, None),            # 20K code instruction pairs
        ("openai/openai_humaneval", None, "test"),           # 164 problems — split is "test"
    ],
    "reasoning": [
        # HARD FIRST — multi-step reasoning, not simple logic
        ("allenai/ai2_arc", "ARC-Challenge", None),         # 2.6K challenge set (hard split)
        ("Rowan/hellaswag", None, "validation"),             # 40K commonsense (use validation)
        ("allenai/winogrande", "winogrande_xl", None),       # 40K commonsense reasoning
        ("allenai/openbookqa", "main", None),                # Science reasoning
        ("teknium/OpenHermes-2.5", None, None),             # Diverse reasoning in conversations
    ],
    "creativity": [
        # HARD FIRST — evolved/complex instructions, not basic templates
        ("WizardLMTeam/WizardLM_evol_instruct_V2_196k", None, None),  # 196K evolved (HARDER by design)
        ("teknium/OpenHermes-2.5", None, None),             # 1M diverse instruction pairs
        ("Open-Orca/OpenOrca", None, None),                  # 4.2M conversations — diverse
        ("tatsu-lab/alpaca", None, None),                    # 52K instructions — LAST (easiest)
    ],
    "knowledge": [
        # HARD FIRST — graduate-level, not basic trivia
        ("Idavidrein/gpqa", "gpqa_extended", None),         # Graduate-level science (PhD-level!)
        ("TIGER-Lab/MMLU-Pro", None, "test"),               # Harder MMLU (10 choices, harder questions)
        ("cais/mmlu", "all", "test"),                        # Classic MMLU — split is "test"
        ("cais/mmlu", "all", "validation"),                  # Also try validation split
        ("truthfulqa/truthful_qa", "generation", "validation"),  # Truthfulness — split is "validation"
        ("lukaemon/mmlu", "geography", "test"),              # Geography specifically (weak spot!)
        ("lukaemon/mmlu", "astronomy", "test"),              # Science knowledge (weak spot!)
        ("lukaemon/mmlu", "world_religions", "test"),        # More knowledge topics
        ("lukaemon/mmlu", "global_facts", "test"),           # General knowledge
    ],
    "instruction_following": [
        # HARD FIRST — complex format constraints
        ("HuggingFaceH4/ultrafeedback_binarized", None, "train_prefs"),  # 61K preference pairs
        ("WizardLMTeam/WizardLM_evol_instruct_V2_196k", None, None),  # Evolved = harder instructions
        ("tatsu-lab/alpaca", None, None),                    # 52K instructions
        ("Intel/orca_dpo_pairs", None, None),                # 12K DPO pairs
    ],
}


# ============================================================
# MAIN DATA FINDER
# ============================================================

def find_training_data(
    search_plan: List[Dict],
    data_allocation: Dict[str, int],
    output_dir: str = "td_fuse_outputs/training_data",
    hf_token: Optional[str] = None,
) -> TrainingData:
    """
    Search for and download training data based on the weakness report.

    This is the main entry point that td_loop.py calls.

    Args:
        search_plan: From WeaknessReport.search_plan
        data_allocation: From WeaknessReport.data_allocation
        output_dir: Where to save the formatted data
        hf_token: HuggingFace token for private datasets

    Returns:
        TrainingData with formatted samples ready for training
    """
    print("\n" + "=" * 60)
    print("TD DATASET FINDER v1")
    print("=" * 60)

    if hf_token:
        os.environ["HF_TOKEN"] = hf_token

    all_samples = []
    all_sources = []
    categories_covered = []
    categories_missing = []

    for plan_item in search_plan:
        category = plan_item["category"]
        needed = data_allocation.get(category, 1000)
        keywords = plan_item["keywords"]
        hf_tags = plan_item.get("hf_tags", [])

        print(f"\n--- Searching for: {category} ({needed} samples needed) ---")

        category_samples = []

        # Step 1: Try KNOWN-GOOD fallback datasets FIRST (reliable, parseable)
        fallbacks = FALLBACK_DATASETS.get(category, [])
        if fallbacks:
            print(f"  Trying {len(fallbacks)} known-good datasets for {category}...")
            try:
                from datasets import load_dataset  # Import once outside loop
            except ImportError:
                print("    Warning: datasets library not installed — cannot load datasets")
                fallbacks = []  # Skip the fallback loop
            for fallback_entry in fallbacks:
                if len(category_samples) >= needed:
                    break

                # Unpack: (name, config, split_override) — split_override can be None
                ds_name = fallback_entry[0]
                config = fallback_entry[1]
                split_override = fallback_entry[2] if len(fallback_entry) > 2 else None
                use_split = split_override or "train"

                remaining = needed - len(category_samples)
                try:
                    if config:
                        ds = load_dataset(ds_name, config, streaming=True, split=use_split)
                    else:
                        ds = load_dataset(ds_name, streaming=True, split=use_split)

                    count = 0
                    stall = 0
                    for item in ds:
                        if count >= remaining:
                            break
                        sample = _format_hf_item(item, category, ds_name)
                        if sample:
                            category_samples.append(sample)
                            count += 1
                            stall = 0
                        else:
                            stall += 1
                            if stall > 500:
                                break

                    if count > 0:
                        print(f"    Got {count} from: {ds_name}")
                        all_sources.append(DatasetResult(
                            name=ds_name,
                            source="huggingface_known",
                            url=f"https://huggingface.co/datasets/{ds_name}",
                            description=f"Known-good dataset for {category}",
                            size_estimate=count,
                            relevance_score=0.8,
                            category=category,
                        ))
                except Exception as e:
                    print(f"    {ds_name} failed: {e}")
                    continue

        # Step 2: If still not enough, search HuggingFace for more
        if len(category_samples) < needed:
            print(f"  Got {len(category_samples)}/{needed} from known datasets. Searching HuggingFace...")
            hf_results = search_huggingface(keywords, hf_tags, category)

            if hf_results:
                print(f"  Found {len(hf_results)} datasets on HuggingFace:")
                for r in hf_results[:3]:
                    print(f"    - {r.name} (relevance: {r.relevance_score:.2f})")

            # Download from best matches
            for result in hf_results:
                if len(category_samples) >= needed:
                    break

                # Skip datasets we already used as fallbacks
                if any(result.name == fb_entry[0] for fb_entry in fallbacks):
                    continue

                remaining = needed - len(category_samples)
                samples = download_huggingface_dataset(
                    dataset_name=result.name,
                    category=category,
                    num_samples=remaining,
                )
                category_samples.extend(samples)
                if samples:
                    all_sources.append(result)

        # Track what we got
        if category_samples:
            all_samples.extend(category_samples)
            categories_covered.append(category)
            print(f"  Total for {category}: {len(category_samples)} samples")
        else:
            categories_missing.append(category)
            print(f"  WARNING: No data found for {category}!")

    # Shuffle all samples together
    random.shuffle(all_samples)

    # Save to disk
    output_path = Path(output_dir)
    try:
        output_path.mkdir(parents=True, exist_ok=True)
    except OSError as e:
        print(f"\n  ERROR: Cannot create output dir {output_dir}: {e}")
        # Still return what we collected even if can't save
        actual_allocation = {}
        for s in all_samples:
            actual_allocation[s.category] = actual_allocation.get(s.category, 0) + 1
        return TrainingData(
            samples=all_samples, sources_used=all_sources,
            allocation=actual_allocation, total_samples=len(all_samples),
            categories_covered=categories_covered, categories_missing=categories_missing,
        )

    # Save as JSONL (one sample per line)
    data_file = output_path / "training_data.jsonl"
    try:
        with open(data_file, "w") as f:
            for sample in all_samples:
                f.write(json.dumps(format_sample_to_dict(sample)) + "\n")
        print(f"\nSaved {len(all_samples)} samples to {data_file}")
    except OSError as e:
        print(f"\n  WARNING: Failed to save training data: {e} (disk may be full)")

    # Save metadata
    meta_file = output_path / "data_metadata.json"
    actual_allocation = {}
    for s in all_samples:
        actual_allocation[s.category] = actual_allocation.get(s.category, 0) + 1

    meta = {
        "total_samples": len(all_samples),
        "allocation": actual_allocation,
        "sources": [
            {"name": s.name, "source": s.source, "category": s.category, "url": s.url}
            for s in all_sources
        ],
        "categories_covered": categories_covered,
        "categories_missing": categories_missing,
    }
    try:
        with open(meta_file, "w") as f:
            json.dump(meta, f, indent=2)
        print(f"Saved metadata to {meta_file}")
    except OSError as e:
        print(f"  WARNING: Failed to save metadata: {e}")

    return TrainingData(
        samples=all_samples,
        sources_used=all_sources,
        allocation=actual_allocation,
        total_samples=len(all_samples),
        categories_covered=categories_covered,
        categories_missing=categories_missing,
    )


# ============================================================
# STANDALONE USAGE
# ============================================================

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python td_data_finder.py <weakness_report.json> [output_dir]")
        print("\nSearches for training data based on the weakness report.")
        sys.exit(1)

    # Load weakness report
    from td_weakness import load_report
    report = load_report(sys.argv[1])

    output_dir = sys.argv[2] if len(sys.argv) > 2 else "td_fuse_outputs/training_data"

    # Find data
    training_data = find_training_data(
        search_plan=report.search_plan,
        data_allocation=report.data_allocation,
        output_dir=output_dir,
    )

    print(f"\nDone! Got {training_data.total_samples} samples across {len(training_data.categories_covered)} categories.")
    if training_data.categories_missing:
        print(f"Missing data for: {', '.join(training_data.categories_missing)}")