import asyncio, io, time
from concurrent.futures import ThreadPoolExecutor
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from google.genai import Client, types
import httpx

class Peer(BaseHTTPRequestHandler):
  attempts: int = 0

  def do_POST(self) -> None:
    self.rfile.read(int(self.headers['Content-Length']))
    start = self.headers['X-Goog-Upload-Command'] == 'start'
    if not start:
      Peer.attempts += 1
    payload = b'{"file":{"name":"files/local"}}'
    self.send_response(200)
    if start:
      self.send_header('X-Goog-Upload-URL',
                       f'http://{self.headers["Host"]}/upload/v1beta/files')
    elif Peer.attempts == 2:
      self.send_header('X-Goog-Upload-Status', 'final')
    self.send_header('Content-Type', 'application/json')
    self.send_header('Content-Length', str(len(payload)))
    self.end_headers()
    self.wfile.write(payload)

async def check_upload(base_url: str) -> None:
  with Client(
          vertexai=False, api_key='local-only',
          http_options=types.HttpOptions(
              base_url=base_url,
              httpx_async_client=httpx.AsyncClient(trust_env=False),
          ),
      ) as client, io.BytesIO(b'local') as payload:
    async with client.aio:
      upload = asyncio.create_task(client.aio.files.upload(
          file=payload, config={'mime_type': 'text/plain'}))
      gap = 0.0
      while not upload.done():
        last = time.monotonic()
        await asyncio.sleep(0.01)
        gap = max(gap, time.monotonic() - last)
      print((await upload).name, round(gap, 3))
      assert gap < 0.2

with ThreadingHTTPServer(('127.0.0.1', 0), Peer) as server, ThreadPoolExecutor(max_workers=1) as pool:
  pool.submit(server.serve_forever)
  try:
    asyncio.run(check_upload(f'http://127.0.0.1:{server.server_port}'))
  finally:
    server.shutdown()