"""Controlled private-fact task generator (deterministic, seeded).

Generates fictional 'project ledger' documents whose facts cannot be inferred from pretraining: invented entity names,
random identifiers, amounts, dates, units, status events (approved/rejected, cancelled vs postponed), timestamped
revisions with an explicit precedence rule, distractor projects with similar names and plausible wrong numbers,
boilerplate filler to reach a target token length. Each document ships with (a) the text, (b) questions by type,
(c) an evaluator-only answer key and fact ledger. Keys are written to a separate directory never exposed to agents.
"""
from __future__ import annotations
import random, json, math, datetime as dt
from dataclasses import dataclass, field, asdict
from typing import Any

SYLL = ["var", "ko", "lith", "zen", "dra", "mel", "qui", "tor", "ban", "sel", "ur", "nak", "fel", "ost", "ri", "gam", "vol", "tis", "pra", "lun",
        "ker", "oby", "mun", "tal", "ix", "ard", "wen", "sol", "bri", "tek", "nov", "ald", "ymi", "ruk", "esh"]
FIRST = ["Teodora", "Milos", "Ingrid", "Rasheed", "Yevgenia", "Kwame", "Solveig", "Tarek", "Bohdan", "Nnamdi", "Lorcan", "Priyamvada", "Oksana", "Haruki",
         "Zéphyrine", "Casimir", "Ludivine", "Ewan", "Marisol", "Ragnvald", "Aurelio", "Dagny", "Idris", "Wilhelmina"]
LAST = ["Qvistrand", "Obuya", "Ferrante", "Mäkelä", "Dubrovnik", "Achterberg", "Nwachukwu", "Sørlie", "Baptiste-Roy", "Zelenko", "Haldorsen", "Petrakis",
        "Ouellet-Marchand", "Kirchhoff", "Abubakar", "Lindqvist", "Tremblay-Côté", "Varga", "Okonkwo", "Bergström", "Ishikawa", "Delacroix", "Mbeki", "Rousseau"]
TOWNS = ["Chibougamau", "Matagami", "Fermont", "Lebel-sur-Quévillon", "Sept-Îles", "Rouyn-Noranda", "Val-d'Or", "Schefferville", "Baie-Comeau", "Amos", "La Tuque",
         "Senneterre", "Gaspé", "Mont-Laurier", "Témiscaming"]
CURR = ["CAD", "EUR", "USD", "CHF", "NOK"]
UNITS = [("MW", "kW", 1000.0), ("km", "m", 1000.0), ("t", "kg", 1000.0), ("ha", "m²", 10000.0)]
FILLER = [
    "The committee reiterates that all figures are reported in nominal terms unless otherwise stated.",
    "Procurement follows the standard three-quote rule for purchases above the departmental threshold.",
    "Quarterly reviews are held on the second Tuesday of the month following the quarter end.",
    "Environmental monitoring reports are archived by the records office for a period of seven years.",
    "Travel reimbursements require original receipts and a signed itinerary.",
    "Site access badges must be renewed annually and returned on project completion.",
    "The internal audit function reports to the governance board and not to project leads.",
    "Vendor onboarding includes a conflict-of-interest declaration and a bank verification letter.",
    "All timestamps in this ledger are recorded in local time at the head office.",
    "Minutes are circulated within five business days and approved at the following meeting.",
    "Health and safety briefings are mandatory for all visitors before entering an active site.",
    "Capital expenditure is tracked separately from operating costs in the ledger system.",
    "Contingency reserves may only be released with written approval from the finance director.",
    "Printed copies of this memo are uncontrolled; the electronic record prevails.",
    "Project codes follow the pattern of a name, a hyphen and a sequence number.",
    "Where two entries conflict, the entry with the later timestamp supersedes the earlier one.",
]

def _name(rng: random.Random) -> str:
    n = rng.choice([2, 3])
    s = "".join(rng.choice(SYLL) for _ in range(n))
    return s.capitalize()

def _ident(rng: random.Random) -> str:
    L = "ABCDEFGHJKLMNPQRSTUVWXYZ"
    return f"{rng.choice(L)}{rng.choice(L)}-{rng.randint(10,99)}{rng.choice(L)}-{rng.randint(1000,9999)}"

def _lot(rng: random.Random) -> str:
    return f"Lot {rng.randint(10,999)}-{rng.choice('ABCDEFGHKLMNPRSTVWXZ')}"

