#!/usr/bin/env python3
"""CertiSigma Census — Performance Benchmark Suite.

Measures hash throughput, scan pipeline, and manifest operations.
All data is generated in a temporary directory and cleaned up on exit.

Usage::

    python scripts/benchmark.py
    python scripts/benchmark.py --files 10000 --workers 4
    python scripts/benchmark.py --output results.json

Results are printed as JSON to stdout (or written to --output).
"""

from __future__ import annotations

import argparse
import json
import os
import platform
import sys
import tempfile
import time
from pathlib import Path

from certisigma import hash_file
from certisigma_census import __version__
from certisigma_census.manifest import Manifest, ManifestEntry
from certisigma_census.scanner import walk_files, scan_directory

MAX_FILES = 100_000
MAX_ENTRIES = 1_000_000


def _generate_files(directory: Path, count: int, size_bytes: int = 1024) -> list[Path]:
    """Generate *count* files of *size_bytes* each in *directory*.

    Each file has unique content (index prefix) to avoid hash collisions.
    """
    directory.mkdir(parents=True, exist_ok=True)
    padding = os.urandom(min(size_bytes, 4096))
    files = []
    for i in range(count):
        p = directory / f"file_{i:06d}.bin"
        header = f"FILE{i:010d}".encode()
        with open(p, "wb") as f:
            f.write(header)
            remaining = size_bytes - len(header)
            if remaining > 0:
                f.write(padding[:remaining])
        files.append(p)
    return files


def bench_hash_throughput(tmpdir: Path, file_count: int) -> dict:
    """Measure SHA-256 hashing throughput."""
    files = _generate_files(tmpdir / "hash_bench", file_count, size_bytes=10240)
    total_bytes = sum(f.stat().st_size for f in files)

    t0 = time.perf_counter()
    for f in files:
        hash_file(str(f))
    elapsed = time.perf_counter() - t0

    return {
        "name": "hash_throughput",
        "files": file_count,
        "total_bytes": total_bytes,
        "elapsed_s": round(elapsed, 3),
        "throughput_mb_s": round(total_bytes / (1024 * 1024) / elapsed, 1) if elapsed > 0 else 0,
    }


def bench_walk_files(tmpdir: Path, file_count: int) -> dict:
    """Measure filesystem walk speed."""
    d = tmpdir / "walk_bench"
    _generate_files(d, file_count, size_bytes=64)

    t0 = time.perf_counter()
    result = list(walk_files(d))
    elapsed = time.perf_counter() - t0

    return {
        "name": "walk_files",
        "files": len(result),
        "elapsed_s": round(elapsed, 3),
        "files_per_s": round(len(result) / elapsed, 0) if elapsed > 0 else 0,
    }


def bench_scan_sequential(tmpdir: Path, file_count: int) -> dict:
    """Measure full scan pipeline with 1 worker."""
    d = tmpdir / "scan_seq"
    _generate_files(d, file_count, size_bytes=1024)

    t0 = time.perf_counter()
    m = scan_directory(d, workers=1, show_progress=False)
    elapsed = time.perf_counter() - t0
    total = m.total
    m.close()

    return {
        "name": "scan_sequential",
        "files": total,
        "elapsed_s": round(elapsed, 3),
        "files_per_s": round(total / elapsed, 0) if elapsed > 0 else 0,
    }


def bench_scan_parallel(tmpdir: Path, file_count: int, workers: int) -> dict:
    """Measure full scan pipeline with N workers."""
    d = tmpdir / "scan_par"
    _generate_files(d, file_count, size_bytes=1024)

    t0 = time.perf_counter()
    m = scan_directory(d, workers=workers, show_progress=False)
    elapsed = time.perf_counter() - t0
    total = m.total
    m.close()

    return {
        "name": "scan_parallel",
        "files": total,
        "workers": workers,
        "elapsed_s": round(elapsed, 3),
        "files_per_s": round(total / elapsed, 0) if elapsed > 0 else 0,
    }


