import hmac, threading, secrets, socket, struct, click, json, sys
from pathlib import Path

_CONN_LIMIT = threading.BoundedSemaphore(64)
from rich.console import Console
from datetime import datetime

console = Console()

def _recv_exact(sock:socket.socket, n:int) -> bytes:
  buf = b""
  while len(buf) < n:
    chunk = sock.recv(n - len(buf))
    if not chunk:
      raise ConnectionError("Connection closed")
    buf += chunk
  return buf

def _log(addr:tuple, command:str, profile:str|None, status:str, error:str|None=None) -> None:
  ts = datetime.now().strftime("%H:%M:%S")
  addr_str = f"{addr[0]}:{addr[1]}"
  profile_str = f"[dim]{profile}[/dim] " if profile else ""
  if error:
    console.print(f"[dim]{ts}[/dim] {addr_str} {profile_str}[cyan]{command}[/cyan] [red]{status}[/red] {error}")
  else:
    console.print(f"[dim]{ts}[/dim] {addr_str} {profile_str}[cyan]{command}[/cyan] [green]{status}[/green]")

def _proxy_request(client_sock:socket.socket, addr:tuple, profile:str|None, server_token:str|None, auth_keys:list[str]|None, auth_keys_path:"Path|None", nonce:str) -> None:
  from browser_cli.client import _resolve_socket, BrowserNotConnected
  from browser_cli.platform import is_windows

  try:
    header = _recv_exact(client_sock, 4)
    msg_len = struct.unpack("<I", header)[0]
    payload = _recv_exact(client_sock, msg_len)
  except (ConnectionError, OSError):
    return

  def _send_error(msg_id, msg:str) -> None:
    err = json.dumps({"id": msg_id, "success": False, "error": msg}).encode()
    try:
      client_sock.sendall(struct.pack("<I", len(err)) + err)
    except OSError:
      pass

  try:
    msg = json.loads(payload)
  except (json.JSONDecodeError, ValueError):
    _send_error(None, "invalid JSON")
    _log(addr, "?", None, "ERROR", "invalid JSON")
    return

  msg_id = msg.get("id")
  command = msg.get("command", "?")

  # ── auth ──────────────────────────────────────────────────────────────────
  if auth_keys is not None:
    pub = msg.get("pubkey") or ""
    sig = msg.get("sig") or ""
    if not pub or not sig:
      _send_error(msg_id, "unauthorized: pubkey auth required — run 'browser-cli auth keygen' on the client")
      _log(addr, command, None, "DENIED", "missing pubkey/sig")
      return
    if pub not in auth_keys:
      _send_error(msg_id, "unauthorized: untrusted public key")
      _log(addr, command, None, "DENIED", "untrusted key")
      return
    from browser_cli.auth import verify
    if not verify(pub, bytes.fromhex(nonce), msg, sig):
      _send_error(msg_id, "unauthorized: invalid signature")
      _log(addr, command, None, "DENIED", "bad signature")
      return
  elif server_token is not None:
    if not hmac.compare_digest(msg.get("token") or "", server_token):
      _send_error(msg_id, "unauthorized: invalid or missing token")
      _log(addr, command, None, "DENIED", "bad token")
      return

  if command == "browser-cli.targets":
    from browser_cli.client import active_browser_targets
    targets = [
      {"profile": target.profile, "displayName": target.display_name}
      for target in active_browser_targets(include_remotes=False)
    ]
    data = json.dumps({"id": msg_id, "success": True, "data": targets}).encode()
    client_sock.sendall(struct.pack("<I", len(data)) + data)
    _log(addr, command, None, "OK")
    return

  if command == "browser-cli.auth.keys":
    if auth_keys_path is None:
      _send_error(msg_id, "no authorized keys file configured on this server")
      _log(addr, command, None, "ERROR", "no authorized keys file")
      return
    from browser_cli.auth import load_authorized_keys_with_names
    entries = [{"pubkey": pk, "name": name} for pk, name in load_authorized_keys_with_names(auth_keys_path)]
    data = json.dumps({"id": msg_id, "success": True, "data": entries}).encode()
    client_sock.sendall(struct.pack("<I", len(data)) + data)
    _log(addr, command, None, "OK")
    return

  if command == "browser-cli.auth.trust":
    if auth_keys_path is None:
      _send_error(msg_id, "no authorized keys file configured on this server")
      _log(addr, command, None, "ERROR", "no authorized keys file")
      return
    from browser_cli.auth import add_authorized_key
    args = msg.get("args") or {}
    pubkey = str(args.get("pubkey") or "")
    name = str(args.get("name") or "")
    if len(pubkey) != 64:
      _send_error(msg_id, "invalid pubkey: expected 64 hex characters")
      _log(addr, command, None, "ERROR", "invalid pubkey")
      return
    added = add_authorized_key(auth_keys_path, pubkey, name)
    data = json.dumps({"id": msg_id, "success": True, "data": {"added": added}}).encode()
    client_sock.sendall(struct.pack("<I", len(data)) + data)
    _log(addr, command, None, "OK" if added else "ALREADY_TRUSTED")
    return

  resolved_profile = msg.get("_route") or profile

  strip = {"token", "_route", "pubkey", "sig"}
  if strip & msg.keys():
    clean_payload = json.dumps({k: v for k, v in msg.items() if k not in strip}).encode()
    clean_header = struct.pack("<I", len(clean_payload))
  else:
    clean_payload = payload
    clean_header = header

  try:
    sock_path = _resolve_socket(resolved_profile)
  except BrowserNotConnected as e:
    _send_error(msg_id, str(e))
    _log(addr, command, resolved_profile, "ERROR", "browser not connected")
    return

  try:
    if is_windows():
      from multiprocessing.connection import Client as PipeClient
      with PipeClient(sock_path, family="AF_PIPE") as pipe:
        pipe.send_bytes(clean_payload)
        resp = pipe.recv_bytes()
      client_sock.sendall(struct.pack("<I", len(resp)) + resp)
    else:
      with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as local:
        local.connect(sock_path)
        local.sendall(clean_header + clean_payload)
        resp_header = _recv_exact(local, 4)
        resp_len = struct.unpack("<I", resp_header)[0]
        resp_payload = _recv_exact(local, resp_len)
      client_sock.sendall(resp_header + resp_payload)

    resp_data = json.loads(resp_payload if not is_windows() else resp)
    if resp_data.get("success", True):
      _log(addr, command, resolved_profile, "OK")
    else:
      _log(addr, command, resolved_profile, "ERROR", resp_data.get("error", ""))
  except (OSError, json.JSONDecodeError) as e:
    _send_error(msg_id, str(e))
    _log(addr, command, resolved_profile, "ERROR", str(e))

