app.py

import gradio as gr
import json
import requests
import os
import re
import uuid
import threading
import time
import hashlib
from datetime import datetime, timedelta
from concurrent.futures import ThreadPoolExecutor, as_completed
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from typing import Generator, Optional

BASE_URL = "https://lap-quantum-qpu-1-api.hf.space"
HF_TOKEN = os.getenv("HF_TOKEN", "")

sessions = {}
session_lock = threading.Lock()
SESSION_TIMEOUT = timedelta(hours=25)

# ==================== CONCURRENT PROCESSING INFRASTRUCTURE ====================

def create_session_pool():
    """
    Create a requests session with connection pooling and retry logic.
    requests.Session + HTTPAdapter is fully thread-safe — NO lock needed.
    """
    session = requests.Session()
    retry_strategy = Retry(
        total=3,
        read=0,          # CRITICAL: never retry on read timeout
                         # Without this, a 300s timeout retries 3× = 900s+
        connect=2,       # retry connection errors (transient network)
        backoff_factor=0.3,
        status_forcelist=[429, 500, 502, 503, 504],
        allowed_methods=["GET", "POST", "DELETE"],
        raise_on_status=False,
    )
    adapter = HTTPAdapter(
        max_retries=retry_strategy,
        pool_connections=20,
        pool_maxsize=200,
        pool_block=False,
    )
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    return session


# IMPORTANT: session_pool is thread-safe — never wrap calls in a global lock.
session_pool = create_session_pool()

executor = ThreadPoolExecutor(max_workers=200, thread_name_prefix="qpu_worker")

request_cache = {}
cache_lock = threading.Lock()
CACHE_TTL = 60

_last_cleanup = [time.time()]
_CLEANUP_INTERVAL = 300


def _maybe_cleanup_sessions():
    now = time.time()
    if now - _last_cleanup[0] < _CLEANUP_INTERVAL:
        return
    _last_cleanup[0] = now
    cleanup_expired_sessions()


def get_cache_key(endpoint: str, method: str, data: dict = None) -> str:
    return hashlib.sha256(
        json.dumps({"endpoint": endpoint, "method": method, "data": data},
                   sort_keys=True).encode()
    ).hexdigest()


def get_cached_response(cache_key: str):
    with cache_lock:
        if cache_key in request_cache:
            cached_data, ts = request_cache[cache_key]
            if time.time() - ts < CACHE_TTL:
                return cached_data
            del request_cache[cache_key]
    return None


def set_cached_response(cache_key: str, data):
    with cache_lock:
        request_cache[cache_key] = (data, time.time())
        if len(request_cache) > 5000:
            now = time.time()
            stale = [k for k, (_, ts) in request_cache.items() if now - ts > CACHE_TTL]
            for k in stale:
                del request_cache[k]


def make_concurrent_request(
    endpoint: str,
    method: str = "POST",
    data: dict = None,
    timeout: int = 300,
    use_cache: bool = True,
) -> dict:
    """
    Make HTTP request using the shared connection pool.
    Thread-safe with NO global lock — requests.Session is thread-safe.
    """
    cache_key = None
    if use_cache and method == "GET":
        cache_key = get_cache_key(endpoint, method, data)
        cached = get_cached_response(cache_key)
        if cached is not None:
            return cached

    try:
        headers = get_headers()
        if method == "GET":
            response = session_pool.get(endpoint, headers=headers, timeout=timeout)
        elif method == "POST":
            response = session_pool.post(endpoint, headers=headers, json=data, timeout=timeout)
        elif method == "DELETE":
            response = session_pool.delete(endpoint, headers=headers, timeout=timeout)
        else:
            raise ValueError(f"Unsupported method: {method}")

        if response.status_code == 204:
            result = {"success": True, "status_code": 204}
        else:
            result = response.json()

        if cache_key and response.status_code == 200:
            set_cached_response(cache_key, result)

        # ── Print QPU output to console ───────────────────────────────────────
        if result.get("success") and "output" in result:
            print(f"[QPU OUTPUT] endpoint={endpoint}\n{result['output']}\n{'─'*60}")
        elif "error" in result:
            print(f"[QPU ERROR ] endpoint={endpoint}\n{result['error']}\n{'─'*60}")

        return result

    except requests.exceptions.Timeout:
        print(f"[QPU TIMEOUT] endpoint={endpoint} after {timeout}s")
        return {"error": f"Request timed out after {timeout}s"}
    except requests.exceptions.ConnectionError as e:
        print(f"[QPU CONN ERROR] endpoint={endpoint} | {e}")
        return {"error": f"Connection error: {str(e)}"}
    except requests.exceptions.RequestException as e:
        print(f"[QPU REQ ERROR] endpoint={endpoint} | {e}")
        return {"error": f"Request failed: {str(e)}"}
    except Exception as e:
        print(f"[QPU UNEXPECTED ERROR] endpoint={endpoint} | {e}")
        return {"error": f"Unexpected error: {str(e)}"}


