#!/usr/bin/env python3
"""Shared, dependency-free helpers for the S3-real preflight."""

from __future__ import annotations

import hashlib
import json
import math
from pathlib import Path
from typing import Any, Iterable, Iterator, Sequence


SOURCE_BLOCK_SIZE = 64
TARGET_BLOCK_SIZE = 16
SUBBLOCKS_PER_SOURCE = SOURCE_BLOCK_SIZE // TARGET_BLOCK_SIZE


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


def iter_jsonl(path: Path) -> Iterator[dict[str, Any]]:
    with path.open() as stream:
        for line_number, line in enumerate(stream, 1):
            if not line.strip():
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as error:
                raise ValueError(f"{path}:{line_number}: invalid JSON") from error
            if not isinstance(row, dict):
                raise ValueError(f"{path}:{line_number}: row must be an object")
            yield row


def timestamp(row: dict[str, Any]) -> float:
    return float(row["timestamp"])


def input_length(row: dict[str, Any]) -> int:
    return int(row["input_length"])


def output_length(row: dict[str, Any]) -> int:
    return max(1, int(row["output_length"]))


def parse_hash_ids(value: Any) -> list[int]:
    if isinstance(value, list):
        raw = value
    elif isinstance(value, str):
        stripped = value.strip()
        if not stripped:
            return []
        if stripped.startswith("["):
            raw = json.loads(stripped)
        else:
            delimiter = "|" if "|" in stripped else ","
            raw = [part.strip() for part in stripped.split(delimiter) if part.strip()]
    elif value is None:
        return []
    else:
        raw = [value]
    try:
        return [int(item) for item in raw]
    except (TypeError, ValueError) as error:
        raise ValueError(f"hash_ids must contain integers, got {value!r}") from error


def percentile(values: Sequence[float | int], fraction: float) -> float | None:
    if not values:
        return None
    if not 0 <= fraction <= 1:
        raise ValueError("percentile fraction must be in [0, 1]")
    ordered = sorted(float(value) for value in values)
    if len(ordered) == 1:
        return ordered[0]
    position = (len(ordered) - 1) * fraction
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return ordered[lower]
    return ordered[lower] * (upper - position) + ordered[upper] * (position - lower)


def distribution(values: Sequence[float | int]) -> dict[str, float | int | None]:
    return {
        "count": len(values),
        "min": min(values) if values else None,
        "p50": percentile(values, 0.50),
        "p90": percentile(values, 0.90),
        "p95": percentile(values, 0.95),
        "p99": percentile(values, 0.99),
        "max": max(values) if values else None,
        "mean": sum(values) / len(values) if values else None,
    }


def common_prefix_length(left: Sequence[Any], right: Sequence[Any]) -> int:
    count = 0
    for a, b in zip(left, right):
        if a != b:
            break
        count += 1
    return count


def zigzag(value: int) -> int:
    """Map every signed integer injectively to a non-negative integer."""
    return 2 * value if value >= 0 else -2 * value - 1


def expanded_block_id(
    source_hash: int,
    subblock_index: int,
    *,
    subblocks_per_source: int = SUBBLOCKS_PER_SOURCE,
) -> int:
    if not 0 <= subblock_index < subblocks_per_source:
        raise ValueError(f"subblock index out of range: {subblock_index}")
    return zigzag(int(source_hash)) * subblocks_per_source + subblock_index


def expand_hash_ids(
    source_hashes: Sequence[int],
    num_tokens: int,
    *,
    source_block_size: int = SOURCE_BLOCK_SIZE,
    target_block_size: int = TARGET_BLOCK_SIZE,
) -> list[int]:
    if num_tokens <= 0:
        raise ValueError(f"input_length must be positive, got {num_tokens}")
    if source_block_size <= 0 or target_block_size <= 0:
        raise ValueError("block sizes must be positive")
    if source_block_size % target_block_size:
        raise ValueError("source block size must be divisible by target block size")
    subblocks_per_source = source_block_size // target_block_size
    required_source_blocks = math.ceil(num_tokens / source_block_size)
    if len(source_hashes) != required_source_blocks:
        raise ValueError(
            f"strict {source_block_size}-block contract failed: "
            f"input_length={num_tokens} requires {required_source_blocks} source blocks, "
            f"got {len(source_hashes)}"
        )
    target_blocks = math.ceil(num_tokens / target_block_size)
    expanded = [
        expanded_block_id(
            source_hash,
            subblock,
            subblocks_per_source=subblocks_per_source,
        )
        for source_hash in source_hashes
        for subblock in range(subblocks_per_source)
    ]
    return expanded[:target_blocks]


