import time
import json
import hashlib
import platform
import os
from datetime import datetime

import numpy as np
import gradio as gr

APP_TITLE = "Coherent Compute Engine"

# -----------------------------
# Core coherent update (NumPy)
# -----------------------------
def coherent_step(Psi, E, L):
    phase = 0.997 * Psi + 0.003 * E
    drive = np.tanh(phase)
    Psi = 0.999 * Psi + 0.001 * drive
    E   = 0.995 * E   + 0.004 * Psi
    L   = 0.998 * L   + 0.001 * (Psi * E)
    return Psi, E, L


# -----------------------------
# Metrics
# -----------------------------
def compute_coherence(a, b):
    num = float(np.dot(a, b))
    den = float(np.linalg.norm(a) * np.linalg.norm(b) + 1e-9)
    return float(abs(num / den))


# -----------------------------
# Optional baseline (context only, safety-capped)
# -----------------------------
def python_loop_baseline(n, steps, cap=50_000):
    n_cap = min(int(n), int(cap))
    x = [0.5] * n_cap
    t0 = time.time()
    for _ in range(int(steps)):
        for i in range(n_cap):
            x[i] = 0.999 * x[i] + 0.001
    elapsed = time.time() - t0
    items = n_cap * int(steps)
    return items / elapsed / 1e9, n_cap


# -----------------------------
# Receipt hashing
# IMPORTANT: sha256 is computed over receipt JSON WITHOUT the sha256 field
# -----------------------------
def compute_receipt_sha256(receipt_obj: dict) -> str:
    obj = dict(receipt_obj)
    obj.pop("sha256", None)
    payload = json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    return hashlib.sha256(payload).hexdigest()


# -----------------------------
# Main engine
# -----------------------------
def run_engine(n_osc, steps, include_baselines):
    n = int(n_osc)
    steps = int(steps)

    Psi = np.random.rand(n).astype(np.float32)
    E   = np.random.rand(n).astype(np.float32)
    L   = np.random.rand(n).astype(np.float32)

    # sample slice for coherence proxy (keeps runtime sane)
    sample = min(n, 200_000)
    Psi_start = Psi[:sample].copy()

    t0 = time.time()
    for _ in range(steps):
        Psi, E, L = coherent_step(Psi, E, L)
    elapsed = time.time() - t0

    Psi_end = Psi[:sample].copy()

    coherence = compute_coherence(Psi_start, Psi_end)
    mean_energy = float(np.mean(E))

    items = n * steps
    throughput_items = items / elapsed
    throughput_b = throughput_items / 1e9

    result = {
        "Throughput (B items/s)": f"{throughput_b:.3f}",
        "Throughput (items/s)": f"{int(throughput_items):,}",
        "Coherence (|C|)": f"{coherence:.5f}",
        "Mean Energy": f"{mean_energy:.5f}",
        "Elapsed Time (s)": f"{elapsed:.2f}",
        "Oscillators": f"{n:,}",
        "Steps": f"{steps}",
        "Engine": "numpy",
        "CPU": platform.processor() or "",
        "Cores Available": os.cpu_count() or 0,
    }

    if include_baselines:
        py_b, py_ncap = python_loop_baseline(n, steps)
        result.update({
            "Baseline: numpy (B items/s)": f"{throughput_b:.3f}",
            "Baseline: python_loop (B items/s)": f"{py_b:.3f}",
            "Baseline: python_loop items measured": f"{py_ncap:,}",
            "Speedup vs python_loop (x)": f"{(throughput_b / max(py_b, 1e-12)):.1f}",
            "Speedup vs numpy (x)": f"{(throughput_b / max(throughput_b, 1e-12)):.2f}",
            "Note": "Baselines are context-only; python loop is safety-capped and measured live."
        })

    receipt = {
        "timestamp_utc": datetime.utcnow().isoformat() + "Z",
        "app": APP_TITLE,
        "results": result,
        "platform": platform.platform(),
        "python": platform.python_version(),
    }

    receipt["sha256"] = compute_receipt_sha256(receipt)

    os.makedirs("receipts", exist_ok=True)
    fname = f"receipts/receipt_{datetime.utcnow().strftime('%Y-%m-%dT%H-%M-%S')}.json"
    with open(fname, "w", encoding="utf-8") as f:
        json.dump(receipt, f, indent=2, ensure_ascii=False)

    result["Receipt SHA-256 (in file)"] = receipt["sha256"]

    # Return pretty results + receipt file
    return json.dumps(result, indent=2), fname


# -----------------------------
# Receipt verifier
# -----------------------------
def verify_receipt(uploaded_file):
    if uploaded_file is None:
        return "No file uploaded.", ""

    path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
    try:
        with open(path, "r", encoding="utf-8") as f:
            receipt = json.load(f)
    except Exception as e:
        return f"Invalid JSON file: {e}", ""

    claimed = str(receipt.get("sha256", "")).strip()
    computed = compute_receipt_sha256(receipt)

    ok = (claimed != "" and claimed.lower() == computed.lower())

    verdict = {
        "verified": bool(ok),
        "claimed_sha256": claimed,
        "computed_sha256": computed,
        "note": "Verified means: sha256(receipt_without_sha256_field) matches the sha256 stored in the receipt."
    }

    banner = "✅ Receipt verified (SHA-256 match)" if ok else "❌ Receipt FAILED verification (SHA-256 mismatch)"
    return banner, json.dumps(verdict, indent=2)


# -----------------------------
# UI
# -----------------------------
with gr.Blocks(title=APP_TITLE) as demo:
    gr.Markdown(
        f"""
# {APP_TITLE}

Everything you see below is computed **right now**, on this machine.

**What an “item” is**  
• One coherent state update of **[Ψ, E, L]** per oscillator per step

**What you get**  
• Real throughput (items/sec)  
• Stability proxy (|C|)  
• Energy behaviour  
• A downloadable **receipt with SHA-256** (verification-first)  
• A built-in **receipt verifier** (upload → verify → done)

No precomputed results. No estimates.
"""
    )

    with gr.Tabs():
        with gr.Tab("Run"):
            n_slider = gr.Slider(250_000, 8_000_000, value=6_400_000, step=250_000,
                                 label="Number of Oscillators")
            steps_slider = gr.Slider(50, 1_000, value=650, step=25,
                                     label="Simulation Steps")

            baseline_toggle = gr.Checkbox(
                label="Include baselines (context only)",
                value=True
            )

            run_btn = gr.Button("Run Engine")

            output = gr.Code(label="Results", language="json")
            receipt_file = gr.File(label="Receipt (download)")

            run_btn.click(
                fn=run_engine,
                inputs=[n_slider, steps_slider, baseline_toggle],
                outputs=[output, receipt_file]
            )

        with gr.Tab("Verify receipt"):
            gr.Markdown(
                """
Upload a receipt JSON generated by this Space.

This verifier recomputes the SHA-256 (over the receipt JSON **without** the `sha256` field) and checks it matches the stored value.
"""
            )
            up = gr.File(label="Upload receipt (.json)")
            verify_btn = gr.Button("Verify")
            banner = gr.Textbox(label="Verification result", interactive=False)
            detail = gr.Code(label="Details", language="json")
            verify_btn.click(fn=verify_receipt, inputs=[up], outputs=[banner, detail])

demo.launch()