from __future__ import annotations
import random
import torch
from . import bio, real_pii, synthetic_pii
from .taxonomy import TAG_TO_ID

PAD_ID = bio.PAD_ID
IGNORE = -100


def pii_stream(
    n_synth: int,
    seed: int = 0,
    use_real: bool = True,
    real_cap: int = 20000,
    long_frac: float = 0.3,
):
    syn = synthetic_pii.generate(n_synth, seed=seed, long_frac=long_frac)
    if not use_real:
        yield from syn
        return
    reals = [
        real_pii.load_nemotron(max_rows=real_cap),
        real_pii.load_gretel_finance(max_rows=real_cap),
    ]
    rng = random.Random(seed)
    pools = [syn] + reals
    alive = list(pools)
    while alive:
        src = rng.choice(alive)
        try:
            yield next(src)
        except StopIteration:
            alive.remove(src)


def make_pii_batch(examples, max_len: int):
    rows = [bio.spans_to_bio(t, s, max_len, TAG_TO_ID) for t, s in examples]
    length = min(max((len(ids) for ids, _ in rows), default=1), max_len)
    b = len(rows)
    ids = torch.full((b, length), PAD_ID, dtype=torch.long)
    mask = torch.zeros((b, length), dtype=torch.long)
    tags = torch.full((b, length), IGNORE, dtype=torch.long)
    for i, (row_ids, row_tags) in enumerate(rows):
        n = min(len(row_ids), length)
        ids[i, :n] = torch.tensor(row_ids[:n])
        mask[i, :n] = 1
        tags[i, :n] = torch.tensor(row_tags[:n])
    return (ids, mask, tags)


def iter_pii_batches(
    batch_size: int, context_lengths, n_synth: int, seed: int = 0, use_real: bool = True
):
    lengths = context_lengths or [512]
    rng = random.Random(seed + 1)
    buf, stream = ([], pii_stream(n_synth, seed=seed, use_real=use_real))
    for ex in stream:
        buf.append(ex)
        if len(buf) >= batch_size:
            yield make_pii_batch(buf, rng.choice(lengths))
            buf = []
    if buf:
        yield make_pii_batch(buf, rng.choice(lengths))