"""Training-stream composition: balanced real shards mixed with synthetic data.

The first real-corpus integration read shards back to back, so a run saw every
WordNet example before its first CVE and the domain mix at any step was one
domain.  :class:`BalancedShardStream` instead samples a shard per example with
temperature-scaled weights, ``p_i ∝ n_i ** temperature``: 1.0 is proportional
to size, 0.0 is uniform over shards, and the default 0.5 keeps small domains
visible without letting them repeat far more often than large ones.  A bounded
shuffle buffer breaks up the within-shard order (records are written in source
order, so neighbours share a graph neighbourhood).
"""

from __future__ import annotations

import gzip
import json
import pickle
import random
import re
from collections.abc import Iterable, Iterator, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from graph.synthetic.generator import ReasoningExample

#: Sources that ``scripts/benchmark_real.py`` evaluates on.  Their corpus shards
#: are sampled from the *full* graphs, test triples included, by the same kind
#: of path walk the benchmark uses, so training on them would score the model
#: on queries it has seen.  Excluded unless a caller opts in.
BENCHMARK_SOURCES: tuple[str, ...] = ("wn18rr", "fb15k237", "hotpotqa")

#: Development leftovers in ``training/real_corpus`` (smoke, profiling and
#: speed-test shards) that duplicate slices of real shards.
SCRATCH_SHARD = re.compile(r"(smoke|profile|speed)", re.IGNORECASE)


def shard_size(path: Path) -> int:
    """Record count from the shard's manifest, else by counting lines."""
    manifest = Path(f"{path}.manifest.json")
    if manifest.exists():
        payload = json.loads(manifest.read_text(encoding="utf-8"))
        for keys in (("stats", "examples"), ("output", "records"), ("records",)):
            value: object = payload
            for key in keys:
                value = value.get(key) if isinstance(value, dict) else None
            if isinstance(value, int) and value > 0:
                return value
    opener = gzip.open if path.suffix == ".gz" else open
    with opener(path, "rt", encoding="utf-8") as handle:
        return sum(1 for line in handle if line.strip())


def expand_shards(paths: Sequence[str | Path]) -> list[Path]:
    """Accept shard files or directories (searched for ``*.jsonl.gz``)."""
    shards: list[Path] = []
    for entry in map(Path, paths):
        if entry.is_dir():
            shards.extend(sorted(entry.rglob("*.jsonl.gz")))
        else:
            shards.append(entry)
    if not shards:
        raise ValueError(f"no real-corpus shards found in {list(map(str, paths))}")
    return shards


def training_shards_from_manifest(path: str | Path) -> list[Path]:
    """Return the prepared manifest's included training shard paths."""
    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    entries = payload.get("shards", [])
    if not isinstance(entries, list):
        raise TypeError(f"{path}: training manifest has no shard list")
    selected: list[Path] = []
    for entry in entries:
        if not isinstance(entry, dict) or not entry.get("include", False):
            continue
        shard_path = Path(str(entry.get("path", "")))
        if not shard_path.exists():
            raise FileNotFoundError(f"{path}: missing training shard {shard_path}")
        selected.append(shard_path)
    if not selected:
        raise ValueError(f"{path}: no included training shards")
    return selected