def _date(rng: random.Random, y0=2029, y1=2034) -> dt.date:
    return dt.date(rng.randint(y0, y1), rng.randint(1, 12), rng.randint(1, 28))

def fmt_date(d: dt.date) -> str:
    return d.strftime("%-d %B %Y")

def fmt_amt(x: int) -> str:
    return f"{x:,}"

@dataclass
class Fact:
    fact_id: str
    entity: str
    attribute: str
    value: str
    kind: str          # identifier | amount | date | status | person | place | quantity
    sentence: str
    char_start: int = -1
    char_end: int = -1

@dataclass
class Question:
    question_id: str
    qtype: str
    question: str
    answer: str
    answer_type: str      # exact | numeric | yesno | abstain | choice
    tolerance: float
    supporting_fact_ids: list
    answer_position_frac: float   # relative position of the main supporting sentence in the document (0..1)
    aliases: list = field(default_factory=list)

@dataclass
class Project:
    name: str
    code: str
    proposer: str
    proposed_date: dt.date
    proposed_amount: int
    currency: str
    site_lot: str
    town: str
    supplier_id: str
    approved: bool
    decision_date: dt.date
    revisions: list            # [(date, amount)]
    cancelled: bool
    postponed_to: dt.date | None
    quantity: tuple            # (value, unit_big, unit_small, factor)
    decimal_value: str         # unusual decimal, e.g. "17.0463"

def make_project(rng: random.Random, base: str | None = None, used: set | None = None) -> Project:
    used = used if used is not None else set()
    while True:
        nm = base or _name(rng)
        code = f"{nm}-{rng.randint(1, 19)}"
        if code not in used:
            used.add(code); break
    proposed = _date(rng)
    amt = rng.randint(180_000, 9_900_000)
    amt = amt - amt % rng.choice([1, 5, 25])
    dec = rng.randint(1, 3)
    revisions = []
    d = proposed
    for _ in range(rng.choice([0, 1, 2, 2])):
        d = d + dt.timedelta(days=rng.randint(7, 120))
        amt2 = int(amt * rng.uniform(0.82, 1.18)); amt2 -= amt2 % 5
        revisions.append((d, amt2))
    decision = d + dt.timedelta(days=rng.randint(3, 60))
    approved = rng.random() < 0.7
    cancelled = (rng.random() < 0.3)
    postponed = None if (cancelled or rng.random() < 0.4) else decision + dt.timedelta(days=rng.randint(30, 400))
    q = rng.choice(UNITS)
    qv = round(rng.uniform(0.4, 48.0), rng.choice([1, 2]))
    cur = "CAD" if rng.random() < 0.6 else rng.choice(CURR)
    return Project(nm, code, f"{rng.choice(FIRST)} {rng.choice(LAST)}", proposed, amt, cur, _lot(rng), rng.choice(TOWNS), _ident(rng),
                   approved, decision, revisions, cancelled, postponed, (qv, q[0], q[1], q[2]), f"{rng.uniform(1, 99):.4f}")

def current_amount(p: Project) -> int:
    return p.revisions[-1][1] if p.revisions else p.proposed_amount

