"""
A signed-in account: its hub session, its bundle key `M`, its node identities.

Two lives. **Remembered** (after `login`): everything is in the store, and a
run needs no passphrase — an expired access token is renewed with the refresh
token, a refused one falls back to the device's hub key. **Ephemeral** (no
store, or another account named on the command line): signed in with the
passphrase for this run only, nothing written, the session ended on the hub
when the run ends.

The refresh token rotates and a spent one presented again revokes the whole
family, so every renewal happens under the store's lock, after re-reading what
another run may have just written.
"""

import asyncio
import contextlib
import logging
import socket
import time
from collections.abc import AsyncIterator, Callable
from dataclasses import dataclass

from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

from ..errors import AUTH, CliError
from .hub import Hub, HubError
from .keys import (
    Identity,
    auth_key,
    b64,
    bundle_master_key,
    passphrase_half,
    unb64,
)
from .store import KEY_FILE, Store

log = logging.getLogger(__name__)

# Renew this long before the access token runs out, so a request never leaves
# with a token that expires on the way.
RENEW_MARGIN = 300

PassphraseSource = Callable[[], str]


class NoPassphrase(CliError):
    """Raised by a passphrase source that has none to give."""

    def __init__(self, message: str):
        super().__init__(AUTH, message)


def no_passphrase() -> str:
    """The source for an account that must never ask: one being signed out."""
    raise NoPassphrase("No passphrase is asked for to sign out.")


@dataclass
class Pending:
    identity: Identity
    code: str
    expires_at: float