def bench_manifest_add(count: int) -> dict:
    """Measure manifest insert throughput."""
    m = Manifest(root_dir="/bench")
    try:
        entries = [
            ManifestEntry(
                hash_hex=f"{i:064x}", path=f"dir/file_{i}.bin",
                size=1024, mtime=float(i),
            )
            for i in range(count)
        ]
        t0 = time.perf_counter()
        for e in entries:
            m.add(e)
        elapsed = time.perf_counter() - t0
    finally:
        m.close()
    return {
        "name": "manifest_add",
        "entries": count,
        "elapsed_s": round(elapsed, 3),
        "entries_per_s": round(count / elapsed, 0) if elapsed > 0 else 0,
    }


def bench_manifest_query(count: int) -> dict:
    """Measure manifest lookup throughput."""
    m = Manifest(root_dir="/bench")
    try:
        for i in range(count):
            m.add(ManifestEntry(
                hash_hex=f"{i:064x}", path=f"f_{i}.bin",
                size=1024, mtime=float(i),
            ))
        hashes = [f"{i:064x}" for i in range(count)]

        t0 = time.perf_counter()
        for h in hashes:
            m.get(h)
        elapsed = time.perf_counter() - t0
    finally:
        m.close()
    return {
        "name": "manifest_query",
        "entries": count,
        "elapsed_s": round(elapsed, 3),
        "queries_per_s": round(count / elapsed, 0) if elapsed > 0 else 0,
    }


def bench_manifest_save_load(tmpdir: Path, count: int) -> dict:
    """Measure manifest save/load round-trip."""
    m = Manifest(root_dir="/bench")
    try:
        for i in range(count):
            m.add(ManifestEntry(
                hash_hex=f"{i:064x}", path=f"f_{i}.bin",
                size=1024, mtime=float(i),
            ))
        db_path = tmpdir / "bench_manifest.db"

        t0 = time.perf_counter()
        m.save(db_path)
        elapsed_save = time.perf_counter() - t0
    finally:
        m.close()

    t0 = time.perf_counter()
    loaded = Manifest.load(db_path)
    elapsed_load = time.perf_counter() - t0
    total = loaded.total
    loaded.close()

    return {
        "name": "manifest_save_load",
        "entries": count,
        "save_s": round(elapsed_save, 3),
        "load_s": round(elapsed_load, 3),
        "total_s": round(elapsed_save + elapsed_load, 3),
        "verified_count": total,
    }


def main():
    parser = argparse.ArgumentParser(description="Census performance benchmarks")
    parser.add_argument("--files", type=int, default=1000, help="File count for hash/scan benchmarks")
    parser.add_argument("--entries", type=int, default=10000, help="Entry count for manifest benchmarks")
    parser.add_argument("--workers", type=int, default=os.cpu_count() or 4, help="Workers for parallel scan")
    parser.add_argument("--output", type=str, help="Write results to JSON file")
    args = parser.parse_args()

    if args.files > MAX_FILES:
        parser.error(f"--files exceeds limit ({MAX_FILES})")
    if args.entries > MAX_ENTRIES:
        parser.error(f"--entries exceeds limit ({MAX_ENTRIES})")
    if args.files < 1 or args.entries < 1:
        parser.error("--files and --entries must be >= 1")
    if args.workers < 1 or args.workers > 32:
        parser.error("--workers must be between 1 and 32")

    with tempfile.TemporaryDirectory(prefix="census-bench-") as tmpdir:
        tmpdir = Path(tmpdir)

        benchmarks = [
            bench_hash_throughput(tmpdir, args.files),
            bench_walk_files(tmpdir, args.files),
            bench_scan_sequential(tmpdir, args.files),
            bench_scan_parallel(tmpdir, args.files, args.workers),
            bench_manifest_add(args.entries),
            bench_manifest_query(args.entries),
            bench_manifest_save_load(tmpdir, args.entries),
        ]

    report = {
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "census_version": __version__,
        "python_version": platform.python_version(),
        "platform": f"{platform.system()} {platform.machine()}",
        "cpu_count": os.cpu_count(),
        "parameters": {"files": args.files, "entries": args.entries, "workers": args.workers},
        "benchmarks": benchmarks,
    }

    output = json.dumps(report, indent=2, default=str)
    if args.output:
        Path(args.output).write_text(output + "\n", encoding="utf-8")
        print(f"Results written to {args.output}", file=sys.stderr)
    else:
        print(output)


if __name__ == "__main__":
    main()