def render_project(p: Project, rng: random.Random, facts: list, shuffle_sentences: bool) -> str:
    S = []
    def F(attr, val, kind, sent):
        facts.append(Fact(f"{p.code}:{attr}", p.code, attr, val, kind, sent)); S.append(sent)
    F("proposer", p.proposer, "person", f"Project {p.code} was proposed by {p.proposer} on {fmt_date(p.proposed_date)} with an initial budget of {fmt_amt(p.proposed_amount)} {p.currency}.")
    facts.append(Fact(f"{p.code}:proposed_date", p.code, "proposed_date", fmt_date(p.proposed_date), "date", S[-1]))
    facts.append(Fact(f"{p.code}:proposed_amount", p.code, "proposed_amount", f"{fmt_amt(p.proposed_amount)} {p.currency}", "amount", S[-1]))
    for i, (d, a) in enumerate(p.revisions):
        F(f"revision_{i+1}", f"{fmt_amt(a)} {p.currency}", "amount", f"On {fmt_date(d)} the budget of {p.code} was revised to {fmt_amt(a)} {p.currency}.")
    verb = "approved" if p.approved else "rejected"
    F("decision", verb, "status", f"The {verb} decision for {p.code} was recorded on {fmt_date(p.decision_date)}.")
    F("site", f"{p.site_lot} in {p.town}", "place", f"The site of {p.code} is {p.site_lot} in {p.town}.")
    F("supplier_id", p.supplier_id, "identifier", f"The supplier identifier for {p.code} is {p.supplier_id}.")
    if p.cancelled:
        F("cancelled", "yes", "status", f"{p.code} was cancelled; it was not postponed.")
    elif p.postponed_to is not None:
        F("cancelled", "no", "status", f"{p.code} was NOT cancelled; its launch was postponed to {fmt_date(p.postponed_to)}.")
        facts.append(Fact(f"{p.code}:postponed_to", p.code, "postponed_to", fmt_date(p.postponed_to), "date", S[-1]))
    else:
        F("cancelled", "no", "status", f"{p.code} remains on schedule and was neither cancelled nor postponed.")
    qv, ub, us, fac = p.quantity
    F("quantity", f"{qv} {ub}", "quantity", f"The rated capacity of {p.code} is {qv} {ub}.")
    F("decimal", p.decimal_value, "identifier", f"The calibration constant filed for {p.code} is {p.decimal_value}.")
    if shuffle_sentences:
        rng.shuffle(S)
    return " ".join(S)

def _norm_num(x: float) -> str:
    if abs(x - round(x)) < 1e-9:
        return str(int(round(x)))
    return f"{x:.6g}"

def build_questions(projects: list[Project], distractor_names: list[str], rng: random.Random, doc_text: str, facts: list[Fact]) -> list[Question]:
    Q = []
    pos = {}
    for f in facts:
        i = doc_text.find(f.sentence)
        f.char_start, f.char_end = i, i + len(f.sentence)
        pos[f.fact_id] = i / max(1, len(doc_text))
    def add(qtype, text, ans, atype, fids, tol=0.0, aliases=None):
        Q.append(Question(f"q{len(Q):02d}_{qtype}", qtype, text, ans, atype, tol, fids, (pos.get(fids[0], 0.0) if fids else -1.0), aliases or []))
    # choose target projects for each type (deterministic order by rng)
    P = list(projects); rng.shuffle(P)
    p = P[0]; add("single_fact", f"What is the supplier identifier for {p.code}?", p.supplier_id, "exact", [f"{p.code}:supplier_id"])
    p = P[1 % len(P)]; add("single_fact", f"Where is the site of {p.code}? Give the lot and the town.", f"{p.site_lot} in {p.town}", "exact", [f"{p.code}:site"], aliases=[p.site_lot, f"{p.site_lot}, {p.town}"])
    p = P[2 % len(P)]; add("exact_copy", f"What is the calibration constant filed for {p.code}? Copy it exactly.", p.decimal_value, "exact", [f"{p.code}:decimal"])
    p = P[3 % len(P)]; add("exact_copy", f"Who proposed {p.code}? Give the full name.", p.proposer, "exact", [f"{p.code}:proposer"])
    p = P[4 % len(P)]; add("multi_hop", f"What was the initial budget, in {p.currency}, of the project proposed by {p.proposer}?", fmt_amt(p.proposed_amount), "numeric", [f"{p.code}:proposer", f"{p.code}:proposed_amount"], 0.0)
    p = P[5 % len(P)]; add("multi_hop", f"What is the supplier identifier of the project whose site is {p.site_lot} in {p.town}?", p.supplier_id, "exact", [f"{p.code}:site", f"{p.code}:supplier_id"])
    # negation: cancelled vs postponed
    cand = [x for x in projects if not x.cancelled and x.postponed_to is not None] or projects
    p = rng.choice(cand); add("negation", f"Was {p.code} cancelled? Answer yes or no.", "yes" if p.cancelled else "no", "yesno", [f"{p.code}:cancelled"])
    cand = [x for x in projects if x.cancelled] or projects
    p = rng.choice(cand); add("negation", f"Was {p.code} postponed? Answer yes or no.", "yes" if (p.postponed_to is not None) else "no", "yesno", [f"{p.code}:cancelled"])
    # revision / precedence
    cand = [x for x in projects if len(x.revisions) >= 1] or projects
    p = rng.choice(cand); add("revision", f"According to the ledger's precedence rule, what is the current budget of {p.code}, in {p.currency}? Give the number only.", fmt_amt(current_amount(p)), "numeric", [f"{p.code}:revision_{len(p.revisions)}" if p.revisions else f"{p.code}:proposed_amount"], 0.0)
    cand = [x for x in projects if len(x.revisions) >= 2] or cand
    p = rng.choice(cand); add("revision", f"What was the first revised budget of {p.code}, in {p.currency}? Give the number only.", fmt_amt(p.revisions[0][1]) if p.revisions else fmt_amt(p.proposed_amount), "numeric", [f"{p.code}:revision_1" if p.revisions else f"{p.code}:proposed_amount"], 0.0)
    # temporal order
    a, b = P[0], P[1 % len(P)]
    first = a.code if a.decision_date < b.decision_date else b.code
    add("temporal", f"Which decision was recorded first: the one for {a.code} or the one for {b.code}? Answer with the project code.", first, "exact", [f"{a.code}:decision", f"{b.code}:decision"])
    # aggregation: sum of initial budgets of all projects sharing the most common currency (>= 2 projects when possible)
    curs = {}
    for x in projects: curs.setdefault(x.currency, []).append(x)
    cur = max(curs, key=lambda k: len(curs[k]))
    grp = curs[cur]
    add("aggregation", f"What is the total of the initial budgets of all projects whose budget is denominated in {cur}? Give the number only.", fmt_amt(sum(x.proposed_amount for x in grp)), "numeric", [f"{x.code}:proposed_amount" for x in grp], 0.0)
    # unit conversion
    p = P[6 % len(P)]; qv, ub, us, fac = p.quantity
    add("unit_conversion", f"What is the rated capacity of {p.code} expressed in {us}? Give the number only.", _norm_num(qv * fac), "numeric", [f"{p.code}:quantity"], 1e-6)
    # unanswerable
    dn = distractor_names[0] if distractor_names else "Zorblat-3"
    add("unanswerable", f"What is the supplier identifier for {dn}? If the document does not say, answer exactly: unknown.", "unknown", "abstain", [])
    p = P[7 % len(P)]; add("unanswerable", f"What is the postal code of the head office that issued this ledger? If the document does not say, answer exactly: unknown.", "unknown", "abstain", [])
    return Q