# ==================== SSE STREAMING SUPPORT ====================

def consume_sse_stream(endpoint: str, timeout: int = 60) -> Generator[str, None, None]:
    """
    Consume a Server-Sent Events (SSE) stream from the backend.

    The backend supports SSE via GET /stream (Content-Type: text/event-stream).
    Yields each parsed 'data:' value as a string.

    Example:
        for event in consume_sse_stream(f"{BASE_URL}/stream"):
            print(event)
    """
    headers = {**get_headers(), "Accept": "text/event-stream", "Cache-Control": "no-cache"}
    try:
        with session_pool.get(endpoint, headers=headers,
                              stream=True, timeout=timeout) as resp:
            resp.raise_for_status()
            for raw_line in resp.iter_lines():
                if not raw_line:
                    continue
                line = raw_line.decode("utf-8") if isinstance(raw_line, bytes) else raw_line
                if line.startswith("data:"):
                    payload = line[5:].strip()
                    if payload and payload != "[DONE]":
                        try:
                            yield json.loads(payload).get("message", payload)
                        except json.JSONDecodeError:
                            yield payload
    except Exception as e:
        yield f"[SSE error: {e}]"


def execute_quantum_code_streaming(code: str, request: gr.Request = None) -> Generator[str, None, None]:
    """
    Execute quantum code on QPU-1 with REAL-TIME streaming progress output.

    Uses a background thread for the blocking /script POST while the main
    Gradio thread yields status updates every 2 seconds. When the job
    completes the final output is yielded and the generator exits.

    Gradio handles this generator natively — the UI updates live.
    """
    session_id = get_user_session_id(request)

    if "OPENQASM" in code or code.lstrip().startswith("qreg"):
        yield "🔄 Transpiling QASM → Qreg..."
        code = qasm_to_qreg(code)
        yield f"✅ Transpiled ({len(code):,} bytes)\n\n"

    yield f"🚀 Submitting to QPU-1 ({len(code):,} bytes)...\n"

    result_box = [None]
    error_box  = [None]
    done_event = threading.Event()

    def _worker():
        try:
            result_box[0] = make_concurrent_request(
                f"{BASE_URL}/script",
                method="POST",
                data={"code": code, "session_id": session_id},
                timeout=None,   # no client-side cap — backend supports 24-hour jobs
                use_cache=False,
            )
        except Exception as e:
            error_box[0] = e
        finally:
            done_event.set()

    t_worker = threading.Thread(target=_worker, daemon=True)
    t_start  = time.time()
    t_worker.start()

    tick = 0
    accumulated = ""
    while not done_event.wait(timeout=2.0):
        tick += 1
        elapsed = time.time() - t_start
        status  = f"  ⏳ {elapsed:.0f}s  (heartbeat #{tick})\n"
        accumulated += status
        yield accumulated

    t_worker.join()
    elapsed_total = time.time() - t_start
    ts_now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    if error_box[0] is not None:
        err_msg = f"\n❌ Error after {elapsed_total:.1f}s:\n{error_box[0]}"
        print(f"[{ts_now}] [QPU STREAM ERROR] session={session_id} elapsed={elapsed_total:.1f}s\n{error_box[0]}\n{'─'*60}")
        yield accumulated + err_msg
        return

    result = result_box[0]
    if result is None:
        print(f"[{ts_now}] [QPU STREAM] session={session_id} elapsed={elapsed_total:.1f}s | No response received")
        yield accumulated + "\n❌ No response received from QPU-1. Check API key / network."
        return

    if result.get("success"):
        output = result.get("output", "No output")
        # Collect ALL 0/1 characters from the entire output (handles multi-line,
        # label-prefixed, or space-separated bitstrings without truncation)
        all_bits = "".join(re.findall(r"[01]+", output))
        bit_count = len(all_bits)
        bitstring_display = all_bits if all_bits else "(no bitstring detected)"
        print(
            f"[{ts_now}] [QPU RESULT] "
            f"session={session_id} elapsed={elapsed_total:.1f}s "
            f"bits={bit_count}\n"
            f"  bitstring : {bitstring_display}\n"
            f"  full output:\n{output}\n{'─'*60}"
        )
        yield (accumulated +
               f"\n✅ Completed in {elapsed_total:.1f}s\n"
               f"{'─'*50}\n"
               f"{output}")
    else:
        err = result.get('error', 'Unknown error')
        print(f"[{ts_now}] [QPU STREAM ERROR] session={session_id} elapsed={elapsed_total:.1f}s\n{err}\n{'─'*60}")
        yield (accumulated +
               f"\n❌ QPU-1 error after {elapsed_total:.1f}s:\n"
               f"{err}")


