import threading, secrets, socket, struct, click, json, sys
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) -> 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", "?")

  if server_token is not None:
    if msg.get("token") != 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

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

  strip = {"token", "_route"}
  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 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) -> None:
  with client_sock:
    _proxy_request(client_sock, addr, profile, server_token)

@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 token authentication.")
@click.pass_context
def cmd_serve(ctx, host, port, token, no_auth):
  """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 no_auth:
    server_token = None
  else:
    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:
    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 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), daemon=True).start()
  except KeyboardInterrupt:
    console.print("[yellow]Stopped.[/yellow]")
  finally:
    server.close()