import os
import time

import gradio as gr
from starlette.requests import Request
from starlette.responses import JSONResponse, Response

DOWNLOAD_SIZE_BYTES = 10 * 1024 * 1024
_RANDOM_BLOB = os.urandom(DOWNLOAD_SIZE_BYTES)

NO_CACHE_HEADERS = {
    "Cache-Control": "no-store, no-cache, must-revalidate, max-age=0",
    "Pragma": "no-cache",
    "Expires": "0",
}


def register_backend_routes(blocks: gr.Blocks) -> None:
    """Attach REST endpoints to the underlying Gradio server once it exists."""
    server_app = getattr(blocks, "server_app", None)
    if server_app is None:
        return

    if getattr(server_app.state, "speedtest_routes_installed", False):
        return

    server_app.state.speedtest_routes_installed = True

    async def ping_test() -> Response:
        headers = {**NO_CACHE_HEADERS, "Content-Length": "4", "X-Server-Timestamp": str(time.time())}
        return Response(content="pong", media_type="text/plain", headers=headers)

    async def download_test() -> Response:
        headers = {
            **NO_CACHE_HEADERS,
            "Content-Length": str(DOWNLOAD_SIZE_BYTES),
            "Content-Disposition": f'attachment; filename="speedtest_{DOWNLOAD_SIZE_BYTES}.bin"',
        }
        return Response(content=_RANDOM_BLOB, media_type="application/octet-stream", headers=headers)

    async def upload_test(request: Request) -> JSONResponse:
        payload = await request.body()
        headers = dict(NO_CACHE_HEADERS)
        return JSONResponse({"received_bytes": len(payload)}, headers=headers)

    server_app.add_api_route("/ping_test", ping_test, methods=["GET"])
    server_app.add_api_route("/download_test", download_test, methods=["GET"])
    server_app.add_api_route("/upload_test", upload_test, methods=["POST"])