class Account:
    def __init__(self, hub: Hub, username: str, passphrase: PassphraseSource,
                 store: Store | None = None):
        self.hub = hub
        self.username = username
        self._passphrase = passphrase
        self.store = store            # None: ephemeral
        # Called with (group, DeviceLinkPending, first=) when this machine waits
        # for a device approval; returns whether to keep waiting.
        self.on_device_code = None
        self.state: dict = {}
        self.secrets: dict = {}

    @property
    def remembered(self) -> bool:
        return self.store is not None

    @property
    def user_id(self) -> str:
        return self.state["user_id"]

    @property
    def master(self) -> bytes | None:
        m = self.secrets.get("master")
        return unb64(m) if m else None

    @property
    def legacy(self) -> bytes | None:
        a = self.secrets.get("legacy")
        return unb64(a) if a else None

    # ── Opening ──────────────────────────────────────────────────────────────

    @classmethod
    def open(cls, hub: Hub, username: str | None, passphrase: PassphraseSource,
             store: Store) -> "Account":
        """The remembered account when it is the one asked for, else an ephemeral one."""
        stored = store.peek()
        same_hub = stored.get("hub") == hub.url
        same_user = username is None or username == stored.get("username")
        if stored and same_hub and same_user and store.load():
            acc = cls(hub, store.state["username"], passphrase, store)
            acc.state, acc.secrets = store.state, store.secrets
            return acc
        if not username:
            raise CliError(AUTH, "Not signed in. Run `meshbay-cli login`, or pass --user.")
        return cls(hub, username, passphrase)

    @contextlib.asynccontextmanager
    async def _locked(self) -> AsyncIterator[None]:
        """The store's lock, taken off the loop; re-reads what is stored."""
        if self.store is None:
            yield
            return
        cm = self.store.locked()
        await asyncio.to_thread(cm.__enter__)
        try:
            if not self.store.load():
                raise CliError(AUTH, "Signed out by another run; run `meshbay-cli login`.")
            self.state, self.secrets = self.store.state, self.store.secrets
            yield
        finally:
            cm.__exit__(None, None, None)

    def _save(self) -> None:
        """Under the lock. `sign_in` replaces both dicts, so hand over the current ones."""
        self.store.state, self.store.secrets = self.state, self.secrets
        self.store.save()

    async def update(self, mutate: Callable[[dict, dict], None]) -> None:
        """Change state and secrets, and write them when remembered."""
        async with self._locked():
            mutate(self.state, self.secrets)
            if self.store is not None:
                self._save()

    # ── Signing in ───────────────────────────────────────────────────────────

    async def sign_in(self, *, remember: bool) -> None:
        """With the passphrase: tokens, `M`, and — remembered — the device's hub key."""
        passphrase = self._passphrase()
        ak = await asyncio.to_thread(auth_key, self.username, passphrase)
        previous = dict(self.state)
        r = await self.hub.login(self.username, ak, previous.get("known_browser"))
        token = r["access_token"]
        me = await self.hub.me(token)
        if not r.get("bundle_pepper"):
            raise CliError(AUTH, "The hub did not provide the bundle pepper.")
        a = await asyncio.to_thread(passphrase_half, self.username, passphrase)
        master = bundle_master_key(a, unb64(r["bundle_pepper"]), me["user_id"])

        same_account = previous.get("user_id") == me["user_id"]
        state = {
            "hub": self.hub.url,
            "username": self.username,
            "user_id": me["user_id"],
            "known_browser": r.get("known_browser") or previous.get("known_browser"),
            # A node's key does not depend on who signed in.
            "pins": previous.get("pins", {}) if previous.get("hub") == self.hub.url else {},
        }
        if same_account and previous.get("browser_access"):
            state["browser_access"] = previous["browser_access"]
        secrets = {
            "access_token": token,
            "access_exp": time.time() + int(r.get("expires_in", 0)),
            "refresh_token": r["refresh_token"],
            "master": b64(master),
            "pepper_version": int(r.get("bundle_pepper_version") or 1),
            "legacy": b64(a),
            "identities": self.secrets.get("identities", {}) if same_account else {},
        }
        if remember:
            device = (self.secrets.get("device_sk") if same_account else None)
            sk = (Ed25519PrivateKey.from_private_bytes(unb64(device)) if device
                  else Ed25519PrivateKey.generate())
            reg = await self.hub.register_device(
                token, b64(sk.public_key().public_bytes_raw()),
                f"meshbay-cli on {socket.gethostname()}", ak)
            secrets["device_sk"] = b64(sk.private_bytes_raw())
            state["device_id"] = reg["id"]
        self.state, self.secrets = state, secrets

    async def remember(self, store: Store) -> str | None:
        """Write this account to `store`, replacing what it held. A notice, or None."""
        self.store = store
        async with self._locked_fresh():
            store.state, store.secrets = self.state, self.secrets
            store.save(new=True)
        return KEY_FILE if store.state.get("key_store") == KEY_FILE else None

    @contextlib.asynccontextmanager
    async def _locked_fresh(self) -> AsyncIterator[None]:
        cm = self.store.locked()
        await asyncio.to_thread(cm.__enter__)
        try:
            yield
        finally:
            cm.__exit__(None, None, None)

    # ── The access token ─────────────────────────────────────────────────────

    def _fresh(self) -> bool:
        return bool(self.secrets.get("access_token")) and \
            self.secrets.get("access_exp", 0) - RENEW_MARGIN > time.time()

    async def token(self, *, rejected: str | None = None) -> str:
        """A valid access token. `rejected`: the one the hub just refused."""
        if rejected is None and self._fresh():
            return self.secrets["access_token"]
        if not self.remembered:
            if not self.secrets:
                await self.sign_in(remember=False)
            else:
                await self._renew()
            return self.secrets["access_token"]
        async with self._locked():
            # Another run may have renewed while this one waited for the lock;
            # renewing again would present a refresh token it has just spent.
            current = self.secrets.get("access_token")
            if not self._fresh() or current == rejected:
                await self._renew()
                self._save()
        return self.secrets["access_token"]

    async def _renew(self) -> None:
        rt = self.secrets.get("refresh_token")
        if rt:
            try:
                r = await self.hub.refresh(rt)
                self._take(r)
                return
            except HubError as e:
                log.debug("refresh refused: %s", e)
        sk = self.secrets.get("device_sk")
        if sk:
            try:
                r = await self.hub.device_auth(
                    self.username, Ed25519PrivateKey.from_private_bytes(unb64(sk)))
                self._take(r)
                return
            except CliError as e:
                log.debug("device sign-in refused: %s", e.message)
        try:
            await self.sign_in(remember=self.remembered)
        except NoPassphrase as e:
            if self.remembered:
                raise CliError(AUTH, "The saved session has ended; run "
                                     "`meshbay-cli login`, or pass --password.") from e
            raise

    def _take(self, r: dict) -> None:
        self.secrets["access_token"] = r["access_token"]
        self.secrets["access_exp"] = time.time() + int(r.get("expires_in", 0))
        self.secrets["refresh_token"] = r["refresh_token"]

    async def call(self, fn, *args):
        """A hub call with the access token, renewed once on a 401."""
        token = await self.token()
        try:
            return await fn(token, *args)
        except HubError as e:
            if e.status != 401:
                raise
        return await fn(await self.token(rejected=token), *args)

    # ── Identities and pins ──────────────────────────────────────────────────

    def identity(self, node_pk: str) -> Identity | None:
        keys = self.secrets.get("identities", {}).get(node_pk)
        return Identity.from_pkcs8(keys) if keys else None

    async def keep_identity(self, node_pk: str, identity: Identity) -> None:
        def put(state: dict, secrets: dict) -> None:
            secrets.setdefault("identities", {})[node_pk] = identity.to_pkcs8()
        await self.update(put)

    async def forget_identity(self, node_pk: str) -> None:
        def drop(state: dict, secrets: dict) -> None:
            secrets.get("identities", {}).pop(node_pk, None)
        await self.update(drop)

    # A device request this machine filed and that waits for an approval: the
    # keys it named, and the code shown for them. Kept so a later run shows the
    # same code and joins with the same keys once approved.

    def pending_identity(self, node_pk: str) -> "Pending | None":
        p = self.secrets.get("pending", {}).get(node_pk)
        if not p or p["expires_at"] < time.time():
            return None
        return Pending(Identity.from_pkcs8(p["keys"]), p["code"], p["expires_at"])

    async def keep_pending(self, node_pk: str, identity: Identity, code: str,
                           expires_at: float) -> None:
        def put(state: dict, secrets: dict) -> None:
            secrets.setdefault("pending", {})[node_pk] = {
                "keys": identity.to_pkcs8(), "code": code, "expires_at": expires_at}
        await self.update(put)

    async def adopt_pending(self, node_pk: str) -> None:
        """The node admitted the pending keys: they are this machine's identity there."""
        def move(state: dict, secrets: dict) -> None:
            p = secrets.get("pending", {}).pop(node_pk, None)
            if p:
                secrets.setdefault("identities", {})[node_pk] = p["keys"]
        await self.update(move)

    # An invitation's code waiting for its node: kept when the node was not
    # reachable at `join`, sent by the next connection to it, dropped once used
    # or refused. A link's code names its node; an invitation's does not.

    def join_code(self, group_id: str, node_pk: str) -> str | None:
        c = self.secrets.get("join_codes", {}).get(group_id)
        if not c or (c.get("node_pk") and c["node_pk"] != node_pk):
            return None
        return c["code"]

    async def keep_join_code(self, group_id: str, code: str, node_pk: str | None) -> None:
        def put(state: dict, secrets: dict) -> None:
            secrets.setdefault("join_codes", {})[group_id] = {"code": code,
                                                              "node_pk": node_pk}
        await self.update(put)

    async def drop_join_code(self, group_id: str) -> None:
        def drop(state: dict, secrets: dict) -> None:
            secrets.get("join_codes", {}).pop(group_id, None)
        if group_id in self.secrets.get("join_codes", {}):
            await self.update(drop)

    @property
    def browser_access(self) -> str | None:
        """This machine's setting: "on", "off", or None (unset — leave the node
        as it is). Set only by `browser-access`; never from the hub."""
        return self.state.get("browser_access")

    @property
    def pepper_version(self) -> int:
        return int(self.secrets.get("pepper_version") or 1)

    def pin(self, node_id: str) -> str | None:
        return self.state.get("pins", {}).get(node_id)

    async def set_pin(self, node_id: str, node_pk: str) -> None:
        def put(state: dict, secrets: dict) -> None:
            state.setdefault("pins", {})[node_id] = node_pk
        await self.update(put)

    # ── Ending ───────────────────────────────────────────────────────────────

    async def close(self) -> None:
        """An ephemeral session ends on the hub with the run."""
        if not self.remembered and self.secrets.get("refresh_token"):
            with contextlib.suppress(Exception):
                await self.hub.logout(self.secrets["refresh_token"])

    async def sign_out(self) -> list[str]:
        """End the session, retire the device's hub key, erase the store. Warnings."""
        warnings = []
        device_id = self.state.get("device_id")
        if device_id:
            try:
                await self.call(self.hub.delete_device, device_id)
            except (HubError, CliError) as e:
                warnings.append(
                    "The hub did not retire this machine's device key "
                    f"({getattr(e, 'detail', None) or getattr(e, 'message', e)}); remove "
                    f"\"meshbay-cli on {socket.gethostname()}\" from your account's devices.")
        rt = self.secrets.get("refresh_token")
        if rt:
            with contextlib.suppress(Exception):
                await self.hub.logout(rt)
        if self.store is not None:
            async with self._locked_fresh():
                self.store.erase()
        return warnings