"""Reproducible CPU/Arrow benchmark on one pinned Hugging Face shard.

Install with ``python -m pip install ".[arrow,benchmark]"`` and accept the
dataset's access terms before running this script.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import statistics
import time
from datetime import datetime, timezone

import pyarrow as pa
from datasets import load_dataset
from huggingface_hub import HfApi, get_token, login

from range_curate.examples.general_economy import DefaultTaxonomy, RelevanceEngine, __version__

DATASET = "Brianferrell787/financial-news-multisource"
REVISION = "c25780f336280adb57c64bda7aed605d065c672d"
SHARD = "data/sp500_daily_headlines/sp500_daily_headlines.000.parquet"


def utc_timestamp(value: str) -> datetime:
    parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=timezone.utc)
    return parsed.astimezone(timezone.utc)


def load_seed() -> tuple[pa.RecordBatch, str]:
    if get_token() is None:
        login()
    info = HfApi().dataset_info(DATASET, revision=REVISION, token=True)
    if info.sha != REVISION:
        raise RuntimeError(f"dataset revision drifted: {info.sha}")

    dataset = load_dataset(
        DATASET,
        name="data",
        data_files={"train": SHARD},
        split="train",
        revision=REVISION,
        token=True,
    ).select_columns(["date", "text"])

    texts = pa.array(dataset["text"], type=pa.string())
    timestamps = pa.array(
        [utc_timestamp(value) for value in dataset["date"]],
        type=pa.timestamp("us", tz="UTC"),
    )
    seed = pa.record_batch([texts, timestamps], names=["text", "timestamp"])

    digest = hashlib.sha256()
    for date, text in zip(dataset["date"], dataset["text"], strict=True):
        row = json.dumps(
            [date, text],
            ensure_ascii=False,
            separators=(",", ":"),
        ).encode("utf-8")
        digest.update(len(row).to_bytes(8, "little"))
        digest.update(row)
    return seed, digest.hexdigest()


def repeated_batches(
    seed: pa.RecordBatch,
    total_rows: int,
    *,
    batch_rows: int = 100_000,
) -> tuple[pa.RecordBatch, ...]:
    output = []
    emitted = 0
    while emitted < total_rows:
        start = emitted % seed.num_rows
        take = min(batch_rows, total_rows - emitted, seed.num_rows - start)
        output.append(seed.slice(start, take))
        emitted += take
    return tuple(output)


def reader(seed: pa.RecordBatch, batches: tuple[pa.RecordBatch, ...]):
    return pa.RecordBatchReader.from_batches(seed.schema, batches)


def semantic_digest(facts, total_rows: int) -> str:
    digest = hashlib.sha256(f"rows={total_rows}\n".encode())
    for row in facts.to_pyarrow().to_pylist():
        payload = json.dumps(
            row,
            sort_keys=True,
            separators=(",", ":"),
            default=str,
        ).encode("utf-8")
        digest.update(len(payload).to_bytes(8, "little"))
        digest.update(payload)
    return digest.hexdigest()


def measure(
    engine: RelevanceEngine,
    seed: pa.RecordBatch,
    total_rows: int,
    *,
    workers: int,
    repetitions: int,
) -> dict[str, object]:
    prepared = repeated_batches(seed, total_rows)
    warmup_rows = min(total_rows, 10_000)
    warmup = repeated_batches(seed, warmup_rows)
    engine.screen_arrow(reader(seed, warmup), workers=workers, row_count=warmup_rows)

    elapsed = []
    digests = []
    exceptions = []
    for _ in range(repetitions):
        started = time.perf_counter()
        facts = engine.screen_arrow(
            reader(seed, prepared),
            workers=workers,
            row_count=total_rows,
        )
        elapsed.append(time.perf_counter() - started)
        digests.append(semantic_digest(facts, total_rows))
        exceptions.append(facts.exception_count)

    if len(set(digests)) != 1 or len(set(exceptions)) != 1:
        raise RuntimeError("repetitions produced different semantic outputs")

    median = statistics.median(elapsed)
    return {
        "backend": "cpu-rust-arrow",
        "rows": total_rows,
        "physical_rows": seed.num_rows,
        "replication_factor": round(total_rows / seed.num_rows, 3),
        "workers": workers,
        "repetitions": repetitions,
        "median_seconds": round(median, 6),
        "rows_per_second": round(total_rows / median),
        "exception_rows": exceptions[0],
        "output_sha256": digests[0],
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--rows",
        type=int,
        nargs="+",
        default=[10_000, 100_000, 1_000_000],
    )
    parser.add_argument("--workers", type=int, default=min(10, os.cpu_count() or 1))
    parser.add_argument("--repetitions", type=int, default=3)
    args = parser.parse_args()

    seed, seed_sha256 = load_seed()
    taxonomy = DefaultTaxonomy()
    engine = RelevanceEngine(taxonomy=taxonomy, workers=args.workers)

    print(
        json.dumps(
            {
                "package": __version__,
                "taxonomy": taxonomy.fingerprint,
                "dataset": DATASET,
                "revision": REVISION,
                "shard": SHARD,
                "physical_rows": seed.num_rows,
                "seed_sha256": seed_sha256,
            },
            indent=2,
        )
    )
    for total_rows in args.rows:
        print(
            json.dumps(
                measure(
                    engine,
                    seed,
                    total_rows,
                    workers=args.workers,
                    repetitions=args.repetitions,
                ),
                sort_keys=True,
            )
        )


if __name__ == "__main__":
    main()