"""File-based Ed25519 keys and authorized_keys helpers."""
from __future__ import annotations

from pathlib import Path

from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives.serialization import (
  Encoding,
  NoEncryption,
  PrivateFormat,
  PublicFormat,
  load_pem_private_key,
)

from browser_cli.auth.agent import AgentKey

def generate_keypair() -> tuple[bytes, str]:
  """Return (private_key_pem_bytes, public_key_hex)."""
  private_key = Ed25519PrivateKey.generate()
  pem = private_key.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption())
  public_hex = private_key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw).hex()
  return pem, public_hex

def load_private_key(path: Path) -> Ed25519PrivateKey:
  return load_pem_private_key(path.read_bytes(), password=None)

def public_key_hex(key: Ed25519PrivateKey | AgentKey) -> str:
  if isinstance(key, AgentKey):
    return key.pubkey_bytes.hex()
  return key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw).hex()

def _parse_authorized_line(line: str) -> tuple[str, str, list[str] | None] | None:
  """Parse one authorized_keys line into (pubkey, name, categories).

  Line format: ``<pubkey> [name words...] [allow:cat,cat,...]``. The optional
  ``allow:`` token may appear anywhere after the pubkey (conventionally last);
  the remaining words form the name. ``categories`` is None when no ``allow:``
  token is present (the key falls back to the server-wide policy), or a list of
  category strings (possibly empty) otherwise.
  """
  line = line.strip()
  if not line or line.startswith("#"):
    return None
  tokens = line.split()
  pubkey = tokens[0]
  categories: list[str] | None = None
  name_tokens: list[str] = []
  for tok in tokens[1:]:
    if tok.startswith("allow:"):
      categories = [c for c in tok[len("allow:"):].split(",") if c]
    else:
      name_tokens.append(tok)
  return pubkey, " ".join(name_tokens), categories

def format_authorized_line(pub_hex: str, name: str = "", categories: list[str] | None = None) -> str:
  """Render an authorized_keys line. Inverse of :func:`_parse_authorized_line`."""
  parts = [pub_hex]
  if name:
    parts.append(name)
  if categories is not None:
    parts.append("allow:" + ",".join(categories))
  return " ".join(parts)

def load_authorized_keys_with_names(path: Path) -> list[tuple[str, str]]:
  """Return list of (pubkey_hex, name) pairs. Name is empty string if not set."""
  return [(pubkey, name) for pubkey, name, _cats in load_authorized_keys_with_policies(path)]

def load_authorized_keys_with_policies(path: Path) -> list[tuple[str, str, list[str] | None]]:
  """Return list of (pubkey_hex, name, categories) triples. categories is None when unset."""
  if not path.exists():
    return []
  result = []
  for line in path.read_text(encoding="utf-8").splitlines():
    parsed = _parse_authorized_line(line)
    if parsed is not None:
      result.append(parsed)
  return result

def load_authorized_keys(path: Path) -> list[str]:
  return [pubkey for pubkey, _name in load_authorized_keys_with_names(path)]

def add_authorized_key(path: Path, pub_hex: str, name: str = "", categories: list[str] | None = None) -> bool:
  """Append pub_hex to authorized_keys. Returns False if already present."""
  path.parent.mkdir(parents=True, exist_ok=True)
  existing = {pubkey for pubkey, _name in load_authorized_keys_with_names(path)}
  if pub_hex in existing:
    return False
  line = format_authorized_line(pub_hex, name, categories) + "\n"
  with open(path, "a", encoding="utf-8") as file:
    file.write(line)
  return True

def set_authorized_key_policy(path: Path, identifier: str, categories: list[str] | None) -> tuple[str, str] | None:
  """Update the per-key policy for a trusted key.

  ``identifier`` may be the full public key or an exact key name. ``categories``
  is written as the ``allow:`` token; ``None`` removes the token so the key uses
  the server default. Returns ``(pubkey, name)`` for the updated key, ``None`` if
  no key matched, and raises ``ValueError`` for ambiguous names.
  """
  if not path.exists():
    return None

  wanted = identifier.strip()
  lines = path.read_text(encoding="utf-8").splitlines(keepends=True)
  matches: list[tuple[int, str, str, str]] = []

  for index, line in enumerate(lines):
    parsed = _parse_authorized_line(line)
    if parsed is None:
      continue
    pubkey, name, _cats = parsed
    if pubkey.lower() == wanted.lower() or (name and name == wanted):
      newline = "\n" if line.endswith("\n") else ""
      matches.append((index, pubkey, name, newline))

  if not matches:
    return None
  if len(matches) > 1:
    raise ValueError(f"ambiguous key name: {identifier!r} matches {len(matches)} keys")

  index, pubkey, name, newline = matches[0]
  lines[index] = format_authorized_line(pubkey, name, categories) + newline
  path.write_text("".join(lines), encoding="utf-8")
  return pubkey, name