#!/usr/bin/env python3
"""Run one registry operation and write a bounded, source-pinned benchmark report."""

import argparse
import hashlib
import json
import os
import platform
import resource
import selectors
import signal
import sqlite3
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

MAX_MANIFEST_BYTES = 1024 * 1024
MAX_OUTPUT_BYTES = 16 * 1024 * 1024


def digest(data):
    return hashlib.sha256(data).hexdigest()


def file_digest(path):
    value = hashlib.sha256()
    with path.open("rb") as stream:
        while chunk := stream.read(1024 * 1024):
            value.update(chunk)
    return value.hexdigest()


def unique_object(pairs):
    value = {}
    for key, item in pairs:
        if key in value:
            raise ValueError("duplicate JSON key")
        value[key] = item
    return value


def read_manifest(path):
    if path.is_symlink() or not path.is_file():
        raise ValueError("source manifest must be a regular file")
    data = path.read_bytes()
    if len(data) > MAX_MANIFEST_BYTES:
        raise ValueError("source manifest exceeds 1 MiB")
    manifest = json.loads(data, object_pairs_hook=unique_object)
    required = ("schema", "source", "snapshot", "sha256", "bytes")
    if not isinstance(manifest, dict) or any(key not in manifest for key in required):
        raise ValueError("source manifest is incomplete")
    if (not isinstance(manifest["sha256"], str)
            or len(manifest["sha256"]) != 64
            or any(char not in "0123456789abcdef" for char in manifest["sha256"])):
        raise ValueError("source manifest has invalid SHA-256")
    if not isinstance(manifest["bytes"], int) or manifest["bytes"] <= 0:
        raise ValueError("source manifest has invalid byte count")
    return {
        "manifest_sha256": digest(data),
        "schema": manifest["schema"],
        "source": manifest["source"],
        "snapshot": manifest["snapshot"],
        "source_object_sha256": manifest["sha256"],
        "compressed_bytes": manifest["bytes"],
    }


def database_counts(path):
    if path is None or not path.exists():
        return {"bytes": 0, "records": 0, "facts": 0}
    if path.is_symlink() or not path.is_file():
        raise ValueError("database must be a regular file")
    connection = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
    try:
        page_count = connection.execute("PRAGMA page_count").fetchone()[0]
        page_size = connection.execute("PRAGMA page_size").fetchone()[0]
        tables = {
            row[0]
            for row in connection.execute(
                "SELECT name FROM sqlite_schema WHERE type='table'"
            )
        }
        records = connection.execute("SELECT count(*) FROM records").fetchone()[0] \
            if "records" in tables else 0
        facts = connection.execute("SELECT count(*) FROM facts").fetchone()[0] \
            if "facts" in tables else 0
        return {
            "bytes": page_count * page_size,
            "records": records,
            "facts": facts,
        }
    finally:
        connection.close()


def bounded_file(path):
    if path.stat().st_size > MAX_OUTPUT_BYTES:
        return None
    return path.read_bytes()


def hardware():
    cpu = "unknown"
    cpuinfo = Path("/proc/cpuinfo")
    if cpuinfo.is_file():
        for line in cpuinfo.read_text(errors="replace").splitlines():
            if line.startswith("model name"):
                cpu = line.partition(":")[2].strip()
                break
    memory_kib = None
    meminfo = Path("/proc/meminfo")
    if meminfo.is_file():
        first = meminfo.read_text(errors="replace").splitlines()[0].split()
        if len(first) >= 2 and first[0] == "MemTotal:":
            memory_kib = int(first[1])
    return {
        "operating_system": platform.system(),
        "kernel": platform.release(),
        "architecture": platform.machine(),
        "cpu": cpu,
        "logical_cpus": os.cpu_count(),
        "memory_kib": memory_kib,
    }


def run_bounded(command, stdout_path, stderr_path):
    process = subprocess.Popen(
        command,
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        start_new_session=True,
    )
    streams = selectors.DefaultSelector()
    streams.register(process.stdout, selectors.EVENT_READ, (stdout_path, 0))
    streams.register(process.stderr, selectors.EVENT_READ, (stderr_path, 0))
    exceeded = False
    outputs = {
        stdout_path: stdout_path.open("xb"),
        stderr_path: stderr_path.open("xb"),
    }
    try:
        while streams.get_map():
            for key, _events in streams.select():
                chunk = os.read(key.fileobj.fileno(), 64 * 1024)
                path, written = key.data
                if not chunk:
                    streams.unregister(key.fileobj)
                    key.fileobj.close()
                    continue
                accepted = min(len(chunk), MAX_OUTPUT_BYTES - written)
                if accepted > 0:
                    outputs[path].write(chunk[:accepted])
                    written += accepted
                    streams.modify(key.fileobj, selectors.EVENT_READ, (path, written))
                if accepted != len(chunk) and not exceeded:
                    exceeded = True
                    try:
                        os.killpg(process.pid, signal.SIGKILL)
                    except ProcessLookupError:
                        pass
        return process.wait(), exceeded
    finally:
        streams.close()
        for output in outputs.values():
            output.close()