def token_payload(tokens: Sequence[int]) -> bytes:
    return len(tokens).to_bytes(2, "little") + b"".join(
        int(token).to_bytes(4, "little", signed=False) for token in tokens
    )


def token_block_identity_records(
    token_ids: Sequence[int], block_size: int = TARGET_BLOCK_SIZE
) -> list[tuple[int, bytes]]:
    """Return parent-sensitive block IDs and independent collision witnesses."""
    if block_size <= 0:
        raise ValueError("block_size must be positive")
    parent = b"FRONTIER_EXACT_TRACE_ROOT"
    records: list[tuple[int, bytes]] = []
    for start in range(0, len(token_ids), block_size):
        payload = token_payload(token_ids[start : start + block_size])
        identity_input = parent + b"\0" + payload
        parent = hashlib.blake2b(identity_input, digest_size=16).digest()
        records.append(
            (
                int.from_bytes(parent, "big", signed=False),
                hashlib.sha256(identity_input).digest(),
            )
        )
    return records


def token_block_identities(
    token_ids: Sequence[int], block_size: int = TARGET_BLOCK_SIZE
) -> list[int]:
    return [identity for identity, _ in token_block_identity_records(token_ids, block_size)]


def token_block(block_id: int, *, vocab_size: int, token_offset: int = 1024) -> list[int]:
    """Encode one block identity injectively as exactly sixteen valid token IDs."""
    base = vocab_size - token_offset
    if base < 256:
        raise ValueError("vocab has too few non-reserved token IDs")
    value = int(block_id)
    if value < 0:
        raise ValueError("expanded block IDs must be non-negative")
    digits = []
    for _ in range(TARGET_BLOCK_SIZE):
        digits.append(token_offset + value % base)
        value //= base
    if value:
        raise ValueError("block identity exceeds the injective 16-token code space")
    return digits


def synthetic_tokens(
    block_ids: Sequence[int], num_tokens: int, *, vocab_size: int, token_offset: int = 1024
) -> list[int]:
    tokens = [
        token
        for block_id in block_ids
        for token in token_block(block_id, vocab_size=vocab_size, token_offset=token_offset)
    ][:num_tokens]
    if len(tokens) != num_tokens:
        raise ValueError(f"synthetic token length mismatch: {len(tokens)} != {num_tokens}")
    return tokens


def stable_int(value: Any) -> int:
    raw = json.dumps(value, sort_keys=True, separators=(",", ":")).encode()
    return int.from_bytes(hashlib.sha256(raw).digest()[:8], "big") & ((1 << 63) - 1)


def resolve_session_root(row: dict[str, Any], root_of: dict[Any, Any]) -> Any:
    chat = row.get("chat_id")
    parent = row.get("parent_chat_id")
    has_parent = parent not in (None, "", -1, "-1")
    root = root_of.get(parent, parent) if has_parent else chat
    if chat is not None:
        root_of[chat] = root
    return root


def session_uniform(*, seed: int, window_id: str, session_root: Any) -> float:
    payload = json.dumps(
        {"seed": seed, "window_id": window_id, "session_root": session_root},
        sort_keys=True,
        separators=(",", ":"),
    ).encode()
    digest = hashlib.blake2b(payload, digest_size=8).digest()
    return int.from_bytes(digest, "big") / float(1 << 64)


def session_id(row: dict[str, Any]) -> int:
    return stable_int(
        row.get("session_root", row.get("chat_id", row.get("parent_chat_id", "missing-chat")))
    )


def batched(iterable: Iterable[Any], size: int) -> Iterator[list[Any]]:
    batch: list[Any] = []
    for item in iterable:
        batch.append(item)
        if len(batch) == size:
            yield batch
            batch = []
    if batch:
        yield batch