#!/usr/bin/env bash
# =============================================================================
# Pyvorin Edge — Benchmark Proof Orchestrator
# =============================================================================
# Run this on your Pi 5 to generate a tamper-evident, hardware-attested
# benchmark report that customers can verify independently.
#
# Usage:
#   cd ~/pyvorin-edge
#   bash pifiles/prove-benchmark.sh
#
# Outputs:
#   ~/pyvorin-edge/proof/benchmark_attested.json   — signed, attested results
#   ~/pyvorin-edge/proof/benchmark_report.md       — human-readable summary
#   ~/pyvorin-edge/proof/public_key.b64            — key to verify signature
#   ~/pyvorin-edge/proof/RECORDING.txt             — asciinema instructions
# =============================================================================

set -euo pipefail

INSTALL_DIR=""
SKIP_RECORDING=0

# Colours
if [[ -t 1 ]]; then
    GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; BLUE='\033[0;34m'; BOLD='\033[1m'; NC='\033[0m'
else
    GREEN=''; YELLOW=''; RED=''; BLUE=''; BOLD=''; NC=''
fi

log_info()  { echo -e "${BLUE}[INFO]${NC}  $*"; }
log_ok()    { echo -e "${GREEN}[OK]${NC}    $*"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC}  $*"; }
log_err()   { echo -e "${RED}[ERROR]${NC} $*" >&2; }
log_title() { echo -e "\n${BOLD}$*${NC}"; }

usage() {
cat <<'USAGE'
Pyvorin Edge — Benchmark Proof Orchestrator

Usage:
  bash prove-benchmark.sh

Options:
  --dir PATH          Install directory (default: auto-detect)
  --skip-recording    Skip asciinema screen recording setup
  --help              Show this help

What this does:
  1. Generates a one-time Ed25519 key pair for signing
  2. Collects hardware attestation (CPU temp, serial, model, etc.)
  3. Runs all 8 vertical benchmarks with telemetry
  4. Signs the JSON output cryptographically
  5. Generates a human-readable markdown report
  6. Provides instructions for screen recording
USAGE
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --dir)          INSTALL_DIR="$2"; shift 2 ;;
        --skip-recording) SKIP_RECORDING=1; shift ;;
        --help|-h)      usage; exit 0 ;;
        *) log_err "Unknown option: $1"; usage; exit 1 ;;
    esac
done

# Auto-detect install dir
if [[ -z "${INSTALL_DIR}" ]]; then
    if [[ -f "pyproject.toml" && -d "edge_sdk" && -d "venv" ]]; then
        INSTALL_DIR="$(pwd)"
    elif [[ -d "${HOME}/pyvorin-edge" ]]; then
        INSTALL_DIR="${HOME}/pyvorin-edge"
    else
        log_err "Cannot find a pyvorin-edge install."
        exit 1
    fi
fi

cd "${INSTALL_DIR}"
source "${INSTALL_DIR}/venv/bin/activate"

PROOF_DIR="${INSTALL_DIR}/proof"
mkdir -p "${PROOF_DIR}"

log_title "Pyvorin Edge — Benchmark Proof Orchestrator"
log_info "Directory:  ${INSTALL_DIR}"
log_info "Proof dir:  ${PROOF_DIR}"

# ---------------------------------------------------------------------------
# Step 1 — Generate signing key pair
# ---------------------------------------------------------------------------
log_title "Step 1/6 — Generate Ed25519 signing key"
KEY_PY=$(cat <<'PYEOF'
import base64
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives import serialization

priv = Ed25519PrivateKey.generate()
pub = priv.public_key()

priv_b64 = base64.b64encode(priv.private_bytes(
    encoding=serialization.Encoding.Raw,
    format=serialization.PrivateFormat.Raw,
    encryption_algorithm=serialization.NoEncryption()
)).decode('ascii')

pub_b64 = base64.b64encode(pub.public_bytes(
    encoding=serialization.Encoding.Raw,
    format=serialization.PublicFormat.Raw
)).decode('ascii')

print("PRIV:" + priv_b64)
print("PUB:" + pub_b64)
PYEOF
)