def generate_document(seed: int, n_projects: int, target_tokens: int, chars_per_token: float = 4.0) -> dict:
    rng = random.Random(seed)
    used = set()
    projects = [make_project(rng, used=used) for _ in range(n_projects)]
    # distractors: similar names (same base, different number) that appear only in a 'proposals under review' list without facts
    distractor_names = [f"{p.name}-{(int(p.code.split('-')[-1]) + rng.randint(1,5))}" for p in projects[: max(1, n_projects // 2)]]
    facts: list[Fact] = []
    blocks = [render_project(p, rng, facts, shuffle_sentences=False) for p in projects]
    header = (f"INTERNAL LEDGER MEMO — Capital Projects Register, issue {rng.randint(100,999)}. "
              "Where two entries conflict, the entry with the later date supersedes the earlier one. "
              "Projects listed under 'proposals under review' have no recorded attributes in this issue.")
    review = "Proposals under review (no attributes recorded): " + ", ".join(distractor_names) + "."
    # interleave filler to reach target length
    target_chars = int(target_tokens * chars_per_token)
    body = [header, review]
    cur = sum(len(b) for b in body) + sum(len(b) for b in blocks)
    n_fill = max(0, int((target_chars - cur) / 95))
    fill_between = max(1, n_fill // max(1, len(blocks)))
    for b in blocks:
        body.append(b)
        body.append(" ".join(rng.choice(FILLER) for _ in range(fill_between)))
    doc = "\n\n".join(body)
    questions = build_questions(projects, distractor_names, rng, doc, facts)
    for q in questions:
        for fid in q.supporting_fact_ids:
            assert any(f.fact_id == fid for f in facts), fid
    return {
        "document_id": f"syn_{seed:05d}_{target_tokens}",
        "seed": seed, "n_projects": n_projects, "target_tokens": target_tokens,
        "text": doc, "char_len": len(doc),
        "questions": [asdict(q) for q in questions],
        "facts": [asdict(f) for f in facts],
        "distractor_names": distractor_names,
    }