class BalancedShardStream:
    """Endless, seeded, temperature-balanced stream over real-corpus shards."""

    def __init__(
        self,
        paths: Sequence[str | Path],
        *,
        temperature: float = 0.5,
        seed: int = 0,
        shuffle_buffer: int = 2048,
        exclude_roles: Iterable[str] = ("evaluation_only",),
        exclude_sources: Iterable[str] = BENCHMARK_SOURCES,
        single_pass: bool = False,
        start: int = 0,
    ) -> None:
        if not 0.0 <= temperature <= 1.0:
            raise ValueError("temperature must be in [0, 1]")
        roles, sources = set(exclude_roles), set(exclude_sources)
        self.excluded: list[tuple[str, str]] = []
        self.shards = []
        for path in expand_shards(paths):
            manifest = _manifest(path)
            source = manifest.get("source") or {}
            reason = (
                "scratch shard"
                if SCRATCH_SHARD.search(path.name)
                else "evaluation role"
                if _role(path, manifest) in roles
                else "benchmark source"
                if source.get("source_id") in sources
                else ""
            )
            if reason:
                self.excluded.append((str(path), reason))
            else:
                self.shards.append(path)
        if not self.shards:
            raise ValueError(f"every shard was excluded: {self.excluded}")
        self.sizes = [shard_size(path) for path in self.shards]
        raw = [size**temperature for size in self.sizes]
        total = sum(raw)
        self.weights = [value / total for value in raw]
        self.seed = seed
        self.shuffle_buffer = max(1, shuffle_buffer)
        self.single_pass = single_pass
        #: Single-pass records to skip unparsed, for resuming mid-pass.
        self.start = start
        #: Records drawn from each shard so far, for the training summary.
        self.reads = [0] * len(self.shards)

    def __iter__(self) -> Iterator[ReasoningExample]:
        if self.single_pass:
            yield from self._single_pass()
            return
        from graph.real.corpus import deserialize_example

        rng = random.Random(self.seed)
        readers = [_cycle_lines(path) for path in self.shards]
        indices = range(len(self.shards))
        buffer: list[str] = []
        while True:
            shard = rng.choices(indices, weights=self.weights)[0]
            buffer.append(next(readers[shard]))
            self.reads[shard] += 1
            if len(buffer) >= self.shuffle_buffer:
                slot = rng.randrange(len(buffer))
                buffer[slot], buffer[-1] = buffer[-1], buffer[slot]
                yield deserialize_example(json.loads(buffer.pop()))

    def _single_pass(self) -> Iterator[ReasoningExample]:
        """Every record of every shard exactly once.

        Shards are drawn in proportion to the records they have left, so all
        of them drain together and the domain mix stays even across the pass;
        ``temperature`` does not apply.
        """
        from graph.real.corpus import deserialize_example

        rng = random.Random(self.seed)
        readers = [_read_lines(path) for path in self.shards]
        remaining = list(self.sizes)
        live = [index for index, size in enumerate(remaining) if size]
        buffer: list[str] = []
        emitted = 0
        while live:
            shard = rng.choices(live, weights=[remaining[i] for i in live])[0]
            line = next(readers[shard], None)
            if line is None:
                remaining[shard] = 0
            else:
                buffer.append(line)
                self.reads[shard] += 1
                remaining[shard] -= 1
                if remaining[shard] <= 0:
                    # Manifest counts can lag the file; drain what is left.
                    for extra in readers[shard]:
                        buffer.append(extra)
                        self.reads[shard] += 1
            live = [index for index in live if remaining[index] > 0]
            if len(buffer) >= self.shuffle_buffer or (not live and buffer):
                while buffer and (len(buffer) >= self.shuffle_buffer or not live):
                    slot = rng.randrange(len(buffer))
                    buffer[slot], buffer[-1] = buffer[-1], buffer[slot]
                    line = buffer.pop()
                    emitted += 1
                    if emitted > self.start:
                        yield deserialize_example(json.loads(line))

    def describe(self) -> list[dict[str, object]]:
        return [
            {
                "path": str(path),
                "records": size,
                "weight": round(weight, 5),
                "reads": reads,
            }
            for path, size, weight, reads in zip(
                self.shards, self.sizes, self.weights, self.reads, strict=True
            )
        ]


def _manifest(path: Path) -> dict[str, Any]:
    manifest = Path(f"{path}.manifest.json")
    if not manifest.exists():
        return {}
    payload = json.loads(manifest.read_text(encoding="utf-8"))
    return payload if isinstance(payload, dict) else {}


def _role(path: Path, manifest: dict[str, Any]) -> str:
    source = manifest.get("source") or {}
    if isinstance(source, dict) and source.get("training_role"):
        return str(source["training_role"])
    if manifest.get("training_role"):
        return str(manifest["training_role"])
    return "evaluation_only" if "/eval/" in path.as_posix() else "train"


