"""
data/loader.py — Downloads open-source datasets from Hugging Face and caches
them locally as JSON files for the SelfEvo search tool, QA tool, and task registry.

Datasets used (all open-source / permissively licensed):
  - gsm8k          (Google, MIT) — grade-school math word problems
  - ai2_arc        (AI2, CC BY 4.0) — ARC Challenge science questions
  - trivia_qa      (UW, Apache 2.0) — open-domain trivia Q&A
  - lucasmccabe/logiqa (LogiQA, public research) — logical reasoning
  - cosmos_qa      (AI2, CC BY 4.0) — commonsense reading comprehension

Each run checks if the cache files already exist and skips download if present.
"""

import json
import os
import logging
from pathlib import Path
from typing import Dict, List, Any

logger = logging.getLogger(__name__)

DATA_DIR = Path(__file__).parent
CACHE = {
    "knowledge_base": DATA_DIR / "knowledge_base.json",
    "gsm8k":          DATA_DIR / "gsm8k_samples.json",
    "arc":            DATA_DIR / "arc_samples.json",
    "trivia_qa":      DATA_DIR / "trivia_qa_samples.json",
    "logiqa":         DATA_DIR / "logiqa_samples.json",
}

# ── Max samples to pull per dataset (keep small for fast startup) ─────────────
LIMITS = {
    "gsm8k":     80,
    "arc":       60,
    "trivia_qa": 80,
    "logiqa":    40,
}


# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _save(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    logger.info("Saved %s (%d items)", path.name, len(data) if isinstance(data, list) else len(data.get("entries", [])))


def _load(path: Path) -> Any:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _hf_available() -> bool:
    try:
        import datasets  # noqa: F401
        return True
    except ImportError:
        return False


# ─────────────────────────────────────────────────────────────────────────────
# Individual dataset downloaders
# ─────────────────────────────────────────────────────────────────────────────

def _download_gsm8k(limit: int) -> List[Dict]:
    from datasets import load_dataset
    ds = load_dataset("gsm8k", "main", split="train", trust_remote_code=False)
    samples = []
    for item in ds.select(range(min(limit, len(ds)))):
        q = item["question"].strip()
        # Extract the final numeric answer (after ####)
        ans_raw = item["answer"]
        if "####" in ans_raw:
            ans = ans_raw.split("####")[-1].strip().replace(",", "")
        else:
            # last number in the answer text
            import re
            nums = re.findall(r"[\d,]+\.?\d*", ans_raw.replace(",", ""))
            ans = nums[-1] if nums else ans_raw.strip()
        samples.append({"question": q, "answer": ans, "source": "gsm8k"})
    return samples


def _download_arc(limit: int) -> List[Dict]:
    from datasets import load_dataset
    ds = load_dataset("ai2_arc", "ARC-Challenge", split="train", trust_remote_code=False)
    samples = []
    for item in ds.select(range(min(limit, len(ds)))):
        q = item["question"].strip()
        # Get the correct answer text from choices
        choices = item["choices"]
        label = item["answerKey"]
        labels = choices["label"]
        texts  = choices["text"]
        ans = ""
        for lbl, txt in zip(labels, texts):
            if lbl == label:
                ans = txt
                break
        if ans:
            samples.append({"question": q, "answer": ans, "source": "arc_challenge"})
    return samples


def _download_trivia_qa(limit: int) -> List[Dict]:
    from datasets import load_dataset
    ds = load_dataset("trivia_qa", "rc.nocontext", split="train", trust_remote_code=False)
    samples = []
    for item in ds.select(range(min(limit, len(ds)))):
        q = item["question"].strip()
        # Use the first alias as the canonical answer
        aliases = item["answer"].get("aliases", [])
        ans = aliases[0] if aliases else item["answer"].get("value", "")
        if q and ans:
            samples.append({"question": q, "answer": ans, "source": "trivia_qa"})
    return samples


