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

from rich.console import Console

from browser_cli.client import _recv_exact, _recv_all
from browser_cli.compat import adapt_auth, adapt_request, adapt_response
from browser_cli.version_manager import PROTOCOL_MIN_CLIENT, MAX_MSG_BYTES, parse_version, get_installed_version

_UA_PATTERN = re.compile(r"^browser-cli/\d")
_CONN_LIMIT = threading.BoundedSemaphore(64)
console = Console()


def _framed_send(sock: socket.socket, data: bytes) -> None:
  sock.sendall(struct.pack("<I", len(data)) + data)

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, auth_keys:list[str]|None, auth_keys_path:"Path|None", nonce:str, pq_private_key=None) -> None:
  from browser_cli.client import _resolve_socket, BrowserNotConnected
  from browser_cli.platform import is_windows

  def _send_error(msg_id, msg:str) -> None:
    err = json.dumps({"id": msg_id, "success": False, "error": msg}).encode()
    try:
      _framed_send(client_sock, err)
    except OSError:
      pass

  def _send_ok(msg_id, payload) -> None:
    out = json.dumps({"id": msg_id, "success": True, "data": payload}).encode()
    try:
      _framed_send(client_sock, out)
    except OSError:
      pass

  try:
    header = _recv_exact(client_sock, 4)
    msg_len = struct.unpack("<I", header)[0]
    if msg_len > MAX_MSG_BYTES:
      _send_error(None, f"message too large ({msg_len} bytes)")
      return
    payload = _recv_exact(client_sock, msg_len)
  except (ConnectionError, OSError):
    return

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

  # ── user-agent + version check ────────────────────────────────────────────
  msg_id = msg.get("id")
  ua = msg.get("user_agent") or ""
  if not _UA_PATTERN.match(ua):
    _send_error(msg_id, "forbidden: client required")
    _log(addr, msg.get("command", "?"), None, "DENIED", f"bad user-agent: {ua!r}")
    return
  client_ver = "0"
  try:
    client_ver = ua.split("/", 1)[1]
    if parse_version(client_ver) < parse_version(PROTOCOL_MIN_CLIENT):
      _send_error(msg_id, f"client version {client_ver} is too old; please upgrade to >= {PROTOCOL_MIN_CLIENT}")
      _log(addr, msg.get("command", "?"), None, "DENIED", f"client {client_ver} < min {PROTOCOL_MIN_CLIENT}")
      return
  except (IndexError, ValueError):
    pass

  msg = adapt_auth(msg, client_ver)
  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
    pq_shared_secret = None
    if pq_private_key is not None:
      kex = msg.get("pq_kex") or {}
      pq_required = parse_version(client_ver) >= parse_version("0.9.4")
      if not isinstance(kex, dict) or kex.get("alg") != "ML-KEM-768" or not kex.get("ciphertext"):
        if pq_required:
          _send_error(msg_id, "unauthorized: post-quantum key exchange required")
          _log(addr, command, None, "DENIED", "missing pq kex")
          return
      else:
        try:
          from browser_cli.auth import pq_kex_server_decapsulate
          pq_shared_secret = pq_kex_server_decapsulate(pq_private_key, str(kex["ciphertext"]))
        except Exception:
          _send_error(msg_id, "unauthorized: invalid post-quantum key exchange")
          _log(addr, command, None, "DENIED", "bad pq kex")
          return

    from browser_cli.auth import verify
    if not verify(pub, bytes.fromhex(nonce), msg, sig, pq_shared_secret):
      _send_error(msg_id, "unauthorized: invalid signature")
      _log(addr, command, None, "DENIED", "bad signature")
      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)
    ]
    _send_ok(msg_id, targets)
    _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)]
    _send_ok(msg_id, entries)
    _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 not re.fullmatch(r"[0-9a-f]{64}", pubkey):
      _send_error(msg_id, "invalid pubkey: expected 64 lowercase hex characters")
      _log(addr, command, None, "ERROR", "invalid pubkey")
      return
    added = add_authorized_key(auth_keys_path, pubkey, name)
    _send_ok(msg_id, {"added": added})
    _log(addr, command, None, "OK" if added else "ALREADY_TRUSTED")
    return

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

  # ── strip protocol fields, apply request compat shim, forward ─────────────
  strip = {"token", "_route", "pubkey", "sig", "user_agent", "pq_kex"}
  clean_msg = {k: v for k, v in msg.items() if k not in strip}
  clean_msg = adapt_request(clean_msg, client_ver)
  clean_payload = json.dumps(clean_msg).encode()
  clean_header = struct.pack("<I", len(clean_payload))

  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_payload = pipe.recv_bytes()
      resp_payload = adapt_response(resp_payload, command, client_ver)
      _framed_send(client_sock, resp_payload)
    else:
      with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as local:
        local.connect(sock_path)
        local.sendall(clean_header + clean_payload)
        resp_payload = _recv_all(local)
      resp_payload = adapt_response(resp_payload, command, client_ver)
      _framed_send(client_sock, resp_payload)

    resp_data = json.loads(resp_payload)
    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, 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)
      pq_private_key = None
      challenge_msg = {
        "type": "challenge",
        "nonce": nonce,
        "server_version": get_installed_version(),
        "min_client_version": PROTOCOL_MIN_CLIENT,
      }
      if auth_keys_path is not None:
        from browser_cli.auth import PQ_KEX_ALG, pq_kex_server_keypair
        pq_keypair = pq_kex_server_keypair()
        if pq_keypair is not None:
          pq_private_key, pq_public_key = pq_keypair
          challenge_msg["pq_kex"] = {"alg": PQ_KEX_ALG, "public_key": pq_public_key.hex()}
      challenge = json.dumps(challenge_msg).encode()
      try:
        _framed_send(client_sock, challenge)
      except OSError:
        return
      _proxy_request(client_sock, addr, profile, auth_keys, auth_keys_path, nonce, pq_private_key)
  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("--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). Required unless --no-auth.")
@click.pass_context
def cmd_serve(ctx, host, port, 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}")
  elif no_auth:
    auth_keys_path = None
  else:
    console.print("[red]Error:[/red] --authorized-keys FILE is required. Use --no-auth to explicitly disable auth (dangerous).")
    sys.exit(1)

  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)

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

  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]")
  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, auth_keys_path), daemon=True).start()
  except KeyboardInterrupt:
    console.print("[yellow]Stopped.[/yellow]")
  finally:
    server.close()