#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""On-the-fly numeric distractor rows: ~50 web_search snippet templates × random slots.

Post-tool Hermes shape (call stripped by tool_context_scenario_to_hermes_row).
Gold answer contains only the gold number; snippet always has gold + distractor + fillers.
"""

from __future__ import annotations

import random

from hermes_agent_eval import SEARCH_TOOL, assistant_tool_call, tool_response

_CITIES = [
    ("Cairo", "Egypt"),
    ("Lagos", "Nigeria"),
    ("Seoul", "South Korea"),
    ("Lima", "Peru"),
    ("Oslo", "Norway"),
    ("Lisbon", "Portugal"),
    ("Nairobi", "Kenya"),
    ("Santiago", "Chile"),
    ("Hanoi", "Vietnam"),
    ("Riga", "Latvia"),
    ("Accra", "Ghana"),
    ("Quito", "Ecuador"),
    ("Muscat", "Oman"),
    ("Dhaka", "Bangladesh"),
    ("Jakarta", "Indonesia"),
]

_FX = [
    ("USD", "EUR", "US dollar", "euro", "euros per dollar"),
    ("GBP", "USD", "British pound", "US dollar", "dollars per pound"),
    ("AUD", "USD", "Australian dollar", "US dollar", "dollars per AUD"),
    ("CAD", "USD", "Canadian dollar", "US dollar", "dollars per CAD"),
    ("JPY", "USD", "Japanese yen", "US dollar", "dollars per yen"),
]

_STOCKS = ["ACME", "NOVA", "PEAK", "RIVER", "SOLAR", "ATLAS", "ORBIT"]
_PATHS = [
    "deploy/staging.yaml",
    "config/api.toml",
    "config/client.toml",
    "infra/service.yaml",
    "ops/limits.env",
]


def _fmt_int(n: int) -> str:
    return str(n)


def _fmt_dec(x: float, places: int = 2) -> str:
    return f"{x:.{places}f}"


def _distinct(rng: random.Random, gold: str, make_dist, tries: int = 12) -> str:
    for _ in range(tries):
        d = make_dist()
        if d != gold:
            return d
    return make_dist() + "1" if not make_dist().endswith("1") else make_dist() + "2"


# Each template: builds (question, query, title, snippet, think, answer) via rng.
# ponytail: 50+ surfaces — upgrade = scrape real SERP morphologies later.


def _t_pop_outdated(rng):
    city, country = rng.choice(_CITIES)
    gold_n = rng.randint(18, 180)
    gold = _fmt_int(gold_n)
    dist = _distinct(rng, gold, lambda: _fmt_int(max(1, gold_n // rng.choice([8, 10, 12]))))
    f1, f2 = rng.randint(80, 4000), rng.randint(100, 5000)
    return (
        f"What is the population of {city}?",
        f"population of {city}",
        f"{city} population",
        f"{city} in {country} has a population of about {gold} million people, "
        f"{f1} barber shops and {f2} grocery stores. An outdated blog cites {dist} million.",
        f"Official figure is {gold} million, not the outdated {dist} million.",
        f"{city}'s population is about {gold} million.",
    )


def _t_pop_estimate(rng):
    city, country = rng.choice(_CITIES)
    gold_n = rng.randint(5, 90)
    gold = _fmt_int(gold_n)
    dist = _distinct(rng, gold, lambda: _fmt_int(gold_n * rng.choice([8, 10, 12])))
    f1 = rng.randint(20, 900)
    return (
        f"How many people live in {city}?",
        f"{city} population",
        f"{city} demographics",
        f"Travel notes mention {dist} million residents in {city}, {country}; "
        f"census data reports about {gold} million people and {f1} public schools.",
        f"Census says {gold} million; the travel note {dist} million is wrong.",
        f"About {gold} million people live in {city}.",
    )


def _t_pop_table(rng):
    city, country = rng.choice(_CITIES)
    gold_n = rng.randint(12, 120)
    gold = _fmt_int(gold_n)
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(1, 9)))
    parks = rng.randint(15, 400)
    return (
        f"What is {city}'s population?",
        f"population {city} {country}",
        f"{city} facts",
        f"Quick facts for {city} ({country}): parks={parks}; viral meme population={dist} million; "
        f"statistics office population about {gold} million.",
        f"Statistics office value is {gold} million.",
        f"The population of {city} is about {gold} million.",
    )


def _t_fx_forum(rng):
    a, b, a_name, b_name, unit = rng.choice(_FX)
    gold_f = round(rng.uniform(0.45, 1.55), 2)
    gold = _fmt_dec(gold_f)
    dist = _distinct(rng, gold, lambda: _fmt_dec(round(gold_f * rng.choice([0.3, 0.4, 0.5]), 2)))
    return (
        f"What is the current {a} to {b} exchange rate?",
        f"{a} {b} exchange rate",
        f"{a}{b} rate",
        f"A forum post claims {dist} {unit}; Reuters cites about {gold} {unit} for the {a_name}.",
        f"Credible rate is {gold}, not {dist}.",
        f"The {a} to {b} rate is about {gold}.",
    )


def _t_fx_old_table(rng):
    a, b, a_name, b_name, unit = rng.choice(_FX)
    gold_f = round(rng.uniform(0.5, 1.4), 2)
    gold = _fmt_dec(gold_f)
    dist = _distinct(rng, gold, lambda: _fmt_dec(round(rng.uniform(0.15, 0.35), 2)))
    return (
        f"What is the {a}/{b} rate right now?",
        f"{a} to {b} rate today",
        f"{a} {b}",
        f"An old spreadsheet shows {dist} {unit}; the live rate is about {gold} {unit}.",
        f"Live rate is {gold}.",
        f"The {a} to {b} rate is about {gold}.",
    )


def _t_stock_spam(rng):
    sym = rng.choice(_STOCKS)
    gold_n = rng.randint(40, 480)
    gold = _fmt_int(gold_n)
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(4, 40)))
    vol = rng.randint(100, 9000)
    return (
        f"What is {sym} trading at in US dollars?",
        f"{sym} stock price USD",
        f"{sym} quote",
        f"Spam ads mention {dist} USD; major exchanges show about {gold} USD for {sym}, volume {vol}k.",
        f"Exchange price is {gold} USD.",
        f"{sym} trades near {gold} US dollars.",
    )


def _t_replicas_notes(rng):
    path = rng.choice(_PATHS)
    gold = _fmt_int(rng.choice([2, 3, 4, 5, 6, 8]))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.choice([1, 7, 9, 12])))
    image = f"v{rng.randint(1, 3)}.{rng.randint(0, 9)}.{rng.randint(0, 9)}"
    return (
        f"How many replicas does {path} specify?",
        f"replicas in {path}",
        f"{path} deploy",
        f"Dev notes mention {dist} replica for local; {path} specifies replicas: {gold} and image: {image}.",
        f"Config value is {gold} replicas.",
        f"The replica count in {path} is {gold}.",
    )


def _t_port_file(rng):
    path = rng.choice(_PATHS)
    gold = _fmt_int(rng.randint(2000, 9999))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(80, 900)))
    workers = rng.randint(2, 16)
    return (
        f"What port does the service listen on in {path}?",
        f"listen port {path}",
        f"{path} network",
        f"Comments say legacy port {dist}; {path} has listen_port = {gold} and workers = {workers}.",
        f"Configured listen_port is {gold}.",
        f"The service listens on port {gold}.",
    )


def _t_timeout_ms(rng):
    path = rng.choice(_PATHS)
    gold = _fmt_int(rng.choice([1500, 2000, 3000, 4500, 5000, 8000, 10000]))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.choice([100, 250, 500, 999])))
    retries = rng.randint(1, 5)
    return (
        f"What is timeout_ms in {path}?",
        f"timeout_ms {path}",
        f"{path} client",
        f"README example uses {dist}; {path} sets timeout_ms = {gold} and retries = {retries}.",
        f"File value is {gold}.",
        f"timeout_ms in {path} is {gold}.",
    )


def _t_distance(rng):
    a, _ = rng.choice(_CITIES)
    b, _ = rng.choice(_CITIES)
    if a == b:
        b = "Berlin"
    gold = _fmt_int(rng.randint(200, 9000))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(20, 180)))
    flights = rng.randint(5, 80)
    return (
        f"How far is {a} from {b} in kilometers?",
        f"distance {a} to {b} km",
        f"{a}-{b} distance",
        f"A tweet says {dist} km; airline tables list about {gold} km between {a} and {b}, with {flights} daily flights.",
        f"Table distance is {gold} km.",
        f"{a} is about {gold} kilometers from {b}.",
    )


def _t_temp_c(rng):
    city, country = rng.choice(_CITIES)
    gold = _fmt_int(rng.randint(-5, 38))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(40, 55)))
    humid = rng.randint(30, 95)
    return (
        f"What is the high temperature in {city} today in °C?",
        f"weather high {city} celsius",
        f"{city} weather",
        f"{city}, {country}: forecast high {gold}°C, humidity {humid}%. A wrong mirror site shows {dist}°C.",
        f"Forecast high is {gold}°C.",
        f"The high in {city} today is about {gold}°C.",
    )


def _t_wifi_gen(rng):
    # gold is generation number as digit-free name mostly — use channel count style
    gold = _fmt_int(rng.choice([6, 7]))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.choice([4, 5, 8])))
    year = rng.choice([2019, 2021, 2024])
    return (
        "What is the latest Wi-Fi generation number mentioned?",
        "latest Wi-Fi generation number",
        "Wi-Fi standard",
        f"Marketing copy still pushes Wi-Fi {dist}; industry notes say Wi-Fi {gold} (around {year}) is the recent generation, "
        f"with {rng.randint(160, 320)} MHz channels.",
        f"Recent generation is Wi-Fi {gold}.",
        f"Wi-Fi {gold} is the recent Wi-Fi generation mentioned.",
    )


def _t_http_ver(rng):
    gold = "3"
    dist = rng.choice(["1", "2", "1.1"])
    quic_ports = rng.randint(1000, 9000)
    return (
        "What HTTP version number runs over QUIC according to the snippet?",
        "HTTP version over QUIC",
        "HTTP/3 QUIC",
        f"Legacy docs emphasize HTTP/{dist}; the snippet states HTTP/{gold} runs over QUIC on UDP, example port {quic_ports}.",
        f"The QUIC HTTP version is {gold}.",
        f"HTTP/{gold} is the version that runs over QUIC.",
    )


def _t_workers(rng):
    path = rng.choice(_PATHS)
    gold = _fmt_int(rng.randint(2, 64))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(1, 8)))
    return (
        f"How many workers are configured in {path}?",
        f"workers {path}",
        f"{path} workers",
        f"Old helm chart used workers = {dist}; {path} now has workers = {gold} and queue = {rng.randint(10, 200)}.",
        f"Configured workers={gold}.",
        f"There are {gold} workers in {path}.",
    )


def _t_battery(rng):
    gold = _fmt_int(rng.randint(50, 99))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(10, 40)))
    cycles = rng.randint(100, 900)
    return (
        "What battery percentage does the status report show?",
        "device battery percent status",
        "battery status",
        f"Debug dump lists cycles={cycles}, a stale cache shows {dist}%, live status battery = {gold}%.",
        f"Live battery is {gold}%.",
        f"The battery is at {gold} percent.",
    )


def _t_users_online(rng):
    gold = _fmt_int(rng.randint(200, 20000))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(10, 150)))
    rooms = rng.randint(5, 80)
    return (
        "How many users are online according to the status page?",
        "users online status page",
        "online users",
        f"Homepage badge wrongly shows {dist} online; status API reports {gold} users online across {rooms} rooms.",
        f"Status API says {gold} users.",
        f"About {gold} users are online.",
    )


def _t_build_number(rng):
    gold = _fmt_int(rng.randint(1000, 9999))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(100, 999)))
    return (
        "What is the build number in the release notes snippet?",
        "release build number",
        "build number",
        f"CI artifact name includes build-{dist}; release notes say build number {gold} shipped with {rng.randint(2, 40)} fixes.",
        f"Release build number is {gold}.",
        f"The build number is {gold}.",
    )


def _t_area_km2(rng):
    city, country = rng.choice(_CITIES)
    gold = _fmt_int(rng.randint(50, 2500))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(5, 40)))
    return (
        f"What is the area of {city} in square kilometers?",
        f"{city} area km2",
        f"{city} area",
        f"{city}, {country}: wiki infobox area {gold} km²; a quiz site lists {dist} km² and {rng.randint(2, 20)} districts.",
        f"Infobox area is {gold} km².",
        f"{city} covers about {gold} square kilometers.",
    )


def _t_year_founded_distract(rng):
    # question asks for founding year; distractor is employee count
    org = rng.choice(["Northwind", "Contoso", "Fabrikam", "Tailspin"])
    gold = _fmt_int(rng.randint(1985, 2018))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(50, 9000)))  # employees
    return (
        f"In what year was {org} founded?",
        f"{org} founded year",
        f"{org} history",
        f"{org} employs about {dist} people and was founded in {gold}; HQ has {rng.randint(2, 12)} floors.",
        f"Founded year is {gold}.",
        f"{org} was founded in {gold}.",
    )


def _t_speed_limit(rng):
    gold = _fmt_int(rng.choice([30, 50, 70, 90, 100, 120]))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.choice([20, 40, 60, 80, 110])))
    lanes = rng.randint(1, 4)
    return (
        "What is the speed limit in km/h on the described road?",
        "road speed limit km/h",
        "speed limit",
        f"Tour blog claims {dist} km/h; official signage speed limit {gold} km/h on a {lanes}-lane section.",
        f"Official limit is {gold} km/h.",
        f"The speed limit is {gold} km/h.",
    )


def _t_memory_gb(rng):
    gold = _fmt_int(rng.choice([8, 16, 32, 64, 128]))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.choice([4, 12, 24, 48])))
    cores = rng.randint(4, 64)
    return (
        "How much RAM in GB does the server inventory list?",
        "server RAM GB inventory",
        "server memory",
        f"Old ticket said {dist} GB; inventory lists RAM = {gold} GB and CPU cores = {cores}.",
        f"Inventory RAM is {gold} GB.",
        f"The server has {gold} GB of RAM.",
    )


def _t_invoice_total(rng):
    gold = _fmt_int(rng.randint(120, 9000))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(10, 99)))
    tax = rng.randint(5, 25)
    return (
        "What is the invoice total in dollars?",
        "invoice total dollars",
        "invoice",
        f"Line preview shows tip {dist}; invoice total ${gold} including {tax}% tax.",
        f"Invoice total is {gold}.",
        f"The invoice total is {gold} dollars.",
    )


def _t_latency_ms(rng):
    gold = _fmt_int(rng.randint(20, 400))
    dist = _distinct(rng, gold, lambda: _fmt_int(rng.randint(800, 2000)))
    qps = rng.randint(50, 5000)
    return (
        "What p99 latency in milliseconds is reported?",
        "p99 latency ms",
        "latency",
        f"Alert threshold is {dist} ms; dashboard p99 latency = {gold} ms at {qps} QPS.",
        f"Reported p99 is {gold} ms.",
        f"p99 latency is {gold} milliseconds.",
    )


def _t_population_country(rng):
    # country-level, not city
    countries = [
        ("Egypt", 87),
        ("Peru", 34),
        ("Chile", 19),
        ("Kenya", 55),
        ("Portugal", 10),
        ("Vietnam", 98),
        ("Ghana", 33),
    ]
    country, base = rng.choice(countries)
    gold_n = base + rng.randint(-2, 3)
    gold = _fmt_int(max(5, gold_n))
    dist = _distinct(rng, gold, lambda: _fmt_int(max(1, gold_n // 10)))
    cities_n = rng.randint(20, 200)
    return (
        f"What is the population of {country}?",
        f"population of {country}",
        f"{country} population",
        f"{country} has about {gold} million people and roughly {cities_n} municipalities. "
        f"A mislabeled chart shows {dist} million.",
        f"Population is about {gold} million.",
        f"{country}'s population is about {gold} million.",
    )


# Build ≥50 callables by repeating families with slight wrappers
_TEMPLATE_FNS = []


def _register(fn, n=1):
    for _ in range(n):
        _TEMPLATE_FNS.append(fn)


_register(_t_pop_outdated, 4)
_register(_t_pop_estimate, 4)
_register(_t_pop_table, 3)
_register(_t_population_country, 3)
_register(_t_fx_forum, 4)
_register(_t_fx_old_table, 3)
_register(_t_stock_spam, 3)
_register(_t_replicas_notes, 3)
_register(_t_port_file, 3)
_register(_t_timeout_ms, 3)
_register(_t_distance, 3)
_register(_t_temp_c, 2)
_register(_t_wifi_gen, 2)
_register(_t_http_ver, 2)
_register(_t_workers, 2)
_register(_t_battery, 2)
_register(_t_users_online, 2)
_register(_t_build_number, 2)
_register(_t_area_km2, 2)
_register(_t_year_founded_distract, 2)
_register(_t_speed_limit, 2)
_register(_t_memory_gb, 2)
_register(_t_invoice_total, 2)
_register(_t_latency_ms, 2)


def make_distractor_scenario(rng: random.Random | None = None) -> dict:
    rng = rng or random.Random()
    fn = rng.choice(_TEMPLATE_FNS)
    question, query, title, snippet, think, answer = fn(rng)
    return {
        "id": f"synth_num_{fn.__name__}",
        "category": "grounded_answer",
        "expect": "Answer with gold number from noisy search snippet.",
        "tools": [SEARCH_TOOL],
        "turns": [
            {"from": "human", "value": question},
            {"from": "gpt", "value": assistant_tool_call("web_search", {"query": query})},
            {
                "from": "tool",
                "value": tool_response(
                    "web_search:0",
                    "web_search",
                    {"results": [{"title": title, "snippet": snippet}]},
                ),
            },
        ],
        "gold": f"<think>\n{think}\n</think>\n{answer}",
    }


def iter_numeric_distractor_synth(seed: int = 0):
    """Infinite Hermes-ready scenarios (convert with tool_context_scenario_to_hermes_row)."""
    rng = random.Random(seed)
    i = 0
    while True:
        # re-seed lightly for diversity across workers
        local = random.Random(rng.randint(0, 2**30) ^ i)
        yield make_distractor_scenario(local)
        i += 1


def self_check():
    assert len(_TEMPLATE_FNS) >= 50, len(_TEMPLATE_FNS)
    from sft_data import tool_context_scenario_to_hermes_row
    from embed_answer_score import extract_numbers, gold_answer_text, tool_content_text_from_tool_turns

    rng = random.Random(0)
    for _ in range(30):
        s = make_distractor_scenario(rng)
        row = tool_context_scenario_to_hermes_row(s)
        assert row.get("conversations"), s
        gold = gold_answer_text(s["gold"])
        tool_raw = next(t["value"] for t in s["turns"] if t["from"] == "tool")
        tool = tool_content_text_from_tool_turns([tool_raw])
        gn = extract_numbers(gold)
        tn = extract_numbers(tool)
        assert gn, gold
        assert all(n in tn for n in gn), (gn, tn, tool, s["gold"])
    print(f"numeric_distractor_synth self-check ok templates={len(_TEMPLATE_FNS)}")


if __name__ == "__main__":
    self_check()