def _download_logiqa(limit: int) -> List[Dict]:
    """
    Downloads LogiQA from lucasmccabe/logiqa (public parquet dataset).
    Falls back to empty list if unavailable — non-fatal.
    """
    from datasets import load_dataset
    # lucasmccabe/logiqa is public and parquet-based (no legacy scripts)
    candidates = ["lucasmccabe/logiqa", "EleutherAI/logiqa"]
    ds = None
    for repo in candidates:
        try:
            ds = load_dataset(repo, split="train", trust_remote_code=False)
            logger.info("LogiQA loaded from %s", repo)
            break
        except Exception as exc:
            logger.warning("Could not load LogiQA from %s: %s", repo, exc)
    if ds is None:
        logger.warning("All LogiQA sources failed — using empty fallback.")
        return []
    samples = []
    for item in ds.select(range(min(limit, len(ds)))):
        q = item.get("query") or item.get("question") or item.get("context", "")
        options = item.get("options") or item.get("choices", [])
        label = item.get("correct_option") or item.get("label", 0)
        if isinstance(label, str):
            label_map = {"a": 0, "b": 1, "c": 2, "d": 3}
            label = label_map.get(label.lower(), 0)
        ans = options[label] if isinstance(options, list) and label < len(options) else str(label)
        if isinstance(q, str) and q.strip():
            samples.append({"question": q.strip(), "answer": ans, "options": options, "source": "logiqa"})
    return samples


# ─────────────────────────────────────────────────────────────────────────────
# Knowledge-base builder
# ─────────────────────────────────────────────────────────────────────────────

