"""One reconcile tick — the loop of MONKY-ADR-0028 §D, with Reconciliation v2 vocabulary:

    flock -> load state -> checkin (bearer = Bao token; bootstrap it from the kit's grant if
    absent) -> action:
        down  -> compose down --remove-orphans [-v]      -> report down
        none  -> compose ps; healthy?                    -> report applied (heartbeat)
        apply -> bundle (sha verified) -> refusal checks -> lease -> jwt-tenancy login
                 -> KV reads -> .env 0600 -> promote -> pull -> up -d --remove-orphans
                 -> wait healthy                          -> report applied | failed
    finally: token upkeep (renew-self / re-lease before max TTL), scrub, save state.

Exit codes: 0 ok (or offline in laptop mode) · 75 temporary network failure · 78 AGENT_ENV_MISMATCH
(never retried) · 1 refusal / failure."""

from __future__ import annotations

import logging
import os
import platform
import shutil
import socket
import time
from pathlib import Path

from monky_deployd import __version__
from monky_deployd import bundle as bundlemod
from monky_deployd import state as statemod
from monky_deployd.bao import BaoClient, BaoError, BaoToken, ManifestPathError, kv_data_path
from monky_deployd.bundle import Bundle, BundleError
from monky_deployd.compose import Compose, ComposeError, Container, Docker
from monky_deployd.config import Config
from monky_deployd.redact import REDACTOR
from monky_deployd.state import Lock, State
from monky_deployd.tenancy import (
    Checkin,
    EnvMismatch,
    LeaseShapeUnsupported,
    RateLimited,
    TenancyClient,
    TenancyError,
    Unauthenticated,
)
from monky_deployd.transport import HttpClient, TransportError, build

log = logging.getLogger("monky-deployd")

EX_OK, EX_FAIL, EX_TEMPFAIL, EX_ENV_MISMATCH = 0, 1, 75, 78
LEASE_MIN_GAP_S = 15 * 60  # never re-lease more often than this (tenancy allows 5/h)


class Refusal(Exception):
    """The agent refuses to apply; reported to tenancy as `failed` with the code + names only."""

    def __init__(self, code: str, detail: str):
        super().__init__(f"{code}: {detail}")
        self.code, self.detail = code, detail


class NoCredentials(Exception):
    pass


class TailHandler(logging.Handler):
    """Keeps the (redacted) last 16 KiB of this tick's log for the report's `log_tail`."""

    def __init__(self, limit: int = 16 * 1024):
        super().__init__(logging.INFO)
        self.limit = limit
        self.lines: list[str] = []
        self.size = 0
        self.addFilter(REDACTOR)
        self.setFormatter(logging.Formatter("%(levelname)s %(message)s"))

    def emit(self, record: logging.LogRecord) -> None:
        line = self.format(record)
        self.lines.append(line)
        self.size += len(line) + 1
        while self.size > self.limit and self.lines:
            self.size -= len(self.lines.pop(0)) + 1

    def text(self) -> str:
        return REDACTOR.scrub("\n".join(self.lines))