def measured_artifacts(values):
    result = []
    labels = set()
    for value in values:
        label, separator, raw_path = value.partition("=")
        if not separator or not label or label in labels:
            raise ValueError("measured paths require unique LABEL=PATH values")
        labels.add(label)
        path = Path(raw_path)
        if path.is_symlink() or not path.is_file():
            raise ValueError("measured artifact must be a regular file")
        result.append({
            "label": label,
            "bytes": path.stat().st_size,
            "sha256": file_digest(path),
        })
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--profile", choices=("small", "medium", "provider"), required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--source-manifest", type=Path, action="append", default=[])
    parser.add_argument("--database", type=Path)
    parser.add_argument("--expanded-bytes", type=int)
    parser.add_argument("--measured-path", action="append", default=[])
    parser.add_argument("command", nargs=argparse.REMAINDER)
    args = parser.parse_args()
    if not args.command or args.command[0] != "--" or len(args.command) == 1:
        parser.error("terminate options with -- and provide a command")
    command = args.command[1:]
    if args.output.exists():
        raise FileExistsError("benchmark output already exists")
    if args.expanded_bytes is not None and args.expanded_bytes < 0:
        raise ValueError("expanded bytes cannot be negative")
    sources = [read_manifest(path) for path in args.source_manifest]
    before = database_counts(args.database)
    args.output.mkdir(parents=True)
    stdout_path = args.output / "stdout"
    stderr_path = args.output / "stderr"
    started_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
    usage_before = resource.getrusage(resource.RUSAGE_CHILDREN)
    started = time.monotonic()
    exit_status, output_limit_exceeded = run_bounded(command, stdout_path, stderr_path)
    wall_seconds = time.monotonic() - started
    usage_after = resource.getrusage(resource.RUSAGE_CHILDREN)
    after = database_counts(args.database)
    stdout = bounded_file(stdout_path)
    operation = None
    if stdout is not None:
        try:
            operation = json.loads(stdout, object_pairs_hook=unique_object)
        except (UnicodeDecodeError, json.JSONDecodeError, ValueError):
            pass
    fact_growth = after["facts"] - before["facts"]
    report = {
        "schema": "argand.site-benchmark/v1",
        "profile": args.profile,
        "started_at": started_at,
        "hardware": hardware(),
        "sources": sources,
        "operation": {
            "program": Path(command[0]).name,
            "argument_vector_sha256": digest(b"\0".join(os.fsencode(item) for item in command)),
            "exit_status": exit_status,
            "output_limit_exceeded": output_limit_exceeded,
            "stdout_sha256": file_digest(stdout_path),
            "stderr_sha256": file_digest(stderr_path),
            "result": operation,
        },
        "metrics": {
            "wall_seconds": round(wall_seconds, 6),
            "user_cpu_seconds": round(usage_after.ru_utime - usage_before.ru_utime, 6),
            "system_cpu_seconds": round(usage_after.ru_stime - usage_before.ru_stime, 6),
            "peak_rss_kib": usage_after.ru_maxrss,
            "compressed_source_bytes": sum(source["compressed_bytes"] for source in sources),
            "expanded_source_bytes": args.expanded_bytes,
            "database_before": before,
            "database_after": after,
            "database_growth_bytes": after["bytes"] - before["bytes"],
            "record_growth": after["records"] - before["records"],
            "fact_growth": fact_growth,
            "facts_per_second": round(fact_growth / wall_seconds, 3)
            if fact_growth > 0 and wall_seconds > 0 else None,
            "import_checkpoint_records": 256,
        },
        "artifacts": measured_artifacts(args.measured_path),
    }
    report_path = args.output / "REPORT.json"
    with report_path.open("xb") as stream:
        stream.write(json.dumps(report, indent=2, sort_keys=True).encode() + b"\n")
        stream.flush()
        os.fsync(stream.fileno())
    directory = os.open(args.output, os.O_RDONLY | os.O_DIRECTORY)
    try:
        os.fsync(directory)
    finally:
        os.close(directory)
    return exit_status


if __name__ == "__main__":
    try:
        sys.exit(main())
    except (FileExistsError, OSError, ValueError) as error:
        print(error, file=sys.stderr)
        sys.exit(2)