with gr.Blocks(fill_height=True) as demo:
    gr.HTML(value="""
<style>
  .speedtest-container {
    font-family: 'Segoe UI', Arial, sans-serif;
    max-width: 640px;
    margin: 0 auto;
    padding: 32px 24px 40px;
    background: #ffffff;
    border-radius: 16px;
    box-shadow: 0 12px 24px rgba(15, 23, 42, 0.08);
  }
  .speedtest-heading { margin: 0 0 8px; font-size: 1.9rem; }
  .speedtest-lead { color: #475569; line-height: 1.55; }
  .speedtest-button {
    appearance: none;
    border: none;
    border-radius: 999px;
    padding: 14px 28px;
    font-size: 1rem;
    font-weight: 600;
    color: #fff;
    background: linear-gradient(135deg, #2563eb, #9333ea);
    cursor: pointer;
    transition: transform 0.2s ease, box-shadow 0.2s ease;
  }
  .speedtest-button:disabled { opacity: 0.5; cursor: progress; transform: none; box-shadow: none; }
  .speedtest-button:not(:disabled):hover { transform: translateY(-1px); box-shadow: 0 10px 20px rgba(79, 70, 229, 0.25); }
  .speedtest-results { margin-top: 32px; display: none; }
  .speedtest-grid {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(180px, 1fr));
    gap: 20px;
  }
  .metric-card {
    background: #f8fafc;
    border-radius: 12px;
    padding: 18px 20px;
    text-align: center;
    border: 1px solid #e2e8f0;
  }
  .metric-label { display: block; font-weight: 600; color: #334155; margin-bottom: 6px; }
  .metric-value { font-size: 1.6rem; font-variant-numeric: tabular-nums; color: #111827; }
  .metric-subtext { margin-top: 10px; color: #64748b; font-size: 0.9rem; }
  .speedtest-log {
    margin-top: 28px;
    padding: 16px;
    font-family: 'JetBrains Mono', 'Fira Mono', Menlo, monospace;
    font-size: 0.85rem;
    line-height: 1.45;
    background: #0f172a;
    color: #e2e8f0;
    border-radius: 12px;
    max-height: 220px;
    overflow-y: auto;
    white-space: pre-wrap;
  }
</style>
<div class="speedtest-container">
  <h2 class="speedtest-heading">Wi-Fi / Ethernet Speed Test</h2>
  <p class="speedtest-lead">Run a download, upload, and latency check right from your browser. Results reflect the real connection between this device and Hugging Face Spaces.</p>
  <button id="startBtn" class="speedtest-button">Start Speed Test</button>
  <div id="results" class="speedtest-results">
    <div class="speedtest-grid">
      <div class="metric-card">
        <span class="metric-label">Download</span>
        <span id="dlSpeed" class="metric-value">-</span>
        <div class="metric-subtext">Average throughput in Mbps</div>
      </div>
      <div class="metric-card">
        <span class="metric-label">Upload</span>
        <span id="ulSpeed" class="metric-value">-</span>
        <div class="metric-subtext">Average throughput in Mbps</div>
      </div>
      <div class="metric-card">
        <span class="metric-label">Ping</span>
        <span id="ping" class="metric-value">-</span>
        <div class="metric-subtext">Round-trip latency in ms</div>
      </div>
    </div>
    <div id="log" class="speedtest-log"></div>
  </div>
</div>
<script>
  const startBtn = document.getElementById('startBtn');
  const resultsDiv = document.getElementById('results');
  const dlSpan = document.getElementById('dlSpeed');
  const ulSpan = document.getElementById('ulSpeed');
  const pingSpan = document.getElementById('ping');
  const logView = document.getElementById('log');

  function log(message) {
    const timestamp = new Date().toLocaleTimeString();
    logView.textContent += '[' + timestamp + '] ' + message + '
';
    logView.scrollTop = logView.scrollHeight;
  }

  function formatMbps(bitsPerSecond) {
    if (!Number.isFinite(bitsPerSecond) || bitsPerSecond <= 0) {
      return '0.00 Mbps';
    }
    return (bitsPerSecond / 1e6).toFixed(2) + ' Mbps';
  }

  async function measurePing(attempts = 5) {
    const samples = [];
    for (let i = 0; i < attempts; i++) {
      const start = performance.now();
      const response = await fetch('/ping_test?i=' + i + '&cacheBust=' + Math.random(), { cache: 'no-store' });
      if (!response.ok) {
        throw new Error('Ping request failed with status ' + response.status);
      }
      await response.text();
      const end = performance.now();
      samples.push(end - start);
    }
    const avg = samples.reduce((acc, cur) => acc + cur, 0) / samples.length;
    return avg;
  }

  async function measureDownload() {
    const start = performance.now();
    const response = await fetch('/download_test', { cache: 'no-store' });
    if (!response.ok) {
      throw new Error('Download request failed with status ' + response.status);
    }
    const blob = await response.blob();
    const end = performance.now();
    const durationSeconds = (end - start) / 1000;
    const bitsTransferred = blob.size * 8;
    return { bitsTransferred, durationSeconds };
  }

  async function measureUpload(bytesToSend) {
    const randomData = new Uint8Array(bytesToSend);
    crypto.getRandomValues(randomData);
    const start = performance.now();
    const response = await fetch('/upload_test', {
      method: 'POST',
      headers: { 'Content-Type': 'application/octet-stream' },
      body: randomData,
    });
    const end = performance.now();
    if (!response.ok) {
      throw new Error('Upload request failed with status ' + response.status);
    }
    const json = await response.json();
    if (!json || typeof json.received_bytes !== 'number') {
      throw new Error('Upload response missing received_bytes.');
    }
    const durationSeconds = (end - start) / 1000;
    const bitsTransferred = json.received_bytes * 8;
    return { bitsTransferred, durationSeconds };
  }

  startBtn.addEventListener('click', async () => {
    startBtn.disabled = true;
    resultsDiv.style.display = 'block';
    dlSpan.textContent = 'Testing...';
    ulSpan.textContent = 'Testing...';
    pingSpan.textContent = 'Testing...';
    logView.textContent = '';

    try {
      log('Measuring ping...');
      const pingMs = await measurePing();
      pingSpan.textContent = pingMs.toFixed(2) + ' ms';
      log('Average ping: ' + pingSpan.textContent);

      log('Running download test (10 MiB)...');
      const downloadResult = await measureDownload();
      const downloadBps = downloadResult.bitsTransferred / downloadResult.durationSeconds;
      dlSpan.textContent = formatMbps(downloadBps);
      log('Download duration: ' + downloadResult.durationSeconds.toFixed(2) + ' s');
      log('Download speed: ' + dlSpan.textContent);

      log('Running upload test (10 MiB)...');
      const uploadResult = await measureUpload(downloadResult.bitsTransferred / 8);
      const uploadBps = uploadResult.bitsTransferred / uploadResult.durationSeconds;
      ulSpan.textContent = formatMbps(uploadBps);
      log('Upload duration: ' + uploadResult.durationSeconds.toFixed(2) + ' s');
      log('Upload speed: ' + ulSpan.textContent);

      log('Speed test complete. Start again for a fresh run.');
    } catch (error) {
      console.error(error);
      const message = error && error.message ? error.message : String(error);
      log('Error: ' + message);
      dlSpan.textContent = 'Error';
      ulSpan.textContent = 'Error';
      pingSpan.textContent = 'Error';
    } finally {
      startBtn.disabled = false;
    }
  });
</script>
""")


    def _ensure_routes() -> None:
        register_backend_routes(demo)

    demo.load(fn=_ensure_routes, inputs=None, outputs=None)


def build_app() -> gr.Blocks:
    register_backend_routes(demo)
    return demo


app = build_app()


if __name__ == "__main__":
    import os
    port = int(os.environ.get("PORT", 7860))
    app.queue().launch(server_name="0.0.0.0", server_port=port)