def _read_lines(path: Path) -> Iterator[str]:
    opener = gzip.open if path.suffix == ".gz" else open
    with opener(path, "rt", encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                yield line


def _cycle_lines(path: Path) -> Iterator[str]:
    opener = gzip.open if path.suffix == ".gz" else open
    while True:
        produced = False
        with opener(path, "rt", encoding="utf-8") as handle:
            for line in handle:
                if line.strip():
                    produced = True
                    yield line
        if not produced:
            raise RuntimeError(f"{path} contains no records")


def mix_streams(
    synthetic: Iterable[ReasoningExample],
    real: Iterable[ReasoningExample],
    *,
    count: int,
    real_fraction: float,
) -> Iterator[ReasoningExample]:
    """Deterministically interleave the two streams at ``real_fraction``."""
    if not 0.0 < real_fraction <= 1.0:
        raise ValueError("real_fraction must be in (0, 1]")
    synthetic_iter = iter(synthetic)
    real_iter = iter(real)
    budget = 0.0
    for _ in range(count):
        budget += real_fraction
        if budget >= 1.0:
            budget -= 1.0
            yield next(real_iter)
        else:
            yield next(synthetic_iter)


def mix_exhaustive(
    synthetic: Iterable[ReasoningExample],
    real: Iterable[ReasoningExample],
    *,
    synthetic_count: int,
    real_count: int,
    start: int = 0,
) -> Iterator[ReasoningExample]:
    """Interleave two finite streams evenly until both are exhausted.

    The real share is ``real_count / (synthetic_count + real_count)``.  The
    counts only set the ratio: when one stream ends, the other is drained, so
    no record of either is dropped.  ``start`` resumes the interleaving after
    that many mixed items; the caller must already have advanced each stream
    by the amounts :func:`exhaustive_split` returns.
    """
    fraction = _real_share(synthetic_count, real_count)
    synthetic_iter, real_iter = iter(synthetic), iter(real)
    _, _, budget = exhaustive_split(start, synthetic_count, real_count)
    while True:
        budget += fraction
        if budget >= 1.0:
            budget -= 1.0
            primary, other = real_iter, synthetic_iter
        else:
            primary, other = synthetic_iter, real_iter
        item = next(primary, None)
        if item is None:
            yield from other
            return
        yield item


def _real_share(synthetic_count: int, real_count: int) -> float:
    total = synthetic_count + real_count
    if total <= 0:
        raise ValueError("nothing to mix")
    return real_count / total


def exhaustive_split(
    start: int, synthetic_count: int, real_count: int
) -> tuple[int, int, float]:
    """How :func:`mix_exhaustive` divides its first ``start`` items.

    Returns ``(synthetic_taken, real_taken, budget)`` so a resumed run can skip
    exactly the records it already trained on.
    """
    fraction = _real_share(synthetic_count, real_count)
    synthetic_taken = real_taken = 0
    budget = 0.0
    for _ in range(start):
        budget += fraction
        if budget >= 1.0:
            budget -= 1.0
            if real_taken < real_count:
                real_taken += 1
            else:
                synthetic_taken += 1
        elif synthetic_taken < synthetic_count:
            synthetic_taken += 1
        else:
            real_taken += 1
    return synthetic_taken, real_taken, budget


@dataclass(frozen=True)
class EpochSpec:
    """Everything needed to rebuild a single-pass training stream elsewhere.

    The stream is rebuilt from this in a producer process, so it holds only
    plain values; see :func:`single_pass_stream`.
    """

    synthetic_root: str
    warmup_steps: int
    stage_steps: int
    seed: int
    synthetic_per_step: float
    synthetic_count: int
    real_paths: tuple[str, ...]
    real_count: int
    exclude_sources: tuple[str, ...]
    #: Mixed examples already trained on (a resumed run skips them).
    start: int = 0


class EpochStream:
    """The single-pass stream of an :class:`EpochSpec`, with its shard reads."""

    def __init__(self, spec: EpochSpec) -> None:
        from graph.synthetic.corpus import SyntheticCorpus
        from graph.synthetic.curriculum import CurriculumScheduler

        synthetic_skip, real_skip, _ = exhaustive_split(
            spec.start, spec.synthetic_count, spec.real_count
        )
        scheduler = CurriculumScheduler(
            warmup_steps=spec.warmup_steps, stage_steps=spec.stage_steps
        )
        self.synthetic = SyntheticCorpus(spec.synthetic_root).epoch_stream(
            scheduler,
            seed=spec.seed,
            examples_per_step=max(spec.synthetic_per_step, 1e-9),
            start=synthetic_skip,
        )
        self.real = (
            BalancedShardStream(
                spec.real_paths,
                seed=spec.seed,
                exclude_sources=spec.exclude_sources,
                single_pass=True,
                start=real_skip,
            )
            if spec.real_paths
            else None
        )
        self.spec = spec

    def __iter__(self) -> Iterator[ReasoningExample]:
        if self.real is None:
            return iter(self.synthetic)
        return mix_exhaustive(
            self.synthetic,
            self.real,
            synthetic_count=self.spec.synthetic_count,
            real_count=self.spec.real_count,
            start=self.spec.start,
        )

    def describe(self) -> list[dict[str, object]]:
        return self.real.describe() if self.real is not None else []


class BackgroundStream:
    """:class:`EpochStream` produced in a separate process.

    Regenerating and verifying synthetic examples ran on the training thread
    and cost about a quarter of every step.  A spawned producer overlaps it with
    the GPU work; the order of examples is unchanged.
    """

    def __init__(self, spec: EpochSpec, *, chunk: int = 64, depth: int = 64) -> None:
        import multiprocessing

        context = multiprocessing.get_context("spawn")
        self._queue = context.Queue(maxsize=depth)
        self._process = context.Process(
            target=_produce, args=(spec, self._queue, chunk), daemon=True
        )
        self._process.start()
        self.reads: list[dict[str, object]] = []
        #: Examples handed to the consumer.
        self._yielded = 0

    def __iter__(self) -> Iterator[ReasoningExample]:
        import queue as queue_module

        while True:
            try:
                kind, payload = self._queue.get(timeout=60)
            except queue_module.Empty:
                if not self._process.is_alive():
                    raise RuntimeError(
                        f"data producer exited with code {self._process.exitcode}"
                    ) from None
                continue
            if kind == "examples":
                for example in pickle.loads(payload):
                    self._yielded += 1
                    yield example
            elif kind == "done":
                self.reads = payload
                return
            else:
                raise RuntimeError(f"data producer failed:\n{payload}")

    def describe(self) -> list[dict[str, object]]:
        """Per-shard reads, available once the producer has finished the pass."""
        return self.reads

    def finish(self) -> int:
        """Stop the producer and return how many examples were trained on.

        A complete pass consumed every record of the spec.  The per-shard reads
        are collected when the producer had already reached the end; a run that
        stopped early has no final reads, and the producer is not left running
        to the end of the pass just to report them.
        """
        import queue as queue_module

        # A producer that reached the end of the pass exits by itself once its
        # last messages are flushed; only then is the queue safe to read.  One
        # still mid-pass is terminated, possibly halfway through writing a
        # message, so its queue is abandoned: reading it could block forever,
        # and a pass that did not finish has no end-of-pass reads anyway.
        self._process.join(timeout=5)
        if self._process.is_alive():
            self.close()
            return self._yielded
        while not self.reads:
            try:
                kind, payload = self._queue.get(timeout=5)
            except queue_module.Empty:
                break
            if kind == "done":
                self.reads = payload
            elif kind != "examples":
                raise RuntimeError(f"data producer failed:\n{payload}")
        self.close()
        return self._yielded

    def close(self) -> None:
        if self._process.is_alive():
            self._process.terminate()
        self._process.join(timeout=10)


def _produce(spec: EpochSpec, queue: Any, chunk: int) -> None:
    import traceback

    try:
        stream = EpochStream(spec)
        batch: list[ReasoningExample] = []
        for example in stream:
            batch.append(example)
            if len(batch) >= chunk:
                queue.put(("examples", _pack(batch)))
                batch = []
        if batch:
            queue.put(("examples", _pack(batch)))
        queue.put(("done", stream.describe()))
    except BaseException:
        queue.put(("error", traceback.format_exc()))
        raise


def _pack(examples: list[ReasoningExample]) -> bytes:
    """Serialise a chunk by value.

    Handing tensors to a multiprocessing queue directly makes torch share each
    one through its own file descriptor, and a training pass exhausted the
    process limit within minutes; plain pickle copies them into the payload.
    """
    return pickle.dumps(examples, protocol=pickle.HIGHEST_PROTOCOL)