"""FATHOM deterministic dataset generator — DATA-01..06.

Produces 1000 train + 200 eval examples across 4 task types plus >=450 SFT
warm-start traces for TRL SFTTrainer (chat format).

All randomness is routed through random.Random(seed) instances keyed by
data/seeds.json — re-running this script with the same seeds file MUST
produce byte-identical JSONL outputs. Enforced by tests/test_dataset.py.
"""
from __future__ import annotations

import json
import logging
import os
import random
from pathlib import Path
from typing import Any

log = logging.getLogger("fathom.data")

TASK_TYPES = ("niah", "multi_needle", "extractive", "counting")
DEFAULT_MIX = {"niah": 0.4, "multi_needle": 0.3, "extractive": 0.2, "counting": 0.1}
CONTEXT_LENGTHS = (4096, 16384, 65536, 204800)
NEEDLE_POSITIONS = ("start", "middle", "end")

# Filler corpus: 30 neutral sentences, no proper nouns that could collide with answers
_FILLER = [
    "The committee reviewed the proposed amendments to the existing policy framework.",
    "Several participants noted that further clarification would be necessary.",
    "The quarterly report indicated a steady increase in operational efficiency.",
    "Researchers observed significant variability across the sample population.",
    "The maintenance schedule was updated to reflect recent infrastructure changes.",
    "All participants were required to complete the mandatory orientation session.",
    "The distribution of resources followed a predetermined allocation protocol.",
    "Field observations confirmed the accuracy of the theoretical predictions.",
    "The project timeline was adjusted to accommodate unexpected technical delays.",
    "Compliance with the updated regulations required comprehensive staff training.",
    "The evaluation criteria were established prior to the commencement of testing.",
    "Multiple iterations of the process were necessary to achieve the desired outcome.",
    "The inventory management system was integrated with the existing database.",
    "Periodic assessments were conducted to monitor progress toward stated objectives.",
    "The documentation requirements were clarified during the preliminary review phase.",
    "Stakeholder feedback was incorporated into the revised implementation strategy.",
    "The assessment framework distinguished between formative and summative measures.",
    "Resource allocation decisions were guided by priority rankings established earlier.",
    "The calibration procedure ensured consistency across all measurement instruments.",
    "Preliminary findings suggested that the intervention produced measurable effects.",
    "The oversight committee convened on a monthly basis to review operational metrics.",
    "Participants were divided into cohorts based on predetermined selection criteria.",
    "The verification process involved cross-referencing multiple independent sources.",
    "An audit of the existing procedures identified several areas for improvement.",
    "The configuration parameters were adjusted to optimize system performance.",
    "Baseline measurements were recorded prior to the introduction of any changes.",
    "The scheduling algorithm prioritized tasks based on urgency and available capacity.",
    "A comparative analysis revealed differences between the two methodological approaches.",
    "The deployment process followed a staged rollout to minimize disruption.",
    "All submitted materials were reviewed according to established evaluation rubrics.",
]

_ADJECTIVES = [
    "azure", "crimson", "emerald", "golden", "ivory", "jade", "lavender",
    "magenta", "onyx", "pearl", "ruby", "sapphire", "scarlet", "silver", "teal",
    "violet", "amber", "cobalt", "coral", "indigo",
]
_NOUNS = [
    "vase", "lamp", "clock", "mirror", "chair", "table", "shelf", "frame",
    "carpet", "curtain", "statue", "pillar", "cabinet", "drawer", "bench",
    "chest", "vessel", "column", "panel", "gate",
]
_ITEMS = ["apple", "banana", "cherry", "mango", "peach", "plum", "grape", "lemon"]


def _build_filler(rng: random.Random, target_chars: int) -> str:
    """Tile filler sentences until >= target_chars characters."""
    sentences = list(_FILLER)
    rng.shuffle(sentences)
    result = []
    total = 0
    while total < target_chars:
        for s in sentences:
            result.append(s)
            total += len(s) + 1
            if total >= target_chars:
                break
    return " ".join(result)


def _assert_no_leak(gold_answer: str, context_without_needle: str) -> None:
    """DATA-04 post-check: gold_answer must not appear in filler (without the needle)."""
    if gold_answer.lower().strip() in context_without_needle.lower():
        raise ValueError(
            f"DATA-04 post-check: gold_answer '{gold_answer}' appears verbatim in filler context"
        )