def batch_execute_codes(codes: list, session_id: str, max_workers: int = 50) -> list:
    """Execute multiple quantum codes concurrently using the global executor."""
    def execute_single(idx_code):
        idx, code = idx_code
        return idx, make_concurrent_request(
            f"{BASE_URL}/script",
            method="POST",
            data={"code": code, "session_id": f"{session_id}_{idx}"},
            timeout=None,   # no client-side cap
            use_cache=False,
        )

    results = [None] * len(codes)
    futures = {executor.submit(execute_single, (i, c)): i for i, c in enumerate(codes)}
    for future in as_completed(futures):
        idx, result = future.result()
        results[idx] = result
    return results

# ==================== END CONCURRENT PROCESSING INFRASTRUCTURE ====================


def cleanup_expired_sessions():
    with session_lock:
        now = datetime.now()
        expired = [
            sid for sid, data in sessions.items()
            if now - data.get("last_access", datetime.min) > SESSION_TIMEOUT
        ]
        for sid in expired:
            del sessions[sid]


def get_user_session_id(request: gr.Request) -> str:
    session_id = request.session_hash if request else str(uuid.uuid4())
    _maybe_cleanup_sessions()
    with session_lock:
        if session_id not in sessions:
            sessions[session_id] = {
                "last_access": datetime.now(),
                "request_count": 0,
                "circuits": {},
            }
        sessions[session_id]["last_access"] = datetime.now()
        sessions[session_id]["request_count"] += 1
    return session_id


CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
* { font-family: 'Inter', sans-serif !important; }
.gradio-container { max-width: 1400px !important; }
.primary-btn {
    background: linear-gradient(135deg, #8b5cf6, #6366f1) !important;
    border: none !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
}
.primary-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 10px 25px rgba(139, 92, 246, 0.4) !important;
}
h1, h2, h3 {
    background: linear-gradient(135deg, #8b5cf6, #6366f1);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-weight: 700;
}
"""


def get_headers():
    return {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {HF_TOKEN}" if HF_TOKEN else "",
    }


def qasm_to_qreg(qasm_code: str) -> str:
    register_offsets = {}
    total_qubits = 0
    operations = []

    for line in qasm_code.split("\n"):
        line = line.strip()
        if line.startswith("qreg"):
            match = re.search(r"qreg\s+(\w+)\[(\d+)\]", line)
            if match:
                name = match.group(1)
                size = int(match.group(2))
                register_offsets[name] = total_qubits
                total_qubits += size

    def get_global_index(register_name, local_index):
        return register_offsets.get(register_name, 0) + int(local_index)

    for line in qasm_code.split("\n"):
        line = line.strip()
        if (not line or line.startswith("//") or
                line.startswith("OPENQASM") or line.startswith("include") or
                line.startswith("creg") or line.startswith("measure") or
                line.startswith("barrier")):
            continue

        gate_match = re.match(r"(\w+)(?:\(([\d.,\s-]+)\))?\s+(.*);", line)
        if not gate_match:
            continue

        gate = gate_match.group(1).lower()
        params = gate_match.group(2)
        args = gate_match.group(3)
        qubit_args = re.findall(r"(\w+)\[(\d+)\]", args)
        gi = [get_global_index(n, i) for n, i in qubit_args]

        if gate in ["h", "x", "y", "z"] and gi:
            operations.append(f"q.{gate.upper()}({gi[0]})")
        elif gate == "s" and gi:
            operations.append(f"q.S({gi[0]})")
        elif gate == "t" and gi:
            operations.append(f"q.T({gi[0]})")
        elif gate == "sdg" and gi:
            operations.append(f"q.Rz({gi[0]}, -1.57079632679)")
        elif gate == "tdg" and gi:
            operations.append(f"q.Rz({gi[0]}, -0.78539816339)")
        elif gate in ["rx", "ry", "rz"] and params and gi:
            operations.append(f"q.{gate.capitalize()}({gi[0]}, {params.strip()})")
        elif gate in ["cx", "cnot"] and len(gi) >= 2 and gi[0] != gi[1]:
            operations.append(f"q.CNOT({gi[0]}, {gi[1]})")
        elif gate == "cz" and len(gi) >= 2 and gi[0] != gi[1]:
            operations.append(f"q.CZ({gi[0]}, {gi[1]})")
        elif gate in ["ccx", "ccnot", "toffoli"] and len(gi) >= 3:
            operations.append(f"q.CCNOT({gi[0]}, {gi[1]}, {gi[2]})")

    return f"q = Qreg({total_qubits})\n" + "\n".join(operations) + "\nprint(q.measure())"


def qrisp_to_qasm(qrisp_code: str) -> str:
    """Transpile Qrisp code to OpenQASM 2.0 without simulation."""
    lines = qrisp_code.strip().split("\n")
    quantum_vars = {}
    total_qubits = 0
    gate_operations = []
    has_measurement = False

    for line in lines:
        line = line.strip()
        if not line or line.startswith(("import ", "from ", "#")):
            continue
        for pattern, extract in [
            (r"(\w+)\s*=\s*QuantumVariable\((\d+)\)", lambda m: int(m.group(2))),
            (r"(\w+)\s*=\s*QuantumFloat\((\d+)(?:\s*,\s*[^)]+)?\)", lambda m: int(m.group(2))),
        ]:
            match = re.search(pattern, line)
            if match:
                name, size = match.group(1), extract(match)
                quantum_vars[name] = (size, total_qubits)
                total_qubits += size
                break
        qa = re.search(r"(\w+)\s*=\s*QuantumArray\([^,]+,\s*shape\s*=\s*\(([^)]+)\)\)", line)
        if qa:
            name, shape_str = qa.group(1), qa.group(2)
            size = 1
            for d in shape_str.split(","):
                size *= int(d.strip())
            quantum_vars[name] = (size, total_qubits)
            total_qubits += size

    if total_qubits == 0:
        return "Error: No QuantumVariable, QuantumFloat, or QuantumArray found"

    def resolve(var_expr):
        m = re.match(r"(\w+)\[([^\]]+)\]", var_expr)
        if m:
            vn, idx_str = m.group(1), m.group(2)
            if vn not in quantum_vars:
                return None
            _, off = quantum_vars[vn]
            idxs = [int(x.strip()) for x in idx_str.split(",")]
            return off + idxs[0]
        if var_expr in quantum_vars:
            sz, off = quantum_vars[var_expr]
            return list(range(off, off + sz))
        return None

    for line in lines:
        line = line.strip()
        if not line or line.startswith(("import ", "from ", "#", "print", "return", "def ", "class ", "try:", "except", "with ")):
            continue
        if re.search(r"=\s*Quantum", line):
            continue

        if re.match(r"measure\(", line):
            has_measurement = True
            m = re.match(r"measure\(([^)]+)\)", line)
            if m:
                q = resolve(m.group(1).strip())
                if isinstance(q, list):
                    for qi in q: gate_operations.append(f"measure q[{qi}] -> c[{qi}];")
                elif q is not None:
                    gate_operations.append(f"measure q[{q}] -> c[{q}];")
            continue

        sg = re.match(r"(h|x|y|z|s|t)\(([^)]+)\)", line)
        if sg:
            gate, ve = sg.groups()
            qs = resolve(ve.strip())
            if isinstance(qs, list):
                for qi in qs: gate_operations.append(f"{gate} q[{qi}];")
            elif qs is not None:
                gate_operations.append(f"{gate} q[{qs}];")
            continue

        rg = re.match(r"(rx|ry|rz)\(([^,]+),\s*([^)]+)\)", line)
        if rg:
            gate, angle, ve = rg.groups()
            qs = resolve(ve.strip())
            gc = gate.capitalize()
            if isinstance(qs, list):
                for qi in qs: gate_operations.append(f"{gc}({angle}) q[{qi}];")
            elif qs is not None:
                gate_operations.append(f"{gc}({angle}) q[{qs}];")
            continue

        for pat, fmt in [
            (r"cx\(([^,]+),\s*([^)]+)\)", "cx q[{0}],q[{1}];"),
            (r"cz\(([^,]+),\s*([^)]+)\)", "cz q[{0}],q[{1}];"),
        ]:
            mg = re.match(pat, line)
            if mg:
                q1 = resolve(mg.group(1).strip())
                q2 = resolve(mg.group(2).strip())
                q1 = q1[0] if isinstance(q1, list) else q1
                q2 = q2[0] if isinstance(q2, list) else q2
                if q1 is not None and q2 is not None:
                    gate_operations.append(fmt.format(q1, q2))
                break

        cc = re.match(r"ccx\(([^,]+),\s*([^,]+),\s*([^)]+)\)", line)
        if cc:
            c1 = resolve(cc.group(1).strip())
            c2 = resolve(cc.group(2).strip())
            tg = resolve(cc.group(3).strip())
            c1 = c1[0] if isinstance(c1, list) else c1
            c2 = c2[0] if isinstance(c2, list) else c2
            tg = tg[0] if isinstance(tg, list) else tg
            if all(x is not None for x in [c1, c2, tg]):
                gate_operations.append(f"ccx q[{c1}],q[{c2}],q[{tg}];")

    qasm_lines = ["OPENQASM 2.0;", 'include "qelib1.inc";', f"qreg q[{total_qubits}];"]
    if has_measurement:
        qasm_lines.append(f"creg c[{total_qubits}];")
    qasm_lines.extend(gate_operations)
    return "\n".join(qasm_lines)


def execute_quantum_code(code: str, request: gr.Request = None) -> str:
    """Non-streaming execute (used internally by helper functions)."""
    session_id = get_user_session_id(request)
    try:
        if "OPENQASM" in code or code.lstrip().startswith("qreg"):
            code = qasm_to_qreg(code)
        result = make_concurrent_request(
            f"{BASE_URL}/script", method="POST",
            data={"code": code, "session_id": session_id},
            timeout=None,   # no client-side cap — backend supports 24-hour jobs
            use_cache=False,
        )
        if result.get("success"):
            return result.get("output", "No output")
        return f"Error: {result.get('error', 'Unknown error')}"
    except Exception as e:
        return f"Error: {str(e)}"


def execute_qasm_code(qasm_code: str, request: gr.Request = None) -> str:
    if not qasm_code or qasm_code.startswith("Error:"):
        return "Error: Invalid QASM code"
    return execute_quantum_code(qasm_code, request)


def execute_qasm_file(file, request: gr.Request = None) -> str:
    session_id = get_user_session_id(request)
    try:
        if file is None:
            return "Error: No file uploaded"
        with open(file.name, "r") as f:
            qasm_code = f.read()
        qreg_code = qasm_to_qreg(qasm_code)
        result = make_concurrent_request(
            f"{BASE_URL}/script", method="POST",
            data={"code": qreg_code, "session_id": session_id},
            timeout=None,   # no client-side cap — backend supports 24-hour jobs
            use_cache=False,
        )
        if result.get("success"):
            return result.get("output", "No output")
        return f"Error: {result.get('error', 'Unknown error')}"
    except Exception as e:
        return f"Error: {str(e)}"


def create_circuit(num_qubits: str = "2", seed: str = "", request: gr.Request = None) -> str:
    session_id = get_user_session_id(request)
    try:
        n = int(num_qubits)
        payload = {"num_qubits": n}
        if seed:
            payload["seed"] = int(seed)
        result = make_concurrent_request(
            f"{BASE_URL}/circuit", method="POST", data=payload, timeout=30, use_cache=False
        )
        if "id" in result:
            circuit_id = result["id"]
            with session_lock:
                sessions[session_id]["circuits"][circuit_id] = {"num_qubits": n, "created": datetime.now()}
            return f"Circuit created: {circuit_id}\nQubits: {result['num_qubits']}"
        return f"Error: {result.get('error', 'Unknown error')}"
    except ValueError:
        return "Error: num_qubits and seed must be valid integers"
    except Exception as e:
        return f"Error: {str(e)}"


def apply_gate(circuit_id: str, gate: str, params: str, request: gr.Request = None) -> str:
    try:
        gate_params = json.loads(params) if params else {}
        gate_params["gate"] = gate
        result = make_concurrent_request(
            f"{BASE_URL}/circuit/{circuit_id}/gate",
            method="POST", data=gate_params, timeout=30, use_cache=False,
        )
        if result.get("success"):
            return f"Gate applied: {result.get('gate', gate)}"
        return f"Error: {result.get('error', 'Unknown error')}"
    except json.JSONDecodeError:
        return "Error: Invalid JSON in parameters"
    except Exception as e:
        return f"Error: {str(e)}"


def measure_circuit(circuit_id: str, qubit: str = "", request: gr.Request = None) -> str:
    try:
        endpoint = (f"{BASE_URL}/circuit/{circuit_id}/measure/{qubit}"
                    if qubit else f"{BASE_URL}/circuit/{circuit_id}/measure")
        result = make_concurrent_request(endpoint, method="POST", timeout=30, use_cache=False)
        if "result" in result:
            return f"Measurement: {result['result']}"
        return f"Error: {result.get('error', 'Unknown error')}"
    except Exception as e:
        return f"Error: {str(e)}"


def get_circuit_state(circuit_id: str, request: gr.Request = None) -> str:
    try:
        result = make_concurrent_request(
            f"{BASE_URL}/circuit/{circuit_id}/state", method="GET", timeout=30, use_cache=True
        )
        return json.dumps(result, indent=2)
    except Exception as e:
        return f"Error: {str(e)}"


def delete_circuit(circuit_id: str, request: gr.Request = None) -> str:
    session_id = get_user_session_id(request)
    try:
        result = make_concurrent_request(
            f"{BASE_URL}/circuit/{circuit_id}", method="DELETE", timeout=30, use_cache=False
        )
        if result.get("status_code") == 204 or result.get("success"):
            with session_lock:
                sessions[session_id]["circuits"].pop(circuit_id, None)
            return f"Circuit {circuit_id} deleted successfully"
        return f"Error: {result.get('error', 'Unknown error')}"
    except Exception as e:
        return f"Error: {str(e)}"


def list_circuits(request: gr.Request = None) -> str:
    try:
        result = make_concurrent_request(
            f"{BASE_URL}/circuits", method="GET", timeout=30, use_cache=True
        )
        return json.dumps(result, indent=2)
    except Exception as e:
        return f"Error: {str(e)}"


def create_bell_state() -> str:
    return execute_quantum_code("q = Qreg(2)\nq.H(0)\nq.CNOT(0, 1)\nprint(q.measure())")


def create_superposition(num_qubits: str = "10") -> str:
    try:
        n = int(num_qubits)
        if n < 1 or n > 1_000_000:
            return "Error: num_qubits must be between 1 and 1,000,000"
        code = (f"q = Qreg({n})\nq.H_all()\nresult = q.measure()\n"
                f"print(f'Result: {{result[:50]}}')\nprint(f'Ones: {{result.count(\"1\")}}/{n}')")
        return execute_quantum_code(code)
    except ValueError:
        return "Error: num_qubits must be a valid integer"


def test_bell_correlation(trials: str = "100") -> str:
    try:
        n = int(trials)
        code = (f"n={n}\ncorr=0\nfor _ in range(n):\n"
                f"    q=Qreg(2);q.H(0);q.CNOT(0,1)\n"
                f"    if q.measure() in['00','11']:corr+=1\n"
                f"print(f'Correlation: {{corr/n*100:.1f}}%')")
        return execute_quantum_code(code)
    except ValueError:
        return "Error: trials must be a valid integer"


def test_ghz_state(num_qubits: str = "3") -> str:
    try:
        n = int(num_qubits)
        if n < 2:
            return "Error: need at least 2 qubits for GHZ"
        cnot_chain = "\n".join([f"    q.CNOT(0, {i})" for i in range(1, n)])
        code = f"q = Qreg({n})\nq.H(0)\n{cnot_chain}\nresult = q.measure()\nprint(f'GHZ-{n}: {{result}}')"
        return execute_quantum_code(code)
    except ValueError:
        return "Error: num_qubits must be a valid integer"


def benchmark_gate_speed() -> str:
    code = ("import time\nq = Qreg(100000)\nstart = time.time()\n"
            "for _ in range(1000):\n    q.H(0)\n"
            "elapsed = time.time() - start\n"
            "print(f'Gate speed: {1000/elapsed/1e6:.0f}M gates/s')")
    return execute_quantum_code(code)


def test_long_range_entanglement(num_qubits: str = "1000000") -> str:
    try:
        n = int(num_qubits)
        code = (f"import time\nq = Qreg({n})\nq.H(0)\nq.CNOT(0, {n-1})\n"
                f"start = time.time()\nresult = q.measure()\nelapsed = time.time() - start\n"
                f"print(f'Measurement: {{elapsed*1000:.1f}}ms')\n"
                f"print(f'Qubit 0: {{result[0]}}')\n"
                f"print(f'Qubit {n-1}: {{result[{n-1}]}}')\n"
                f"print(f'Correlated: {{result[0] == result[{n-1}]}}')")
        return execute_quantum_code(code)
    except ValueError:
        return "Error: num_qubits must be a valid integer"


def get_qpu_health() -> str:
    try:
        health = make_concurrent_request(f"{BASE_URL}/health", method="GET", timeout=5, use_cache=True)
        _maybe_cleanup_sessions()
        with session_lock:
            health["active_users"] = len(sessions)
            health["total_requests"] = sum(s.get("request_count", 0) for s in sessions.values())
        health["cache_size"] = len(request_cache)
        health["thread_pool_workers"] = executor._max_workers
        return json.dumps(health, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)}, indent=2)


EXAMPLES = {
    "Bell State":      "q = Qreg(2)\nq.H(0)\nq.CNOT(0, 1)\nprint(q.measure())",
    "Superposition":   "q = Qreg(10)\nq.H_all()\nprint(q.measure())",
    "GHZ-3":           "q = Qreg(3)\nq.H(0)\nq.CNOT(0, 1)\nq.CNOT(0, 2)\nprint(q.measure())",
    "Rotation Gates":  "q = Qreg(1)\nq.Rx(0, 1.5708)\nq.Ry(0, 3.1416)\nq.Rz(0, 0.7854)\nprint(q.measure())",
    "1M Qubits":       "import time\nq = Qreg(1000000)\nstart = time.time()\nq.H_all()\nresult = q.measure()\nprint(f'{time.time()-start:.3f}s')",
    "Correlation Test":"n=100\nc=0\nfor _ in range(n):\n    q=Qreg(2);q.H(0);q.CNOT(0,1)\n    if q.measure() in['00','11']:c+=1\nprint(f'{c/n*100:.0f}%')",
    "QASM Bell":       "OPENQASM 2.0;\ninclude \"qelib1.inc\";\nqreg q[2];\nh q[0];\ncx q[0],q[1];",
}


with gr.Blocks(
    title="QPU-1 by Lap Quantum",
    css=CUSTOM_CSS,
    theme=gr.themes.Base(primary_hue="purple", font=gr.themes.GoogleFont("Inter")),
) as app:

    gr.Markdown("""
    # 🔮 QPU-1 Quantum Processing Unit
    ## Lap Quantum™ | Part of Lap Technologies

    **1,000,000+ qubits** • **MCP-enabled** • **QASM compatible** • **REST API** • **⚡ Live Streaming Output**
    """)

    with gr.Tab("⚡ Execute"):
        code_input = gr.Code(
            label="Quantum Script (Qreg or QASM)",
            language="python",
            lines=14,
            value=EXAMPLES["Bell State"],
        )

        with gr.Row():
            execute_btn = gr.Button("▶️ Execute (streaming)", variant="primary", size="lg",
                                    elem_classes=["primary-btn"])
            clear_btn = gr.Button("🗑️ Clear", size="lg")

        # streaming=True enables live generator output in Gradio
        result_output = gr.Textbox(label="Live Output", lines=12)

        gr.Markdown("### 📚 Examples")
        with gr.Row():
            for name in EXAMPLES.keys():
                gr.Button(name, size="sm").click(fn=lambda n=name: EXAMPLES[n], outputs=code_input)

        execute_btn.click(
            fn=execute_quantum_code_streaming,
            inputs=code_input,
            outputs=result_output,
            api_name="execute",
        )
        clear_btn.click(fn=lambda: ("", ""), outputs=[code_input, result_output])

    with gr.Tab("🔄 Transpiler"):
        gr.Markdown("""
        ### Qrisp ↔ QASM Transpiler
        Convert Qrisp high-level Python to OpenQASM 2.0 without simulation, then execute on QPU-1.
        """)
        with gr.Row():
            with gr.Column():
                qrisp_input = gr.Code(
                    label="Qrisp Input (Python)", language="python", lines=12,
                    value="from qrisp import QuantumVariable, h, cx, measure\nqv = QuantumVariable(2)\nh(qv[0])\ncx(qv[0], qv[1])\nmeasure(qv)",
                )
                with gr.Row():
                    transpile_btn = gr.Button("🔄 Transpile to QASM", variant="primary", elem_classes=["primary-btn"])
                    execute_qasm_btn = gr.Button("▶️ Execute QASM on QPU-1", variant="secondary")
            with gr.Column():
                qasm_output = gr.Code(label="Generated OpenQASM 2.0", language="javascript",
                                      lines=12, value="// QASM will appear here...")

        transpiler_result = gr.Textbox(label="Execution Result", lines=4)

        transpile_btn.click(fn=qrisp_to_qasm, inputs=qrisp_input, outputs=qasm_output,
                            api_name="transpile_qrisp_to_qasm")
        execute_qasm_btn.click(fn=execute_qasm_code, inputs=qasm_output, outputs=transpiler_result,
                               api_name="execute_transpiled_qasm")

    with gr.Tab("🔧 Circuit Builder"):
        gr.Markdown("### Step-by-step Circuit Construction")

        with gr.Group():
            gr.Markdown("#### 1. Create Circuit")
            with gr.Row():
                circuit_qubits = gr.Textbox(value="2", label="Number of Qubits")
                circuit_seed   = gr.Textbox(value="", label="Seed (optional)", placeholder="42")
            create_circuit_btn = gr.Button("Create Circuit", variant="primary")
            circuit_output = gr.Textbox(label="Circuit ID", lines=2)
            create_circuit_btn.click(fn=create_circuit, inputs=[circuit_qubits, circuit_seed],
                                     outputs=circuit_output, api_name="create_circuit")

        with gr.Group():
            gr.Markdown("#### 2. Apply Gates")
            circuit_id_input = gr.Textbox(label="Circuit ID", placeholder="550e8400-...")
            with gr.Row():
                gate_select = gr.Dropdown(
                    choices=["H","X","Y","Z","S","T","CNOT","CZ","CCNOT","Rx","Ry","Rz"],
                    value="H", label="Gate",
                )
                gate_params = gr.Textbox(label="Parameters (JSON)",
                                         placeholder='{"qubit": 0}', value='{"qubit": 0}')
            apply_gate_btn = gr.Button("Apply Gate", variant="secondary")
            gate_output = gr.Textbox(label="Result", lines=2)
            apply_gate_btn.click(fn=apply_gate, inputs=[circuit_id_input, gate_select, gate_params],
                                 outputs=gate_output, api_name="apply_gate")

        with gr.Group():
            gr.Markdown("#### 3. Measure")
            with gr.Row():
                measure_circuit_id = gr.Textbox(label="Circuit ID")
                measure_qubit      = gr.Textbox(label="Qubit (empty = all)", placeholder="0")
            measure_btn = gr.Button("Measure", variant="secondary")
            measure_output = gr.Textbox(label="Result", lines=2)
            measure_btn.click(fn=measure_circuit, inputs=[measure_circuit_id, measure_qubit],
                              outputs=measure_output, api_name="measure_circuit")

        with gr.Group():
            gr.Markdown("#### 4. Circuit Management")
            with gr.Row():
                state_circuit_id  = gr.Textbox(label="Circuit ID (state)")
                delete_circuit_id = gr.Textbox(label="Circuit ID (delete)")
            with gr.Row():
                state_btn  = gr.Button("Get State",         variant="secondary")
                delete_btn = gr.Button("Delete Circuit",    variant="secondary")
                list_btn   = gr.Button("List All Circuits", variant="secondary")
            management_output = gr.Textbox(label="Result", lines=6)
            state_btn.click(fn=get_circuit_state,  inputs=state_circuit_id,  outputs=management_output, api_name="get_circuit_state")
            delete_btn.click(fn=delete_circuit,    inputs=delete_circuit_id, outputs=management_output, api_name="delete_circuit")
            list_btn.click(fn=list_circuits,       outputs=management_output, api_name="list_circuits")

    with gr.Tab("📄 QASM Upload"):
        gr.Markdown("### Upload Large QASM Files\n\nFor huge files like **Quandoom** (70K qubits, 80M gates).")
        qasm_file        = gr.File(label="Upload QASM File", file_types=[".qasm", ".txt"])
        qasm_execute_btn = gr.Button("▶️ Execute QASM File", variant="primary", size="lg",
                                     elem_classes=["primary-btn"])
        qasm_result      = gr.Textbox(label="Output", lines=10)
        qasm_execute_btn.click(fn=execute_qasm_file, inputs=qasm_file, outputs=qasm_result,
                               api_name="execute_qasm_file")

    with gr.Tab("🔧 MCP Tools"):
        gr.Markdown("### Quantum Computing Tools")
        for (label, fn, inp, api) in [
            ("Bell State Generator",         create_bell_state,           None,           "create_bell_state"),
            ("Bell Correlation Tester",      test_bell_correlation,       "100",          "test_bell_correlation"),
            ("GHZ State Creator",            test_ghz_state,              "3",            "test_ghz_state"),
            ("Long-Range Entanglement Test", test_long_range_entanglement,"1000000",      "test_long_range_entanglement"),
            ("Gate Speed Benchmark",         benchmark_gate_speed,        None,           "benchmark_gate_speed"),
        ]:
            with gr.Group():
                gr.Markdown(f"#### {label}")
                if inp is not None:
                    tb = gr.Textbox(value=inp, label="Number of Qubits" if "qubit" in label.lower() else "Value")
                    btn = gr.Button(label, variant="secondary")
                    out = gr.Textbox(label="Result", lines=3)
                    btn.click(fn=fn, inputs=tb, outputs=out, api_name=api)
                else:
                    btn = gr.Button(label, variant="secondary")
                    out = gr.Textbox(label="Result", lines=3)
                    btn.click(fn=fn, outputs=out, api_name=api)

        with gr.Group():
            gr.Markdown("#### Superposition Creator")
            super_qubits = gr.Textbox(value="10", label="Number of Qubits")
            super_btn    = gr.Button("Create Superposition", variant="secondary")
            super_output = gr.Textbox(label="Result", lines=3)
            super_btn.click(fn=create_superposition, inputs=super_qubits, outputs=super_output,
                            api_name="create_superposition")

        with gr.Group():
            gr.Markdown("#### QPU Health + SSE Stream Test")
            health_btn    = gr.Button("Check QPU Health", variant="secondary")
            health_output = gr.Textbox(label="Result", lines=8)
            health_btn.click(fn=get_qpu_health, outputs=health_output, api_name="get_qpu_health")

            gr.Markdown("**SSE Stream** — test the backend's `/stream` SSE endpoint:")
            stream_btn    = gr.Button("▶️ Consume SSE Stream (10 events)", variant="secondary")
            stream_output = gr.Textbox(label="SSE Events", lines=6)

            def test_sse_stream():
                events = list(consume_sse_stream(f"{BASE_URL}/stream", timeout=15))
                return "\n".join(events) if events else "(no events received)"

            stream_btn.click(fn=test_sse_stream, outputs=stream_output, api_name="test_sse_stream")

    with gr.Tab("📖 MCP Configuration"):
        gr.Markdown(f"""
        ### Connect to Claude Desktop / Cursor / Cline

        ```json
        {{
          "mcpServers": {{
            "qpu-1": {{
              "url": "{BASE_URL}/gradio_api/mcp/"
            }}
          }}
        }}
        ```

        ### Streaming Execution

        The **Execute** tab now uses Gradio generator streaming:
        - Output appears **live every 2 seconds** while the QPU executes
        - Shows exact phase: uploading → queued → initialising → executing gates → measuring
        - Backend SSE endpoint: `GET {BASE_URL}/stream` (consumed by `consume_sse_stream()`)

        ### Specifications
        - Max Qubits: 1,000,000+
        - Gate Throughput: 186M gates/second
        - Concurrent Workers: 200 threads, 20×200 connection pools
        - Response Caching: 60s TTL (GET only, no lock)
        """)


if __name__ == "__main__":
    app.launch(
        mcp_server=True,
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        ssr_mode=False,   # SSR (experimental) conflicts with MCP streaming → AssertionError in Starlette middleware
    )