KEYPAIR=$(python -c "${KEY_PY}")
PRIV_KEY=$(echo "${KEYPAIR}" | grep "^PRIV:" | cut -d: -f2)
PUB_KEY=$(echo "${KEYPAIR}" | grep "^PUB:" | cut -d: -f2)

echo "${PUB_KEY}" > "${PROOF_DIR}/public_key.b64"
log_ok "Key pair generated. Public key saved to proof/public_key.b64"

# ---------------------------------------------------------------------------
# Step 2 — Check asciinema
# ---------------------------------------------------------------------------
log_title "Step 2/6 — Screen recording setup"
if [[ "${SKIP_RECORDING}" -eq 1 ]]; then
    log_warn "Screen recording skipped (--skip-recording)"
elif command -v asciinema &>/dev/null; then
    log_ok "asciinema is installed."
    log_info "To record:  asciinema rec proof/benchmark.cast"
    log_info "To upload:  asciinema upload proof/benchmark.cast"
else
    log_warn "asciinema not installed."
    log_info "Install with:  sudo apt install -y asciinema"
    log_info "Or record with:  script -t proof/typescript.log -a proof/typescript"
fi

# ---------------------------------------------------------------------------
# Step 3 — Hardware sanity check
# ---------------------------------------------------------------------------
log_title "Step 3/6 — Hardware sanity check"
python -c "
from edge_sdk.pyvorin_edge.attestation import collect_hardware_info
hw = collect_hardware_info()
print('  Platform:      ' + hw.platform)
print('  Machine:       ' + hw.machine)
print('  CPU model:     ' + hw.cpu_model)
print('  Serial:        ' + hw.serial_number[:20] + '...')
print('  Revision:      ' + hw.hardware_revision)
print('  CPU temp:      %.1f°C' % hw.cpu_temp_c)
print('  Throttled:     ' + (hw.cpu_throttled or 'N/A'))
print('  Uptime:        %.1f hours' % (hw.uptime_seconds / 3600))
print('  Memory:        %d MB total' % (hw.total_memory_kb // 1024))
"

# Warn if throttled
THROTTLED=$(python -c "from edge_sdk.pyvorin_edge.attestation import collect_hardware_info; print(collect_hardware_info().cpu_throttled)" 2>/dev/null || true)
THROTTLED_VAL="${THROTTLED#throttled=}"
if [[ -n "${THROTTLED_VAL}" && "${THROTTLED_VAL}" != "0x0" ]]; then
    log_warn "CPU is throttled (${THROTTLED}) — results may not represent peak performance."
fi

# ---------------------------------------------------------------------------
# Step 4 — Run attested benchmarks
# ---------------------------------------------------------------------------
log_title "Step 4/6 — Run attested benchmarks"
cd "${INSTALL_DIR}/examples/industry_benchmarks"
python run_all.py \
    --attest \
    --sign \
    --private-key "${PRIV_KEY}" \
    --public-key "${PUB_KEY}" \
    --output "${PROOF_DIR}/benchmark_attested.json"
cd "${INSTALL_DIR}"
log_ok "Benchmark complete. Attested JSON: proof/benchmark_attested.json"

# ---------------------------------------------------------------------------
# Step 5 — Generate human-readable report
# ---------------------------------------------------------------------------
log_title "Step 5/6 — Generate markdown report"
python3 - "${PROOF_DIR}" <<'PYEOF'
import json, base64, sys
from pathlib import Path
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey

PROOF_DIR = sys.argv[1]
data = json.loads(Path(PROOF_DIR + '/benchmark_attested.json').read_text())
pub_key_b64 = data.get('public_key_b64', '')
sig_b64 = data.get('signature', '')
integrity = data.get('integrity_hash', 'N/A')
hw = data.get('hardware', {})
rt = data.get('runtime', {})
results = data.get('results', [])

# Verify signature
sig_valid = False
if pub_key_b64 and sig_b64:
    try:
        pub = Ed25519PublicKey.from_public_bytes(base64.b64decode(pub_key_b64))
        payload = {k: v for k, v in data.items() if k not in ('signature', 'public_key_b64')}
        canonical = json.dumps(payload, sort_keys=True, separators=(',', ':')).encode('utf-8')
        pub.verify(base64.b64decode(sig_b64), canonical)
        sig_valid = True
    except Exception:
        pass

lines = [
    '# Pyvorin Edge — Benchmark Proof Report',
    '',
    '**Generated:** ' + rt.get('timestamp_utc', 'N/A') + '  ',
    '**Integrity Hash:** `' + integrity + '`  ',
    '**Signature Valid:** ' + ('YES ✅' if sig_valid else 'NO ❌') + '  ',
    '',
    '---',
    '## Hardware Attestation',
    '',
    '| Property | Value |',
    '|----------|-------|',
    '| Platform | ' + hw.get('platform', 'N/A') + ' |',
    '| Machine  | ' + hw.get('machine', 'N/A') + ' |',
    '| CPU Model | ' + hw.get('cpu_model', 'N/A') + ' |',
    '| CPU Cores | ' + str(hw.get('cpu_count', 'N/A')) + ' |',
    '| Serial Number | `' + str(hw.get('serial_number', 'N/A'))[:24] + '...` |',
    '| Hardware Revision | ' + hw.get('hardware_revision', 'N/A') + ' |',
    '| CPU Temp | ' + str(round(hw.get('cpu_temp_c', 0), 1)) + '°C |',
    '| Throttled | ' + (hw.get('cpu_throttled', 'N/A') or 'No') + ' |',
    '| Uptime | ' + str(round(hw.get('uptime_seconds', 0) / 3600, 1)) + ' hours |',
    '| Total Memory | ' + str(hw.get('total_memory_kb', 0) // 1024) + ' MB |',
    '| Available Memory | ' + str(hw.get('available_memory_kb', 0) // 1024) + ' MB |',
    '',
    '---',
    '## Runtime Environment',
    '',
    '| Property | Value |',
    '|----------|-------|',
    '| Python Version | ' + rt.get('python_version', 'N/A').split()[0] + ' |',
    '| Python Executable | `' + rt.get('python_executable', 'N/A') + '` |',
    '| Pyvorin Edge Version | ' + rt.get('pyvorin_edge_version', 'N/A') + ' |',
    '| Hostname | `' + rt.get('hostname', 'N/A') + '` |',
    '| User | `' + rt.get('user', 'N/A') + '` |',
    '| Working Directory | `' + rt.get('working_directory', 'N/A') + '` |',
    '',
    '---',
    '## Benchmark Results',
    '',
    '| Vertical | Readings | Events | Reduction | Latency p50 |',
    '|----------|----------|--------|-----------|-------------|',
]
for r in results:
    if 'error' in r:
        lines.append('| ' + r.get('vertical', '?') + ' | ERROR | — | — | — |')
    else:
        lines.append(
            '| ' + r['vertical'] + ' | ' + "{:,}".format(r['total_readings']) + ' | ' + str(r['events_triggered']) + ' | '
            + "{:.2f}%".format(r['reduction_percent']) + ' | ' + "{:.4f}".format(r['latency_p50_ms']) + ' ms |'
        )

lines += [
    '',
    '---',
    '## Verification Instructions',
    '',
    '### 1. Reproduce on your own hardware',
    '```bash',
    '# Buy a Raspberry Pi 5, install Raspberry Pi OS, then:',
    'curl -fsSL https://pypi.pyvorin.com/pifiles/install-pyvorin-edge.sh | bash',
    'cd ~/pyvorin-edge',
    'bash pifiles/prove-benchmark.sh',
    '```',
    '',
    '### 2. Verify the signature',
    '```bash',
    'python3 -c "',
    'import json, base64',
    'from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey',
    'data = json.load(open(\"proof/benchmark_attested.json\"))',
    'pub = Ed25519PublicKey.from_public_bytes(base64.b64decode(data[\"public_key_b64\"]))',
    'payload = {k: v for k, v in data.items() if k not in (\"signature\", \"public_key_b64\")}',
    'canonical = json.dumps(payload, sort_keys=True, separators=(\",\", \":\")).encode()',
    'pub.verify(base64.b64decode(data[\"signature\"]), canonical)',
    'print(\"SIGNATURE VALID\")',
    '\"',
    '```',
    '',
    '### 3. Verify the integrity hash',
    '```bash',
    'python3 -c "',
    'import json, hashlib',
    'data = json.load(open(\"proof/benchmark_attested.json\"))',
    'expected = data[\"integrity_hash\"]',
    'payload = {k: v for k, v in data.items() if k != \"integrity_hash\"}',
    'canonical = json.dumps(payload, sort_keys=True, separators=(\",\", \":\")).encode()',
    'actual = hashlib.sha256(canonical).hexdigest()',
    'assert actual == expected, f\"Hash mismatch: {actual} != {expected}\"',
    'print(\"INTEGRITY HASH VALID\")',
    '\"',
    '```',
    '',
    '### 4. Verify hardware identity',
    'Compare the serial number and hardware revision in this report against',
    'the physical device. On the Pi, run:',
    '```bash',
    'vcgencmd get_throttled',
    'cat /proc/device-tree/serial-number',
    'cat /proc/device-tree/model',
    '```',
    '',
    '---',
    '*This report was generated automatically by Pyvorin Edge Benchmark Proof Orchestrator.*',
    '*For questions: support@pyvorin.com*',
]

Path(PROOF_DIR + '/benchmark_report.md').write_text('\n'.join(lines))
print('Report generated: ' + PROOF_DIR + '/benchmark_report.md')
PYEOF
log_ok "Markdown report generated."

# ---------------------------------------------------------------------------
# Step 6 — Recording instructions
# ---------------------------------------------------------------------------
log_title "Step 6/6 — Recording instructions"
cat > "${PROOF_DIR}/RECORDING.txt" <<'RECEOF'
# Pyvorin Edge — Benchmark Recording Guide
# This is a text file. Do not execute with bash. Read it instead.

How to record the benchmark run for video proof:

1. Install asciinema (if not already):
   sudo apt install -y asciinema

2. Record the session:
   cd ~/pyvorin-edge
   asciinema rec proof/benchmark.cast

3. Inside the recording, run:
   bash pifiles/prove-benchmark.sh --skip-recording

4. Exit the recording (Ctrl+D or type "exit"):

5. Upload or share the cast file:
   asciinema upload proof/benchmark.cast
   # This gives you a public URL like https://asciinema.org/a/xxxxxx

Alternative (no asciinema):
   script -t proof/typescript.log -a proof/typescript
   bash pifiles/prove-benchmark.sh --skip-recording
   exit
RECEOF

if [[ "${SKIP_RECORDING}" -eq 0 ]]; then
    log_info "Recording guide saved to proof/RECORDING.txt"
fi

# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------
log_title "============================================================"
log_title "                    PROOF PACKAGE READY"
log_title "============================================================"
echo ""
echo "  Directory:   ${PROOF_DIR}"
echo ""
echo "  Files generated:"
ls -lh "${PROOF_DIR}"
echo ""
log_ok "Your benchmark proof package is ready for customer presentation."
echo ""
echo "  Next steps:"
echo "    cat ${PROOF_DIR}/benchmark_report.md"
echo "    cat ${PROOF_DIR}/public_key.b64"
echo "    python -m http.server 8888 --directory ${PROOF_DIR}   # serve for review"
echo ""
echo "  To record a video demo:"
echo "    cat ${PROOF_DIR}/RECORDING.txt   # read recording instructions"
echo ""

exit 0