def _build_knowledge_base(
    gsm8k_data: List[Dict],
    arc_data:   List[Dict],
    trivia_data:List[Dict],
    logiqa_data:List[Dict],
) -> Dict:
    """
    Merges all dataset QA pairs + hard-coded domain facts into a
    unified knowledge base used by the SearchTool.
    """
    entries = []

    # ── Hard-coded authoritative facts (always present) ───────────────────
    static_facts = [
        # CS / Algorithms
        {"topic": "fibonacci", "content": "Fibonacci sequence: 0,1,1,2,3,5,8,13,21,34,55,89,144,233,377,610,987. fib(10)=55, fib(20)=6765, fib(30)=832040."},
        {"topic": "prime numbers", "content": "Prime numbers: 2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,59,61,67,71,73,79,83,89,97. Sum of primes below 50 = 328."},
        {"topic": "sorting algorithms", "content": "Merge sort: O(n log n) stable. Quicksort: O(n log n) average, O(n^2) worst. Bubble sort: O(n^2). Heapsort: O(n log n)."},
        {"topic": "binary search", "content": "Binary search: O(log n) on sorted arrays. Finds mid=(lo+hi)//2, compares with target, narrows range."},
        {"topic": "dynamic programming", "content": "DP stores subproblem results (memoization/tabulation). Used for: Fibonacci, LCS, Knapsack, Shortest path."},
        {"topic": "big o notation", "content": "O(1) constant. O(log n) logarithmic. O(n) linear. O(n log n) linearithmic. O(n^2) quadratic. O(2^n) exponential."},
        {"topic": "recursion", "content": "Recursion: function calls itself. Must have base case. Stack overflow if infinite. Tail recursion optimizable."},
        {"topic": "gcd", "content": "GCD via Euclidean: gcd(a,b)=gcd(b,a%b) until b=0. gcd(1071,462)=21. gcd(48,18)=6."},
        {"topic": "factorial", "content": "n! = n*(n-1)*…*1. 0!=1, 1!=1, 5!=120, 7!=5040, 10!=3628800, 12!=479001600."},
        {"topic": "graph algorithms", "content": "BFS: O(V+E), shortest path unweighted. DFS: O(V+E), cycle detection. Dijkstra: O(E log V), weighted shortest path. Bellman-Ford: handles negative weights."},
        {"topic": "hash table", "content": "Hash table: O(1) average lookup/insert. Handles collisions via chaining or open addressing. Load factor affects performance."},
        {"topic": "binary tree", "content": "Binary tree traversal: Inorder (left-root-right), Preorder (root-left-right), Postorder (left-right-root). BST: left < root < right."},
        {"topic": "linked list", "content": "Linked list: O(1) insert/delete at head, O(n) search. Doubly linked: bidirectional traversal. Circular: last node points to first."},
        {"topic": "stack queue", "content": "Stack: LIFO, push/pop O(1). Queue: FIFO, enqueue/dequeue O(1). Deque: both ends O(1)."},
        {"topic": "python syntax", "content": "Python list comprehension: [x*2 for x in range(10)]. Lambda: lambda x: x+1. Dict comprehension: {k:v for k,v in items}."},
        # Math
        {"topic": "area geometry", "content": "Circle area=π*r². Triangle area=0.5*base*height. Rectangle=length*width. Trapezoid=0.5*(a+b)*height."},
        {"topic": "pythagoras", "content": "Pythagorean theorem: a²+b²=c². Common triples: (3,4,5), (5,12,13), (8,15,17), (7,24,25)."},
        {"topic": "percentage", "content": "Percentage: X% of Y = X/100 * Y. Percentage increase: ((new-old)/old)*100. Compound interest: A=P(1+r/n)^(nt)."},
        {"topic": "statistics mean median mode", "content": "Mean=sum/count. Median=middle value (sorted). Mode=most frequent. Standard deviation=sqrt(variance)."},
        {"topic": "quadratic equation", "content": "Quadratic ax²+bx+c=0. Solution: x=(-b±√(b²-4ac))/(2a). Discriminant b²-4ac: >0 two real roots, =0 one root, <0 complex roots."},
        {"topic": "logarithm", "content": "log(a*b)=log(a)+log(b). log(a/b)=log(a)-log(b). log(a^n)=n*log(a). ln(e)=1. log10(100)=2."},
        {"topic": "trigonometry", "content": "sin(0)=0, sin(30)=0.5, sin(45)=√2/2, sin(60)=√3/2, sin(90)=1. cos is complement of sin. tan=sin/cos."},
        {"topic": "number theory", "content": "Prime factorization: unique. LCM(a,b)=a*b/GCD(a,b). Modular arithmetic: (a+b)%n=(a%n+b%n)%n."},
        {"topic": "combinations permutations", "content": "Permutations: P(n,r)=n!/(n-r)!. Combinations: C(n,r)=n!/(r!*(n-r)!). C(5,2)=10, C(10,3)=120."},
        {"topic": "distance speed time", "content": "Distance=Speed*Time. Average speed=total distance/total time. Relative speed (same dir)=|v1-v2|, opposite=v1+v2."},
        # Science
        {"topic": "newton laws", "content": "Newton's 1st: inertia. 2nd: F=ma. 3rd: action=reaction. G=6.674×10⁻¹¹ N m²/kg²."},
        {"topic": "periodic table elements", "content": "H(1), He(2), Li(3), C(6), N(7), O(8), Na(11), Mg(12), Al(13), Si(14), P(15), S(16), Cl(17), Ar(18), K(19), Ca(20), Fe(26), Cu(29), Zn(30), Ag(47), Au(79), Hg(80), Pb(82)."},
        {"topic": "photosynthesis", "content": "Photosynthesis: 6CO₂ + 6H₂O + light → C₆H₁₂O₆ + 6O₂. Occurs in chloroplasts. Chlorophyll absorbs red and blue light."},
        {"topic": "electricity", "content": "Ohm's law: V=IR. Power: P=IV=I²R=V²/R. Series: R_total=R1+R2+... Parallel: 1/R_total=1/R1+1/R2+..."},
        {"topic": "states of matter", "content": "Solid: fixed shape/volume. Liquid: fixed volume, variable shape. Gas: variable shape/volume. Plasma: ionized gas."},
        {"topic": "dna genetics", "content": "DNA: double helix, base pairs A-T and G-C. RNA: single strand, U replaces T. Codon: 3 bases = 1 amino acid. 64 codons, 20 amino acids."},
        {"topic": "speed of light", "content": "Speed of light c=3×10⁸ m/s. Light-year=9.46×10¹⁵ m. Sun to Earth: 8 min. Sound in air: 343 m/s at 20°C."},
        {"topic": "chemical reactions", "content": "Exothermic: releases heat (combustion, respiration). Endothermic: absorbs heat (photosynthesis, melting). Catalyst: speeds reaction without being consumed."},
        # Geography / History
        {"topic": "world capitals", "content": "USA: Washington D.C. UK: London. France: Paris. Germany: Berlin. Japan: Tokyo. China: Beijing. India: New Delhi. Russia: Moscow. Brazil: Brasília. Australia: Canberra."},
        {"topic": "world geography", "content": "Largest continent: Asia. Largest ocean: Pacific. Longest river: Nile (6,650 km). Highest mountain: Everest (8,848 m). Amazon is largest by flow."},
        {"topic": "world history", "content": "WW1: 1914-1918. WW2: 1939-1945. French Revolution: 1789. American Independence: 1776. Renaissance: 14th-17th century. Industrial Revolution: 18th-19th century."},
        {"topic": "us presidents", "content": "1st: Washington. 16th: Lincoln. 32nd: FDR. 35th: JFK. 44th: Obama. 45th: Trump. 46th: Biden. 47th: Trump."},
        {"topic": "planet solar system", "content": "Planets (order): Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, Neptune. Largest: Jupiter. Smallest: Mercury. Hottest: Venus."},
        # CS/AI
        {"topic": "machine learning", "content": "Supervised: labeled data (classification, regression). Unsupervised: unlabeled (clustering, dimensionality reduction). Reinforcement: reward signals. Deep learning uses neural networks."},
        {"topic": "neural network", "content": "Layers: input, hidden, output. Activation: ReLU, sigmoid, tanh, softmax. Backpropagation updates weights via gradient descent. CNN for images, RNN/LSTM for sequences, Transformer for NLP."},
        {"topic": "git version control", "content": "git init/clone/add/commit/push/pull. Branch: git checkout -b. Merge: git merge. Rebase: git rebase. Stash: git stash. Reset: git reset --hard."},
        {"topic": "database sql", "content": "SELECT, FROM, WHERE, GROUP BY, HAVING, ORDER BY, LIMIT. JOIN types: INNER, LEFT, RIGHT, FULL. ACID: Atomicity, Consistency, Isolation, Durability."},
        {"topic": "operating system", "content": "Process vs Thread: thread shares memory. Scheduling: FIFO, Round Robin, Priority. Deadlock: mutual exclusion, hold-and-wait, no preemption, circular wait. Virtual memory via paging."},
        {"topic": "networking tcp ip", "content": "OSI layers: Physical, Data Link, Network, Transport, Session, Presentation, Application. TCP: reliable, ordered. UDP: unreliable, fast. HTTP port 80, HTTPS 443, SSH 22, DNS 53."},
        # Logic puzzles
        {"topic": "snail pole", "content": "Snail climbs 3m/day, slides 2m/night. Net 1m/day. For 20m pole: reaches top on day 18 (climbs 3m on day 18 from 17m, reaches 20m without sliding)."},
        {"topic": "water jug puzzle", "content": "3L and 5L jugs to measure 4L: Fill 5L, pour into 3L (2L left in 5L), empty 3L, pour 2L into 3L, fill 5L again, fill 3L from 5L (add 1L needed) → 4L in 5L."},
        {"topic": "hat puzzle logicians", "content": "If A and B don't know their hat color but C does: means A and B each see at least one red hat. C sees two people uncertain → C's hat must be red (all 3 red)."},
        {"topic": "einstein riddle", "content": "Logic grid puzzles: use process of elimination with given clues. Assign attributes (nationality, pet, drink, etc.) to positions via constraint propagation."},
    ]

    entries.extend(static_facts)

    # ── GSM8K: convert to search-friendly topic entries ─────────────────────
    for item in gsm8k_data[:40]:  # use first 40 for KB
        entries.append({
            "topic": f"math word problem",
            "content": f"Q: {item['question']} A: {item['answer']}",
            "source": "gsm8k",
        })

    # ── ARC: science facts from questions ───────────────────────────────────
    for item in arc_data[:30]:
        entries.append({
            "topic": "science question",
            "content": f"Q: {item['question']} A: {item['answer']}",
            "source": "arc_challenge",
        })

    # ── TriviaQA: factual Q&A ────────────────────────────────────────────────
    for item in trivia_data[:40]:
        entries.append({
            "topic": "trivia factual",
            "content": f"Q: {item['question']} A: {item['answer']}",
            "source": "trivia_qa",
        })

    # ── LogiQA: reasoning ────────────────────────────────────────────────────
    for item in logiqa_data[:20]:
        entries.append({
            "topic": "logical reasoning",
            "content": f"Q: {item['question']} A: {item['answer']}",
            "source": "logiqa",
        })

    return {"entries": entries, "total": len(entries)}