class Agent:
    def __init__(self, cfg: Config, *, prune: bool = False, docker: Docker | None = None):
        self.cfg = cfg
        self.prune = prune
        self.docker = docker or Docker(cfg.docker_bin)
        self.transport = build(cfg)
        self.tenancy_http = HttpClient(
            self.transport, cfg.tenancy.scheme, cfg.tenancy.host, cfg.tenancy.port, timeout=cfg.tenancy.timeout_s
        )
        scheme, host, port = cfg.bao_url
        self.bao_http = HttpClient(
            self.transport, scheme, host, port, ca_bundle=cfg.bao.ca_bundle, timeout=cfg.bao.timeout_s
        )
        self.bao = BaoClient(self.bao_http, auth_mount=cfg.bao.mount, role=cfg.bao.role, kv_mount=cfg.bao.kv_mount)
        self.state: State = State(env_id=cfg.env_id)
        self.token: str | None = None
        self.tenancy: TenancyClient | None = None
        self.tail = TailHandler()
        self._staging: list[Path] = []
        self._values: dict[str, str] = {}
        self._releases_this_tick = 0

    # -- entry -------------------------------------------------------------------------------------
    def run_once(self) -> int:
        lock = Lock(self.cfg.lock_path)
        if not lock.acquire():
            log.info("another monky-deployd tick holds %s; skipping", self.cfg.lock_path)
            return EX_OK
        lg = logging.getLogger("monky-deployd")
        if lg.getEffectiveLevel() > logging.INFO:
            lg.setLevel(logging.INFO)  # the tail must see INFO even when the root is quieter
        lg.addHandler(self.tail)
        try:
            self.state = statemod.load(self.cfg.state_path, self.cfg.env_id)
            self.state.agent_version = __version__
            return self._guarded_tick()
        finally:
            try:
                self._maintain_token()
            except Exception as exc:  # never let upkeep mask the tick's own result
                log.warning("token upkeep skipped: %s", REDACTOR.scrub(str(exc)))
            self._scrub()
            try:
                statemod.save(self.cfg.state_path, self.state)
            except OSError as exc:
                log.error("cannot save state: %s", exc)
            lg.removeHandler(self.tail)
            lock.release()

    def _guarded_tick(self) -> int:
        try:
            return self._tick()
        except TransportError as exc:
            self.state.last_error = f"network: {exc}"
            if self.cfg.laptop_mode:
                log.info("offline (%s); laptop mode, nothing to do", exc)
                return EX_OK
            log.warning("temporary network failure: %s", exc)
            return EX_TEMPFAIL
        except EnvMismatch as exc:
            self.state.last_error = f"AGENT_ENV_MISMATCH: {exc.detail}"
            log.error(
                "AGENT_ENV_MISMATCH: this box's token is pinned to another environment (%s); "
                "config says %s. Not retrying — fix the config or re-run the install kit.",
                exc.detail,
                self.cfg.env_id,
            )
            return EX_ENV_MISMATCH
        except Unauthenticated as exc:
            self.state.last_error = f"AGENT_UNAUTHENTICATED: {exc.detail}"
            log.error(
                "AGENT_UNAUTHENTICATED: bearer refused (%s). The deploy grant was superseded or the "
                "token revoked; re-run the install kit (a fresh bootstrap grant).",
                exc.detail,
            )
            return EX_FAIL
        except NoCredentials as exc:
            self.state.last_error = str(exc)
            log.error("%s", exc)
            return EX_FAIL
        except RateLimited as exc:
            self.state.last_error = f"{exc.code}: retry after {exc.retry_after}s"
            log.warning("%s (retry-after %ss)", exc, exc.retry_after)
            return EX_TEMPFAIL
        except (Refusal, BundleError, ManifestPathError, LeaseShapeUnsupported) as exc:
            code = getattr(exc, "code", exc.__class__.__name__)
            detail = getattr(exc, "detail", str(exc))
            self.state.last_error = f"{code}: {detail}"
            self.state.last_result = "failed"
            log.error("refused: %s: %s", code, detail)
            self._report("failed", self.state.desired_sha, detail=f"{code}: {detail}")
            return EX_FAIL
        except (TenancyError, BaoError, ComposeError) as exc:
            self.state.last_error = str(exc)
            self.state.last_result = "failed"
            log.error("failed: %s", exc)
            self._report("failed", self.state.desired_sha, detail=str(exc))
            return EX_FAIL

    # -- the tick ------------------------------------------------------------------------------------
    def _tick(self) -> int:
        cfg = self.cfg
        log.info("tick env=%s site=%s transport=%s", cfg.env_id, cfg.site, self.transport.describe())
        self._ensure_token()
        containers = self._compose().ps() if self.docker.available() else []
        ci = self._checkin(containers)
        self.state.last_checkin_at = time.time()
        self.state.last_action = ci.action
        self.state.desired_sha = ci.desired_sha
        if ci.vault.mount and ci.vault.mount != self.bao.kv_mount:
            log.info("KV mount from tenancy: %s", ci.vault.mount)
            self.bao.kv_mount = ci.vault.mount
        log.info(
            "checkin: action=%s desired=%s applied=%s",
            ci.action,
            _short(ci.desired_sha),
            _short(self.state.applied_sha),
        )
        if ci.action == "down":
            return self._do_down(ci)
        if ci.action == "none":
            return self._do_none(containers)
        return self._do_apply(ci)

    def _checkin(self, containers: list[Container]) -> Checkin:
        assert self.tenancy is not None
        kwargs = dict(
            agent_version=__version__,
            applied_sha=self.state.applied_sha,
            host=self._host_facts(),
            containers=[c.as_dict() for c in containers] or None,
        )
        try:
            return self.tenancy.checkin(**kwargs)
        except Unauthenticated as exc:
            # a superseded grant / revoked token: drop it; a bootstrap grant on disk rescues us once
            log.warning("checkin refused (%s); discarding the stored token", exc.detail)
            statemod.delete(self.cfg.token_path)
            self.state.token = statemod.TokenMeta()
            self.token = None
            if not Path(self.cfg.bootstrap_path).exists():
                raise
            self._ensure_token()
            return self.tenancy.checkin(**kwargs)

    # -- actions ---------------------------------------------------------------------------------------
    def _do_none(self, containers: list[Container]) -> int:
        if not self.state.applied_sha:
            log.info("nothing desired, nothing applied; idle")
            return EX_OK
        healthy = bool(containers) and all(c.ok for c in containers)
        if healthy:
            self._report("applied", self.state.applied_sha, containers=containers)
            log.info("healthy; heartbeat reported for %s", _short(self.state.applied_sha))
            return EX_OK
        bad = ", ".join(f"{c.name}={c.state}/{c.health or '-'}" for c in containers) or "no containers"
        log.warning("unhealthy: %s", bad)
        self.state.last_error = f"unhealthy: {bad}"
        self._report("failed", self.state.applied_sha, containers=containers, detail=f"unhealthy: {bad}")
        return EX_FAIL

    def _do_down(self, ci: Checkin) -> int:
        purge = ci.purge_volumes or self.cfg.volumes_on_absent == "purge"
        if purge and self.cfg.is_prod:
            log.warning("purge_volumes requested on a prod backend: REFUSED, volumes kept")
            purge = False
        compose = self._compose()
        log.info("desired action down: compose down --remove-orphans%s", " -v" if purge else "")
        if self.docker.available():
            compose.down(purge_volumes=purge)
        current = Path(self.cfg.deploy_dir) / "current"
        if current.is_symlink() or current.exists():
            current.unlink()
        self.state.applied_sha = None
        self.state.last_result = "down"
        self._report("down", None)
        return EX_OK

    def _do_apply(self, ci: Checkin) -> int:
        assert self.tenancy is not None
        cfg = self.cfg
        if not ci.desired_sha or not ci.bundle_url:
            raise Refusal("BAD_CHECKIN", "action apply without desired_sha/bundle_url")
        if not self.docker.available():
            raise Refusal("DOCKER_MISSING", f"{cfg.docker_bin} is not on PATH")
        data, hdr_sha = self.tenancy.bundle(ci.bundle_url)
        b = bundlemod.parse(data)
        if b.sha != ci.desired_sha:
            raise Refusal("BUNDLE_SHA_MISMATCH", f"computed {b.sha[:12]} != desired {ci.desired_sha[:12]}")
        if hdr_sha and hdr_sha != b.sha:
            log.warning("bundle header sha %s disagrees with content %s", _short(hdr_sha), _short(b.sha))
        self._refusal_checks(b)
        # secrets: lease -> login -> reads (values never logged; names only)
        entries = b.manifest.get("entries", [])
        if entries:
            token = self._lease_login("apply")
            for e in entries:
                self._values[e["var"]] = self.bao.kv_read(token, e["path"], cfg.env_id, e.get("version"))
            log.info("read %d secret(s): %s", len(entries), ", ".join(sorted(self._values)))
        env_text = bundlemod.render_env(b.env_template, self._values)
        leftover = bundlemod.referenced_vars(env_text)
        if leftover:
            raise Refusal("ENV_INCOMPLETE", "unfilled after render: " + ", ".join(sorted(leftover)))
        release = self._promote(self._stage(b, env_text))
        compose = self._compose()
        log.info("compose pull")
        compose.pull()
        log.info("compose up -d --remove-orphans")
        compose.up()
        ok, containers = compose.wait_healthy(cfg.healthy_timeout_s)
        if not ok:
            bad = ", ".join(f"{c.name}={c.state}/{c.health or '-'}" for c in containers) or "no containers"
            tail = compose.logs_tail(60)
            self.state.last_result = "failed"
            self.state.last_error = f"unhealthy after up: {bad}"
            log.error("not healthy within %ss: %s", cfg.healthy_timeout_s, bad)
            self._report("failed", b.sha, containers=containers, detail=f"unhealthy: {bad}", extra_tail=tail)
            return EX_FAIL
        self.state.remember_applied(b.sha)
        self.state.last_result = "applied"
        self.state.last_error = None
        self._report("applied", b.sha, containers=containers)
        log.info("applied %s (%d container(s) healthy)", _short(b.sha), len(containers))
        if self.prune:
            self._prune(release)
        return EX_OK

    # -- checks ----------------------------------------------------------------------------------------
    def _refusal_checks(self, b: Bundle) -> None:
        cfg = self.cfg
        if b.env_id and b.env_id != cfg.env_id:
            raise Refusal("BUNDLE_ENV_MISMATCH", f"bundle is for {b.env_id}, this box is {cfg.env_id}")
        if b.meta.get("renderer") not in (None, "compose"):
            raise Refusal("BUNDLE_RENDERER", f"renderer {b.meta.get('renderer')!r} is not compose")
        missing = bundlemod.unresolved_vars(b, bundlemod.manifest_vars(b))
        if missing:
            raise Refusal("ENV_INCOMPLETE", "unresolved: " + ", ".join(missing))
        for e in b.manifest.get("entries", []):
            kv_data_path(self.bao.kv_mount, e["path"], cfg.env_id)  # raises ManifestPathError
        findings = bundlemod.privileged_findings(b)
        if findings and not b.flag("allow_privileged"):
            raise Refusal("PRIVILEGED_REFUSED", ", ".join(findings) + " (bundle.json allow_privileged is not set)")
        if self.state.is_rollback(b.sha) and not b.flag("allow_rollback"):
            raise Refusal("ROLLBACK_REFUSED", f"{b.sha[:12]} was applied before; allow_rollback is not set")
        need = b.disk_need_bytes
        if need:
            free = self.docker.free_bytes()
            required = int(need * cfg.disk.factor + cfg.disk.headroom_bytes)
            if free is not None and free < required:
                raise Refusal(
                    "DISK_INSUFFICIENT",
                    f"docker data-root has {free // 2**20} MiB free, bundle needs {required // 2**20} MiB "
                    f"({need // 2**20} MiB x {cfg.disk.factor} + {cfg.disk.headroom_bytes // 2**20} MiB headroom)",
                )

    # -- credentials -------------------------------------------------------------------------------------
    def _ensure_token(self) -> str:
        if self.token:
            return self.token
        tok = statemod.read_token(self.cfg.token_path)
        if tok:
            REDACTOR.add(tok)
            self._adopt(tok, None)
            return tok
        grant = self._read_bootstrap()
        if not grant:
            raise NoCredentials(
                f"no {self.cfg.token_path} and no {self.cfg.bootstrap_path}: run the install kit "
                "(GET /v1/backends/{id}/agent/install) to get a fresh bootstrap grant"
            )
        log.info("bootstrapping: jwt-tenancy login with the install kit's deploy grant")
        bt = self.bao.login(grant)
        self._adopt(bt.client_token, bt, source="bootstrap")
        if statemod.delete(Path(self.cfg.bootstrap_path)):
            log.info("bootstrap grant consumed and deleted")
        return bt.client_token

    def _read_bootstrap(self) -> str | None:
        try:
            grant = Path(self.cfg.bootstrap_path).read_text().strip()
        except FileNotFoundError:
            return None
        if grant:
            REDACTOR.add(grant)
        return grant or None

    def _adopt(self, token: str, bt: BaoToken | None, *, source: str | None = None) -> None:
        """Make `token` the bearer: persist 0600, point the tenancy client at it, record meta."""
        old = self.token
        self.token = token
        if bt is not None or old != token:
            statemod.write_token(self.cfg.token_path, token)
        if bt is not None:
            self.state.token = statemod.TokenMeta(
                accessor=bt.accessor,
                issued_at=time.time(),
                ttl_s=bt.ttl_s,
                renewable=bt.renewable,
                grant_jti=bt.grant_jti,
                source=source,
            )
            if bt.meta.get("env_id") and bt.meta["env_id"] != self.cfg.env_id:
                raise EnvMismatch(403, "AGENT_ENV_MISMATCH", f"token env_id={bt.meta['env_id']}")
        if self.tenancy is None:
            self.tenancy = TenancyClient(self.tenancy_http, self.cfg.env_id, token)
        else:
            self.tenancy.token = token
        if old and old != token:
            self.bao.revoke_self(old)

    def _lease_login(self, reason: str) -> str:
        assert self.tenancy is not None
        lease = self.tenancy.lease(reason)
        self._releases_this_tick += 1
        bt = self.bao.login(lease.login_jwt, mount=lease.mount, role=lease.role)
        log.info("lease: jwt-tenancy login ok (ttl %ss, accessor %s)", bt.ttl_s, bt.accessor)
        self._adopt(bt.client_token, bt, source="lease")
        return bt.client_token

    def _maintain_token(self) -> None:
        """Renew-self when the TTL runs low; re-lease before max TTL or when renewal is refused."""
        if not self.token or self.tenancy is None or self._releases_this_tick:
            return
        cfg = self.cfg
        try:
            info = self.bao.lookup_self(self.token)
        except BaoError as exc:
            if exc.status in (403, 400):
                log.warning("stored token is dead (%s); dropping it", exc)
                statemod.delete(cfg.token_path)
                self.state.token = statemod.TokenMeta()
                self.token = None
            return
        ttl = int(info.get("ttl") or 0)
        renewable = bool(info.get("renewable"))
        creation = info.get("creation_time")
        age = time.time() - float(creation) if creation else (time.time() - (self.state.token.issued_at or time.time()))
        max_ttl = int(info.get("explicit_max_ttl") or 0) or cfg.bao.token_max_ttl_s
        near_max = age > max_ttl - cfg.bao.release_before_s
        last_lease = self.state.token.issued_at or 0
        if near_max or (ttl < cfg.bao.renew_below_s and not renewable):
            if time.time() - last_lease < LEASE_MIN_GAP_S:
                return
            log.info("token age %dh of max %dh: re-leasing", age // 3600, max_ttl // 3600)
            self._lease_login("renew")
            return
        if ttl < cfg.bao.renew_below_s and renewable:
            new_ttl = self.bao.renew_self(self.token)
            log.info("token renewed (ttl %ss -> %ss)", ttl, new_ttl)
            self.state.token.ttl_s = new_ttl
            if new_ttl < cfg.bao.renew_below_s and time.time() - last_lease >= LEASE_MIN_GAP_S:
                log.info("renewal capped by max TTL: re-leasing")
                self._lease_login("renew")

    # -- files -------------------------------------------------------------------------------------------
    def _stage(self, b: Bundle, env_text: str) -> Path:
        deploy = Path(self.cfg.deploy_dir)
        deploy.mkdir(parents=True, exist_ok=True, mode=0o700)
        staging = deploy / f"staging-{b.sha[:12]}-{os.getpid()}"
        if staging.exists():
            shutil.rmtree(staging)
        staging.mkdir(mode=0o700)
        self._staging.append(staging)
        for name, content in b.files.items():
            if name == bundlemod.ENV_TEMPLATE:
                continue
            target = staging / name
            target.parent.mkdir(parents=True, exist_ok=True)
            statemod.write_private(target, content.encode())
            os.chmod(target, 0o600)
        statemod.write_private(staging / ".env", env_text.encode())
        (staging / "SHA256").write_text(b.sha + "\n")
        return staging

    def _promote(self, staging: Path) -> Path:
        deploy = Path(self.cfg.deploy_dir)
        releases = deploy / "releases"
        releases.mkdir(exist_ok=True, mode=0o700)
        sha = (staging / "SHA256").read_text().strip()
        release = releases / sha
        if release.exists():
            shutil.rmtree(release)
        os.replace(staging, release)
        self._staging.remove(staging)
        tmp_link = deploy / ".current.tmp"
        if tmp_link.is_symlink() or tmp_link.exists():
            tmp_link.unlink()
        os.symlink(release, tmp_link)
        os.replace(tmp_link, deploy / "current")
        log.info("promoted release %s", _short(sha))
        return release

    def _prune(self, keep: Path) -> None:
        releases = Path(self.cfg.deploy_dir) / "releases"
        for d in releases.iterdir() if releases.exists() else []:
            if d.resolve() != keep.resolve() and d.is_dir():
                shutil.rmtree(d, ignore_errors=True)
                log.info("pruned release %s", d.name[:12])
        self.docker.image_prune()

    def _compose(self) -> Compose:
        return Compose(self.docker, Path(self.cfg.deploy_dir) / "current", self.cfg.compose_project)

    # -- misc ----------------------------------------------------------------------------------------------
    def _report(
        self,
        result: str,
        sha: str | None,
        *,
        containers: list[Container] | None = None,
        detail: str | None = None,
        extra_tail: str = "",
    ) -> None:
        if self.tenancy is None:
            return
        tail = self.tail.text()
        if extra_tail:
            tail = tail + "\n--- compose logs ---\n" + REDACTOR.scrub(extra_tail)
        try:
            self.tenancy.report(
                sha=sha,
                result=result,
                log_tail=tail,
                containers=[c.as_dict() for c in containers] if containers else None,
                detail=REDACTOR.scrub(detail) if detail else None,
            )
            self.state.last_report_at = time.time()
            self.state.last_result = result
        except (TenancyError, TransportError) as exc:
            log.warning("report %s not delivered: %s", result, exc)

    def _host_facts(self) -> dict:
        os_id = "unknown"
        try:
            kv = dict(line.split("=", 1) for line in Path("/etc/os-release").read_text().splitlines() if "=" in line)
            os_id = f"{kv.get('ID', '?').strip(chr(34))}-{kv.get('VERSION_ID', '?').strip(chr(34))}"
        except OSError:
            pass
        facts = {
            "hostname": socket.gethostname(),
            "os": os_id,
            "kernel": platform.release(),
            "agent_version": __version__,
            "transport": self.cfg.transport,
            "site": self.cfg.site,
            "laptop_mode": self.cfg.laptop_mode,
        }
        if self.docker.available():
            facts["docker"] = self.docker.version()
            facts["compose"] = self.docker.compose_version()
            facts["free_bytes"] = self.docker.free_bytes()
        return facts

    def _scrub(self) -> None:
        for k in list(self._values):
            self._values[k] = ""
        self._values.clear()
        for d in self._staging:
            shutil.rmtree(d, ignore_errors=True)
        self._staging.clear()


def _short(sha: str | None) -> str:
    return sha[:12] if sha else "-"