def _handle_client(client_sock:socket.socket, addr:tuple, profile:str|None, server_token:str|None, auth_keys_path:"Path|None") -> None:
  if not _CONN_LIMIT.acquire(blocking=False):
    client_sock.close()
    return
  client_sock.settimeout(30)
  try:
    with client_sock:
      # reload on every connection so auth trust --remote takes effect immediately
      if auth_keys_path is not None:
        from browser_cli.auth import load_authorized_keys
        auth_keys: list[str] | None = load_authorized_keys(auth_keys_path)
      else:
        auth_keys = None
      nonce = secrets.token_hex(32)
      challenge = json.dumps({"type": "challenge", "nonce": nonce}).encode()
      try:
        client_sock.sendall(struct.pack("<I", len(challenge)) + challenge)
      except OSError:
        return
      _proxy_request(client_sock, addr, profile, server_token, auth_keys, auth_keys_path, nonce)
  finally:
    _CONN_LIMIT.release()

@click.command("serve")
@click.option("--host", default="127.0.0.1", show_default=True, help="Address to bind.")
@click.option("--port", default=8765, show_default=True, type=int, help="TCP port to listen on.")
@click.option("--token", default=None, metavar="TOKEN", help="Auth token (auto-generated if omitted).")
@click.option("--no-auth", is_flag=True, default=False, help="Disable authentication (dangerous).")
@click.option("--authorized-keys", "auth_keys_file", default=None, metavar="FILE",
              help="File of trusted Ed25519 public keys (one hex per line). Enables pubkey auth.")
@click.pass_context
def cmd_serve(ctx, host, port, token, no_auth, auth_keys_file):
  """Expose this browser over TCP so remote hosts can control it."""
  profile = ctx.obj.get("browser") if ctx.obj else None

  if host in ("0.0.0.0", "::"):
    console.print("[yellow]Warning:[/yellow] Binding to all interfaces — anyone who can reach this port controls your browser.")

  if auth_keys_file:
    from browser_cli.auth import load_authorized_keys
    auth_keys_path = Path(auth_keys_file)
    if not load_authorized_keys(auth_keys_path):
      console.print(f"[yellow]Warning:[/yellow] No authorized keys found in {auth_keys_path}")
    server_token = None
  elif no_auth:
    auth_keys_path = None
    server_token = None
  else:
    auth_keys_path = None
    server_token = token or secrets.token_urlsafe(32)

  server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  try:
    server.bind((host, port))
  except OSError as e:
    server.close()
    console.print(f"[red]Cannot bind to {host}:{port}:[/red] {e}")
    sys.exit(1)
  server.listen(16)

  browser_hint = f" (browser: {profile})" if profile else ""
  console.print(f"[green]Serving browser{browser_hint} →[/green] [cyan]{host}:{port}[/cyan]")

  if auth_keys_path is not None:
    from browser_cli.auth import load_authorized_keys
    n = len(load_authorized_keys(auth_keys_path))
    console.print(f"  Auth:   [bold green]Ed25519 pubkey[/bold green] ({n} trusted key{'s' if n != 1 else ''})")
    console.print(f"  CLI:    [dim]browser-cli --remote {host}:{port} tabs list[/dim]")
    console.print(f"  Python: [dim]BrowserCLI(remote=\"{host}:{port}\").tabs_list()[/dim]")
  elif server_token:
    console.print(f"  Token:  [bold yellow]{server_token}[/bold yellow]")
    console.print(f"  CLI:    [dim]browser-cli --remote {host}:{port} --token {server_token} tabs list[/dim]")
    console.print(f"  Python: [dim]BrowserCLI(remote=\"{host}:{port}\", token=\"{server_token}\").tabs_list()[/dim]")
  else:
    console.print(f"  CLI:    [dim]browser-cli --remote {host}:{port} tabs list[/dim]")
    console.print(f"  Python: [dim]BrowserCLI(remote=\"{host}:{port}\").tabs_list()[/dim]")
    console.print("[yellow]  Auth disabled (--no-auth)[/yellow]")

  console.print("Ctrl-C to stop.\n")

  try:
    while True:
      conn, addr = server.accept()
      threading.Thread(target=_handle_client, args=(conn, addr, profile, server_token, auth_keys_path), daemon=True).start()
  except KeyboardInterrupt:
    console.print("[yellow]Stopped.[/yellow]")
  finally:
    server.close()