# ─────────────────────────────────────────────────────────────────────────────
# Main entry point
# ─────────────────────────────────────────────────────────────────────────────

def load_all(force_refresh: bool = False) -> Dict[str, Any]:
    """
    Downloads datasets from Hugging Face and caches them locally.
    Returns dict with loaded data for all datasets.
    Skips download if cache exists (unless force_refresh=True).
    """
    if not _hf_available():
        logger.error(
            "Hugging Face `datasets` library not installed. "
            "Run: pip install datasets"
        )
        return _load_fallback()

    results = {}

    # ── GSM8K ────────────────────────────────────────────────────────────────
    if force_refresh or not CACHE["gsm8k"].exists():
        logger.info("Downloading GSM8K from Hugging Face…")
        data = _download_gsm8k(LIMITS["gsm8k"])
        _save(CACHE["gsm8k"], data)
    else:
        logger.info("Loading GSM8K from cache…")
        data = _load(CACHE["gsm8k"])
    results["gsm8k"] = data

    # ── ARC Challenge ────────────────────────────────────────────────────────
    if force_refresh or not CACHE["arc"].exists():
        logger.info("Downloading ARC Challenge from Hugging Face…")
        data = _download_arc(LIMITS["arc"])
        _save(CACHE["arc"], data)
    else:
        logger.info("Loading ARC from cache…")
        data = _load(CACHE["arc"])
    results["arc"] = data

    # ── TriviaQA ─────────────────────────────────────────────────────────────
    if force_refresh or not CACHE["trivia_qa"].exists():
        logger.info("Downloading TriviaQA from Hugging Face…")
        data = _download_trivia_qa(LIMITS["trivia_qa"])
        _save(CACHE["trivia_qa"], data)
    else:
        logger.info("Loading TriviaQA from cache…")
        data = _load(CACHE["trivia_qa"])
    results["trivia_qa"] = data

    # ── LogiQA (non-fatal) ───────────────────────────────────────────────────
    if force_refresh or not CACHE["logiqa"].exists():
        logger.info("Downloading LogiQA from Hugging Face…")
        try:
            data = _download_logiqa(LIMITS["logiqa"])
            if data:
                _save(CACHE["logiqa"], data)
        except Exception as exc:
            logger.warning("LogiQA download failed, continuing without it: %s", exc)
            data = []
    else:
        logger.info("Loading LogiQA from cache…")
        try:
            data = _load(CACHE["logiqa"])
        except Exception:
            data = []
    results["logiqa"] = data

    # ── Build unified knowledge base ─────────────────────────────────────────
    if force_refresh or not CACHE["knowledge_base"].exists():
        logger.info("Building unified knowledge base…")
        kb = _build_knowledge_base(
            results["gsm8k"],
            results["arc"],
            results["trivia_qa"],
            results["logiqa"],
        )
        _save(CACHE["knowledge_base"], kb)
    else:
        logger.info("Loading knowledge base from cache…")
        kb = _load(CACHE["knowledge_base"])
    results["knowledge_base"] = kb

    total = kb.get("total", len(kb.get("entries", [])))
    logger.info("Data loader complete — %d KB entries across all datasets.", total)
    return results


def _load_fallback() -> Dict[str, Any]:
    """Returns empty structures if datasets library is unavailable."""
    logger.warning("Falling back to empty datasets — install `datasets` for full data.")
    return {
        "gsm8k": [],
        "arc": [],
        "trivia_qa": [],
        "logiqa": [],
        "knowledge_base": {"entries": [], "total": 0},
    }


def get_knowledge_base() -> List[Dict]:
    """Convenience function: load just the KB entries list."""
    if CACHE["knowledge_base"].exists():
        kb = _load(CACHE["knowledge_base"])
        return kb.get("entries", [])
    return []


def get_qa_pairs(source: str = None) -> List[Dict]:
    """
    Returns all QA pairs across datasets.
    Optionally filter by source: 'gsm8k', 'arc_challenge', 'trivia_qa', 'logiqa'.
    """
    all_pairs = []
    for key in ["gsm8k", "arc", "trivia_qa", "logiqa"]:
        if CACHE[key].exists():
            data = _load(CACHE[key])
            if source is None or any(item.get("source") == source for item in data):
                all_pairs.extend(
                    item for item in data
                    if source is None or item.get("source") == source
                )
    return all_pairs


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
    load_all(force_refresh=True)
    print("All datasets downloaded and cached.")