def _gen_niah(rng: random.Random, context_length: int, needle_position: str) -> dict:
    """Needle-in-haystack: single fact extraction."""
    adj = rng.choice(_ADJECTIVES)
    noun = rng.choice(_NOUNS)
    gold_answer = adj
    fact = f"The {noun} is {adj}."
    prompt = f"Question: What color is the {noun} mentioned in the document?"

    target_chars = context_length * 4  # ~4 chars per token estimate
    filler = _build_filler(rng, target_chars)
    words = filler.split()
    total = len(words)

    if needle_position == "start":
        insert_idx = 0
    elif needle_position == "end":
        insert_idx = max(0, total - 20)
    else:  # middle
        insert_idx = total // 2

    fact_words = fact.split()
    words = words[:insert_idx] + fact_words + words[insert_idx:]
    context = " ".join(words)

    _assert_no_leak(gold_answer, context.replace(fact, ""))
    return {"prompt": prompt, "context": context, "gold_answer": gold_answer}


def _gen_multi_needle(rng: random.Random, context_length: int, needle_position: str) -> dict:
    """Multi-needle: sum of 3 integer facts."""
    items = rng.sample(_ITEMS, 3)
    values = [rng.randint(10, 99) for _ in range(3)]
    gold_answer = str(sum(values))
    prompt = f"Question: What is the total cost of {items[0]}, {items[1]}, and {items[2]}?"

    target_chars = context_length * 4
    filler = _build_filler(rng, target_chars)
    words = filler.split()
    total = len(words)

    # Insert 3 facts at distributed positions
    facts = [f"The {items[i]} costs {values[i]}." for i in range(3)]
    positions = [total // 4, total // 2, 3 * total // 4]

    offset = 0
    for i, (fact, pos) in enumerate(zip(facts, positions)):
        insert_at = pos + offset
        fw = fact.split()
        words = words[:insert_at] + fw + words[insert_at:]
        offset += len(fw)

    context = " ".join(words)
    _assert_no_leak(gold_answer, context)
    return {"prompt": prompt, "context": context, "gold_answer": gold_answer}


def _gen_extractive(rng: random.Random, context_length: int, needle_position: str) -> dict:
    """Extractive QA: short-span exact match."""
    years = [str(y) for y in range(1950, 2010)]
    cities = ["Rome", "Vienna", "Geneva", "Brussels", "Lisbon", "Madrid", "Athens",
              "Helsinki", "Stockholm", "Warsaw", "Prague", "Budapest", "Zurich"]
    year = rng.choice(years)
    city = rng.choice(cities)
    gold_answer = city
    fact = f"The {year} agreement was signed in {city}."
    prompt = f"Question: In which city was the {year} agreement signed?"

    target_chars = context_length * 4
    filler = _build_filler(rng, target_chars)
    words = filler.split()
    total = len(words)

    if needle_position == "start":
        insert_idx = 0
    elif needle_position == "end":
        insert_idx = max(0, total - 20)
    else:
        insert_idx = total // 2

    fact_words = fact.split()
    words = words[:insert_idx] + fact_words + words[insert_idx:]
    context = " ".join(words)

    _assert_no_leak(gold_answer, context.replace(fact, ""))
    return {"prompt": prompt, "context": context, "gold_answer": gold_answer}


def _gen_counting(rng: random.Random, context_length: int, needle_position: str) -> dict:
    """Counting: count occurrences of a target word."""
    target_word = rng.choice(_ITEMS)
    count = rng.randint(5, 20)
    gold_answer = str(count)
    prompt = f"Question: How many times does '{target_word}' appear in the document?"

    target_chars = context_length * 4
    filler_words = _build_filler(rng, target_chars).split()
    # Filter out any accidental occurrences of target_word in filler
    filler_words = [w for w in filler_words if w.lower().strip(".,") != target_word]

    # Insert target_word at evenly-spaced positions
    step = max(1, len(filler_words) // (count + 1))
    words = list(filler_words)
    for i in range(count):
        insert_at = min((i + 1) * step, len(words))
        words.insert(insert_at, target_word)

    context = " ".join(words)
    _assert_no_leak(gold_answer, context)
    return {"prompt": prompt, "context": context, "gold_answer": gold_answer}


_GEN_FN = {
    "niah": _gen_niah,
    "multi_needle": _gen_multi_needle,
    "extractive": _gen_extractive,
    "counting": _gen_counting,
}


def _compute_difficulty(context_length: int, needle_position: str, task_type: str) -> str:
    """Deterministic difficulty tier from example attributes."""
    if context_length == 4096 and needle_position == "start" and task_type in ("niah", "extractive"):
        return "trivial"
    elif context_length in (4096, 16384) and needle_position in ("start", "middle"):
        return "easy"
    elif context_length == 65536 or task_type == "multi_needle":
        return "medium"
    else:
        return "hard"


def _write_jsonl(path: Path, rows: list[dict]) -> None:
    """Write JSONL with sorted keys and compact separators for byte-determinism."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, sort_keys=True, separators=(",", ":")) + "\n")


def _pick_task_type(mix: dict, counts: dict) -> str:
    """Pick the task type with the largest gap from target proportions."""
    total = sum(counts.values()) + 1
    best = max(
        mix.keys(),
        key=lambda t: mix[t] - counts.get(t, 0) / total,
    )
    return best


# RLM system prompt for SFT traces (DATA-06)
_RLM_SYSTEM_PROMPT = (
    "You are FATHOM, a recursive language model with a Python REPL sandbox. "
    "You can read a long document via the variable `ctx` and call `llm(prompt, chunk)` "
    "for sub-queries. Think step by step. Emit your final answer inside <answer>...</answer>."
)


def _template_sft_trace(rng: random.Random, example: dict) -> dict:
    """Build a template grep-then-answer SFT trace (deterministic)."""
    task_type = example.get("task_type", "niah")
    gold = example["gold_answer"]
    prompt = example["prompt"]
    ctx_preview = example["context"][:2000]

    if task_type == "counting":
        target = gold  # the count
        # Infer target word from prompt
        import re
        m = re.search(r"'([^']+)'", prompt)
        target_word = m.group(1) if m else "item"
        code = f'count = ctx.count("{target_word}")\nprint(count)'
        tool_output = str(gold)
    elif task_type == "multi_needle":
        code = (
            "import re\n"
            "matches = re.findall(r'costs (\\d+)', ctx)\n"
            "print(sum(int(x) for x in matches))"
        )
        tool_output = str(gold)
    else:
        code = (
            "import re\n"
            "matches = re.findall(r'(?:is|was signed in) ([\\w]+)', ctx[:8192])\n"
            "print(matches[0] if matches else 'not found')"
        )
        tool_output = str(gold)

    messages = [
        {"content": _RLM_SYSTEM_PROMPT, "role": "system"},
        {"content": f"{prompt}\n\n[Document excerpt]:\n{ctx_preview}", "role": "user"},
        {
            "content": f"I'll search the document programmatically.\n```python\n{code}\n```",
            "role": "assistant",
        },
        {"content": tool_output, "role": "tool"},
        {"content": f"Based on the search results, the answer is <answer>{gold}</answer>", "role": "assistant"},
    ]
    return {"messages": messages, "task_id": f"template-sft-{example['task_id']}"}


def _haiku_sft_trace(client: Any, example: dict) -> dict | None:
    """Call Claude Haiku to generate an SFT trace. Returns None on error."""
    try:
        ctx_preview = example["context"][:3000]
        user_msg = f"{example['prompt']}\n\n[Document excerpt]:\n{ctx_preview}"
        resp = client.messages.create(
            model="claude-haiku-4-5",
            max_tokens=1024,
            system=_RLM_SYSTEM_PROMPT,
            messages=[{"role": "user", "content": user_msg}],
        )
        assistant_text = resp.content[0].text
        # Ensure answer tag present
        if "<answer>" not in assistant_text:
            assistant_text += f"\n<answer>{example['gold_answer']}</answer>"
        messages = [
            {"content": _RLM_SYSTEM_PROMPT, "role": "system"},
            {"content": user_msg, "role": "user"},
            {"content": assistant_text, "role": "assistant"},
        ]
        return {"messages": messages, "task_id": f"haiku-sft-{example['task_id']}"}
    except Exception as e:
        log.warning("Haiku API error on seed %s: %s; falling back to template", example.get("seed"), e)
        return None


def generate_sft_traces(
    seed_list: list,
    train_rows: list,
    target_count: int = 500,
    api_key: str | None = None,
) -> list[dict]:
    """Generate SFT traces — Claude Haiku where possible, template fallback. DATA-06."""
    budget = int(os.environ.get("FATHOM_HAIKU_BUDGET", "200")) if api_key else 0
    client = None
    if api_key:
        try:
            import anthropic  # type: ignore
            client = anthropic.Anthropic(api_key=api_key)
        except Exception as e:
            log.warning("anthropic SDK import failed: %s; template-only", e)
            client = None

    # Use trivial/easy rows as basis for traces
    source_rows = [r for r in train_rows if r.get("difficulty") in ("trivial", "easy")]
    if not source_rows:
        source_rows = train_rows

    traces = []
    for i, seed in enumerate(seed_list[:target_count]):
        rng = random.Random(seed)
        example = source_rows[i % len(source_rows)]
        trace = None
        if client is not None and i < budget:
            trace = _haiku_sft_trace(client, example)
        if trace is None:
            trace = _template_sft_trace(rng, example)
        traces.append(trace)

    assert len(traces) >= 450, f"DATA-06 floor: got {len(traces)} traces, need >=450"
    return traces


def generate_all(
    out_dir: str | Path = "data",
    seeds_path: str | Path = "data/seeds.json",
    train_count: int = 1000,
    eval_count: int = 200,
    sft_target_count: int = 500,
    mix: dict | None = None,
) -> dict:
    """Deterministic end-to-end generator. Writes train.jsonl + eval.jsonl + sft_traces.jsonl.

    DATA-01..06 — all randomness routed through seeded RNGs.
    """
    out_dir = Path(out_dir)
    seeds_path = Path(seeds_path)
    mix = mix or DEFAULT_MIX

    with open(seeds_path, "r", encoding="utf-8") as f:
        seeds = json.load(f)

    # Force trivial floor: first 6% of train are trivial (DATA-04)
    trivial_floor = max(60, int(0.06 * train_count))

    def _build_split(seed_list: list, count: int, split: str) -> list[dict]:
        rows = []
        task_counts: dict[str, int] = {t: 0 for t in TASK_TYPES}

        for idx, seed in enumerate(seed_list[:count]):
            rng = random.Random(seed)

            # Force trivial tier for first N examples (DATA-04)
            if split == "train" and idx < trivial_floor:
                task_type = "niah" if idx % 2 == 0 else "extractive"
                context_length = 4096
                needle_position = "start"
            else:
                task_type = _pick_task_type(mix, task_counts)
                context_length = rng.choice(CONTEXT_LENGTHS)
                needle_position = rng.choice(NEEDLE_POSITIONS)

            task_counts[task_type] = task_counts.get(task_type, 0) + 1

            gen_fn = _GEN_FN[task_type]
            try:
                ex = gen_fn(rng, context_length, needle_position)
            except Exception as e:
                log.warning("Skipping example %d due to generation error: %s", idx, e)
                # Retry with a simpler config
                ex = _gen_niah(rng, 4096, "start")
                task_type = "niah"
                context_length = 4096
                needle_position = "start"

            difficulty = _compute_difficulty(context_length, needle_position, task_type)
            row = {
                "context": ex["context"],
                "context_length": context_length,
                "difficulty": difficulty,
                "gold_answer": ex["gold_answer"],
                "needle_position": needle_position,
                "prompt": ex["prompt"],
                "seed": seed,
                "task_id": f"{task_type}-{split}-{idx:04d}",
                "task_type": task_type,
            }
            rows.append(row)
        return rows

    log.info("DATA generating train split (%d examples)...", train_count)
    train_rows = _build_split(seeds["train"], train_count, "train")
    log.info("DATA generating eval split (%d examples)...", eval_count)
    eval_rows = _build_split(seeds["eval"], eval_count, "eval")

    # Self-checks (DATA-02, DATA-04, DATA-05)
    assert len(train_rows) == train_count, f"Expected {train_count} train rows, got {len(train_rows)}"
    assert len(eval_rows) == eval_count, f"Expected {eval_count} eval rows, got {len(eval_rows)}"

    train_ids = {r["task_id"] for r in train_rows}
    eval_ids = {r["task_id"] for r in eval_rows}
    assert not (train_ids & eval_ids), "Train/eval task_id overlap detected (DATA-02)"

    trivial_share = sum(1 for r in train_rows if r["difficulty"] == "trivial") / len(train_rows)
    assert trivial_share >= 0.05, f"Trivial share {trivial_share:.3f} < 0.05 (DATA-04)"

    _write_jsonl(out_dir / "train.jsonl", train_rows)
    _write_jsonl(out_dir / "eval.jsonl", eval_rows)
    log.info(
        "DATA train=%d eval=%d trivial_share=%.3f written",
        len(train_rows), len(eval_rows), trivial_share,
    )

    # SFT traces (DATA-06)
    log.info("DATA generating SFT traces (template-only unless ANTHROPIC_API_KEY set)...")
    sft_traces = generate_sft_traces(
        seeds["sft"],
        train_rows,
        target_count=sft_target_count,
        api_key=os.environ.get("ANTHROPIC_API_KEY"),
    )
    _write_jsonl(out_dir / "sft_traces.jsonl", sft_traces)
    log.info("DATA sft_traces=%d written", len(sft_traces))

    return {"train": len(train_rows), "eval": len(eval_rows), "sft": len(sft_traces)}


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
    result = generate_all()
    print(f"Generated: train={result['train']} eval={result['eval']} sft={result['sft']}")