"""Commands that outlive the reply that started them.

An ordinary `shell_run` is one blocking `conn.run` over a per-call connection
(`ssh.py`): when it hits its timeout the command is killed, so a ten-minute
`apt install` is impossible. A background job is the same command launched
*detached* on the far side -- `setsid`, redirected to a remote logfile and an
exit-file -- so it survives the connection closing. LLeMbas reconnects (a fresh
connection, as always) to read the log and the exit code later.

This is the opposite of `terminal.py`, which survives by *holding* a connection
open. Here we hold nothing: the whole point of `ssh.py`/`base.py` is that no live
connection is kept, and a job that needed one would be a job that broke that.

**The command never touches a quoted shell context.** `sh -c '<cmd>'` shatters
the instant the command contains a `'` -- `git commit -m 'fix'`, `awk '{…}'`,
`sed 's/…/…/'` are the common case, not an edge one, and would also be an
injection hole. So the command is base64-encoded here in Python and decoded on
the far side into a script file; it is bytes, never shell syntax. Only
server-generated hex ids and a fixed root ever reach a path.

Three things make the wrappers correct, and each was got wrong in an earlier
sketch:

* **The child records its own pid via `$$`**, as its first act, under `setsid`
  where it is the session/group leader -- so `job_stop` can `kill -<pid>` the
  whole process group. `echo $!` from the launcher captures the wrong pid.
* **The exit-file is the primary signal.** An empty pid-file means "still
  starting", not "dead"; reading liveness first would race the launch and report
  a job lost the instant it began.
* **The command's exit status comes from the exit-file, never from the wrapper's
  own status** -- which is ~0 from the trailing `rm`. Reading the wrapper's
  status would mark every job a success.
"""

from __future__ import annotations

import asyncio
import base64
import contextlib
import logging
import re
import time
import uuid
from dataclasses import dataclass, field

from lembas.services.agent.base import ExecError, ExecRequest, clean_output

log = logging.getLogger(__name__)

# Where a job's files live on the far side. `${TMPDIR:-/tmp}` so a host that
# puts scratch space elsewhere is honoured, and it clears on reboot -- a job
# does not survive a reboot of its own host either. The chat id namespaces it,
# which is also what makes cross-chat access structurally impossible: a path is
# only ever built from the *calling* chat's id, so a model in one chat cannot
# name another chat's files.
JOB_ROOT = "${TMPDIR:-/tmp}/lembas-jobs"

# A job id is our own short hex; anything else is refused before it reaches a
# path, so `job_output("../../etc/passwd")` cannot walk out of the job root.
_ID = re.compile(r"^[a-f0-9]{12}$")

# How long the fire-and-return launcher waits for the shell to accept the
# command. Not the command's own timeout -- it returns the moment the process is
# detached, which is immediate.
LAUNCH_GRACE = 10.0

# The working set, keyed by job id: what `job_list` shows this session. Mirrored
# to a `Job` row for jobs that are watched, so a restart can rehydrate them.
_JOBS: dict[str, JobState] = {}
# One watcher task per job being polled to completion.
_WATCHERS: dict[str, asyncio.Task] = {}
# One lock per chat, so two jobs finishing at once cannot each start a reply --
# see `wake`.
_WAKE_LOCKS: dict[str, asyncio.Lock] = {}

# Stop watching a job after this. The remote process may keep running; we simply
# stop holding a watcher for it and mark it lost. A job that runs longer than
# this is beyond what auto-wake promises.
MAX_WATCH_SECONDS = 6 * 3600

# How much of a finished job's output is put in front of the model when it is
# woken. Capped so a job that printed a gigabyte does not blow the window.
MAX_COMPLETION_CHARS = 4000


def new_id() -> str:
    return uuid.uuid4().hex[:12]


@dataclass
class JobState:
    """What LLeMbas remembers about one background job, in this process."""

    id: str
    chat_id: str
    command: str
    status: str = "running"  # running | done | killed | lost
    exit_status: int | None = None
    started_at: float = field(default_factory=time.monotonic)
    finished_at: float = 0.0


# --- Paths and the wrappers ----------------------------------------------------
def _dir(chat_id: str) -> str:
    return f'"{JOB_ROOT}/{chat_id}"'


def _file(chat_id: str, job_id: str, ext: str) -> str:
    # Double-quoted so `${TMPDIR:-/tmp}` still expands while the whole path stays
    # one word. The chat id and job id are hex, so nothing here needs escaping.
    return f'"{JOB_ROOT}/{chat_id}/{job_id}.{ext}"'


def _sentinel(job_id: str) -> str:
    return f"__LEMBAS_{job_id}__"


def _inner_script(chat_id: str, job_id: str, command: str) -> str:
    """The detached program: record the pid, run the command, record the status.

    base64-encoded before it leaves, so `command` is bytes and never shell
    syntax. `$$` first, because it is the session leader's pid under setsid and
    `job_stop` kills the group by it. `$?` last, capturing the command's status;
    it is the file `run`'s own exit status must never be read in place of.
    """
    return (
        f"echo $$ > {_file(chat_id, job_id, 'pid')}\n"
        f"{command}\n"
        f"echo $? > {_file(chat_id, job_id, 'exit')}\n"
    )


def _blob(chat_id: str, job_id: str, command: str) -> str:
    raw = _inner_script(chat_id, job_id, command).encode("utf-8")
    return base64.b64encode(raw).decode("ascii")


def _launch_lines(chat_id: str, job_id: str, command: str) -> str:
    """Create the job dir, drop the script, and detach it. No wait."""
    blob = _blob(chat_id, job_id, command)
    return (
        f"mkdir -p {_dir(chat_id)} 2>/dev/null\n"
        f"printf %s '{blob}' | base64 -d > {_file(chat_id, job_id, 'sh')}\n"
        f"setsid sh {_file(chat_id, job_id, 'sh')} "
        f"> {_file(chat_id, job_id, 'log')} 2>&1 < /dev/null &\n"
    )


def launch_command(chat_id: str, job_id: str, command: str) -> str:
    """Fire-and-return: detach the command and stop. Run with a short timeout."""
    return _launch_lines(chat_id, job_id, command) + "printf started\n"


def launch_and_wait_command(chat_id: str, job_id: str, command: str, max_bytes: int) -> str:
    """Detach the command AND wait up to the (asyncssh) timeout for it.

    If it finishes, stdout is the log tail plus a sentinel line carrying the exit
    code, and the files are removed. If asyncssh times out first the channel is
    torn down before the `rm`, so the files survive for a later read and the
    detached process -- new session, redirected, stdin from /dev/null -- keeps
    running. That torn-down-mid-wait case is exactly "it became a background
    job".
    """
    s = _sentinel(job_id)
    pid = _file(chat_id, job_id, "pid")
    exit_ = _file(chat_id, job_id, "exit")
    logf = _file(chat_id, job_id, "log")
    return (
        _launch_lines(chat_id, job_id, command)
        + "while :; do\n"
        f"  [ -f {exit_} ] && break\n"
        f"  __p=$(cat {pid} 2>/dev/null)\n"
        '  [ -n "$__p" ] && ! kill -0 "$__p" 2>/dev/null && break\n'
        # 0.2s: with the feature on, every ordinary command waits one poll for
        # the exit-file, so this is added latency on the hot path. Short enough
        # not to be felt, long enough not to spin.
        "  sleep 0.2\n"
        "done\n"
        f"tail -c {max_bytes} {logf} 2>/dev/null\n"
        f"printf '\\n{s}:'\n"
        f"cat {exit_} 2>/dev/null || printf LOST\n"
        f"rm -f {_file(chat_id, job_id, 'sh')} {pid} {log} {exit_}\n"
    )


def read_command(chat_id: str, job_id: str, max_bytes: int) -> str:
    """The log so far, and whether the job is still running."""
    s = _sentinel(job_id)
    pid = _file(chat_id, job_id, "pid")
    exit_ = _file(chat_id, job_id, "exit")
    return (
        f"tail -c {max_bytes} {_file(chat_id, job_id, 'log')} 2>/dev/null\n"
        f"printf '\\n{s}:'\n"
        f"if [ -f {exit_} ]; then printf 'done '; cat {exit_};\n"
        f'elif __p=$(cat {pid} 2>/dev/null); [ -n "$__p" ] && kill -0 "$__p" 2>/dev/null;'
        " then printf running;\n"
        "else printf lost; fi\n"
    )


def stop_command(chat_id: str, job_id: str) -> str:
    """Kill the whole process group, then record an exit so a reader is not told
    the job is merely lost. A killed process never writes its own exit file."""
    pid = _file(chat_id, job_id, "pid")
    exit_ = _file(chat_id, job_id, "exit")
    return (
        f'__p=$(cat {pid} 2>/dev/null); [ -n "$__p" ] && kill -TERM -"$__p" 2>/dev/null\n'
        "sleep 0.3\n"
        f'[ -n "$__p" ] && kill -KILL -"$__p" 2>/dev/null\n'
        f"[ -f {exit_} ] || echo 143 > {exit_}\n"
        "printf stopped\n"
    )


def cleanup_command(chat_id: str, job_id: str) -> str:
    return (
        f"rm -f {_file(chat_id, job_id, 'sh')} {_file(chat_id, job_id, 'pid')} "
        f"{_file(chat_id, job_id, 'log')} {_file(chat_id, job_id, 'exit')}\n"
    )


# --- Parsing what a wrapper printed --------------------------------------------
@dataclass(frozen=True)
class Completed:
    body: str
    exit_status: int | None  # None ⇒ the job was lost (killed without an exit)


def parse_completed(output: str, job_id: str) -> Completed:
    """Split a launch-and-wait result into the command's output and its status.

    On the *last* sentinel, because the command's own output could contain a
    line that looks like one; everything before it is the body, everything after
    is the exit code the file held.
    """
    marker = f"\n{_sentinel(job_id)}:"
    at = output.rfind(marker)
    if at == -1:
        return Completed(body=output.strip(), exit_status=None)
    body = output[:at].strip()
    tail = output[at + len(marker) :].strip()
    if tail.upper() == "LOST" or not tail:
        return Completed(body=body, exit_status=None)
    try:
        return Completed(body=body, exit_status=int(tail.split()[0]))
    except (ValueError, IndexError):
        return Completed(body=body, exit_status=None)


@dataclass(frozen=True)
class Reading:
    body: str
    status: str  # running | done | lost
    exit_status: int | None


def parse_reading(output: str, job_id: str) -> Reading:
    marker = f"\n{_sentinel(job_id)}:"
    at = output.rfind(marker)
    if at == -1:
        return Reading(body=output.strip(), status="lost", exit_status=None)
    body = output[:at].strip()
    tail = output[at + len(marker) :].strip()
    if tail.startswith("done"):
        parts = tail.split()
        code = int(parts[1]) if len(parts) > 1 and parts[1].lstrip("-").isdigit() else None
        return Reading(body=body, status="done", exit_status=code)
    if tail == "running":
        return Reading(body=body, status="running", exit_status=None)
    return Reading(body=body, status="lost", exit_status=None)


# --- Operations against the machine --------------------------------------------
async def launch(agent, command: str, cwd: str = "") -> JobState:
    """Detach a command and return immediately. Raises ExecError if it will not
    even start."""
    job_id = new_id()
    result = await agent.executor().run(
        ExecRequest(
            command=launch_command(agent.chat_id, job_id, command),
            cwd=cwd,
            timeout=LAUNCH_GRACE,
            max_bytes=agent.max_output,
        )
    )
    if result.timed_out:
        raise ExecError("The machine did not accept the command in time.")
    job = JobState(id=job_id, chat_id=agent.chat_id, command=command)
    _JOBS[job_id] = job
    return job


async def read(agent, job_id: str) -> Reading:
    output, _ = _clean(
        await agent.executor().run(
            ExecRequest(
                command=read_command(agent.chat_id, job_id, agent.max_output),
                timeout=agent.timeout,
                max_bytes=agent.max_output,
            )
        ),
        agent.max_output,
    )
    reading = parse_reading(output, job_id)
    _record(job_id, reading.status, reading.exit_status)
    if reading.status in ("done", "lost"):
        await _cleanup(agent, job_id)
    return reading


async def stop(agent, job_id: str) -> None:
    await agent.executor().run(
        ExecRequest(command=stop_command(agent.chat_id, job_id), timeout=agent.timeout)
    )
    _record(job_id, "killed", 143)


async def _cleanup(agent, job_id: str) -> None:
    with contextlib.suppress(ExecError):
        await agent.executor().run(
            ExecRequest(command=cleanup_command(agent.chat_id, job_id), timeout=agent.timeout)
        )


def _clean(result, limit: int) -> tuple[str, bool]:
    if result.timed_out:
        return result.output, False
    return clean_output(result.output or "", limit=limit)


# --- The registry --------------------------------------------------------------
def register(job: JobState) -> None:
    _JOBS[job.id] = job


def get(job_id: str) -> JobState | None:
    return _JOBS.get(job_id)


def for_chat(chat_id: str) -> list[JobState]:
    return [j for j in _JOBS.values() if j.chat_id == chat_id]


def valid_id(job_id: str) -> bool:
    return bool(_ID.match(job_id or ""))


def _record(job_id: str, status: str, exit_status: int | None) -> None:
    job = _JOBS.get(job_id)
    if job is None or job.status != "running":
        return
    if status in ("done", "lost", "killed"):
        job.status = status
        job.exit_status = exit_status
        job.finished_at = time.monotonic()
        _persist_row(job)


def clear() -> None:
    _JOBS.clear()


# --- Durable record ------------------------------------------------------------
# Best-effort throughout: a job whose row cannot be written (a test with no real
# chat, a transient database hiccup) still runs and is still tracked in-process;
# it just will not survive a restart, which is the row's only purpose.
def _persist_row(job: JobState) -> None:
    from datetime import UTC, datetime

    from lembas.db.models import Job
    from lembas.db.session import session_scope

    try:
        with session_scope() as db:
            row = db.get(Job, job.id)
            if row is None:
                row = Job(id=job.id, chat_id=job.chat_id)
                db.add(row)
            row.command = job.command[:4000]
            row.status = job.status
            row.exit_status = job.exit_status
            row.finished_at = None if job.status == "running" else datetime.now(UTC)
    except Exception:  # noqa: BLE001 - the row is a convenience, not the job
        log.debug("could not persist job %s", job.id, exc_info=True)


# --- The watcher ---------------------------------------------------------------
def _poll_interval(elapsed: float) -> float:
    if elapsed < 30:
        return 3.0
    if elapsed < 300:
        return 10.0
    return 25.0


def start_watch(agent, job: JobState) -> None:
    """Poll a job to completion and, when it finishes, wake the model.

    Only when notify is on -- the watcher's whole job is the wake and the status
    update, and without notify the model reads `job_output` itself, which
    updates the status anyway. Capped by `background_max_jobs`: past it a job
    still runs and can be read, it simply is not watched.

    The credential is copied, not referenced: `generation` clears the agent's
    `spec` when the reply ends, and the watcher outlives the reply. Holding the
    copy for the job's life is the same trade the terminal makes for a held
    shell.
    """
    _persist_row(job)
    if not agent.background_notify or len(_WATCHERS) >= agent.background_max_jobs:
        return
    task = asyncio.create_task(
        _watch(
            dict(agent.spec),
            agent.project_dir,
            job.chat_id,
            job.id,
            job.command,
            agent.max_output,
        )
    )
    _WATCHERS[job.id] = task


async def _watch(
    spec: dict, project_dir: str, chat_id: str, job_id: str, command: str, max_output: int
) -> None:
    from lembas.services.agent.ssh import SshExecutor

    started = time.monotonic()
    try:
        while True:
            await asyncio.sleep(_poll_interval(time.monotonic() - started))
            if time.monotonic() - started > MAX_WATCH_SECONDS:
                _record(job_id, "lost", None)
                return
            try:
                result = await SshExecutor(spec, project_dir).run(
                    ExecRequest(
                        command=read_command(chat_id, job_id, max_output),
                        timeout=30,
                        max_bytes=max_output,
                    )
                )
            except ExecError:
                continue  # transient -- the host is briefly unreachable; retry
            if result.timed_out:
                continue
            output, _ = clean_output(result.output or "", limit=max_output)
            reading = parse_reading(output, job_id)
            if reading.status in ("done", "lost"):
                _record(job_id, reading.status, reading.exit_status)
                with contextlib.suppress(ExecError):
                    await SshExecutor(spec, project_dir).run(
                        ExecRequest(command=cleanup_command(chat_id, job_id), timeout=30)
                    )
                await wake(chat_id, job_id, command, reading.status, reading.exit_status,
                           reading.body)
                return
    except asyncio.CancelledError:
        raise
    except Exception:  # noqa: BLE001 - a watcher that dies must not take others
        log.exception("job watcher for %s raised", job_id)
    finally:
        _WATCHERS.pop(job_id, None)


# --- Waking the model ----------------------------------------------------------
def _lock(chat_id: str) -> asyncio.Lock:
    lock = _WAKE_LOCKS.get(chat_id)
    if lock is None:
        lock = _WAKE_LOCKS[chat_id] = asyncio.Lock()
    return lock


def _completion_text(
    job_id: str, command: str, status: str, exit_status: int | None, output: str
) -> str:
    if status == "done" and exit_status == 0:
        line = "It finished successfully."
    elif status == "done":
        line = f"It exited {exit_status}."
    else:
        line = "It stopped without an exit status (it may have been killed)."
    body = (output or "").strip()[:MAX_COMPLETION_CHARS]
    # A fence for the model's benefit; backticks in the output are neutralised so
    # they cannot close it, the same move `instructions.clean` makes.
    fenced = f"\n\n```\n{body.replace('```', chr(39) * 3)}\n```" if body else ""
    return (
        f"A background job you started has finished — this is a machine event, "
        f"not the person speaking.\n\n"
        f"[job {job_id}] `{command}`\n{line}{fenced}"
    )


async def wake(
    chat_id: str, job_id: str, command: str, status: str, exit_status: int | None, output: str
) -> None:
    """Tell the model a job finished, as a new turn.

    Reuses the queue: if a reply is being written, the completion is left
    `queued` for that reply's `_inject`/`_drain` to deliver; if the chat is idle,
    a fresh reply is started to answer it, the `send_queued_now` move.

    The whole thing is under a per-chat lock, and there is no `await` between the
    running-check and starting the reply, so two jobs finishing at once cannot
    each spin up a generation -- the second sees the first's reply already live
    and leaves its completion for it. That is the invariant the queue exists to
    hold, reached here from outside a request.

    The completion is a user-role turn whose *content* names itself a machine
    event -- `_inject` sends a queued turn verbatim, so the framing cannot live
    there; it lives in the words, the way `execute_plan` quotes the plan.
    """
    from lembas.db.models import ROLE_ASSISTANT, ROLE_USER, Chat
    from lembas.db.session import session_scope
    from lembas.services import chat as chat_service
    from lembas.services import generation as generation_service

    content = _completion_text(job_id, command, status, exit_status, output)
    async with _lock(chat_id):
        running = generation_service.running_for(chat_id) is not None
        assistant_id = ""
        try:
            with session_scope() as db:
                chat = db.get(Chat, chat_id)
                if chat is None:
                    return
                chat_service.create_message(db, chat, ROLE_USER, content, queued=running)
                if not running:
                    assistant = chat_service.create_message(
                        db, chat, ROLE_ASSISTANT, "", complete_=False, model_id=chat.model_id
                    )
                    assistant_id = assistant.id
        except Exception:  # noqa: BLE001 - a failed wake must not crash the watcher
            log.exception("could not wake chat %s for job %s", chat_id, job_id)
            return
        if assistant_id:
            generation_service.ensure(chat_id, assistant_id)


# --- Rehydration and shutdown --------------------------------------------------
def rehydrate() -> None:
    """After a restart, watch again the jobs that were still running.

    Their remote files are keyed deterministically on chat and id, so a fresh
    watcher re-polls them and wakes the model as if nothing happened -- which is
    the whole reason the row exists. Best-effort per job: a host that is down, a
    profile that is gone, a chat that was deleted each just drop that one.
    """
    from lembas.db.models import Chat, SshProfile
    from lembas.db.session import session_scope
    from lembas.services import settings_store
    from lembas.services.agent import ssh as ssh_service

    with session_scope() as db:
        values = settings_store.agents(db)
        if not values.get("enabled") or not values.get("background_notify"):
            return
        max_output = int(values.get("max_output_bytes") or 64 * 1024)
        running = list(db.scalars(_running_rows()))
        for row in running:
            chat = db.get(Chat, row.chat_id)
            if chat is None or not chat.ssh_profile_id:
                continue
            profile = db.get(SshProfile, chat.ssh_profile_id)
            if profile is None or not profile.enabled:
                continue
            spec = ssh_service.spec_from(profile)
            project_dir = chat.project_dir or profile.default_dir or ""
            job = JobState(id=row.id, chat_id=row.chat_id, command=row.command)
            _JOBS[job.id] = job
            if len(_WATCHERS) >= int(values.get("background_max_jobs") or 5):
                break
            _WATCHERS[job.id] = asyncio.create_task(
                _watch(spec, project_dir, row.chat_id, row.id, row.command, max_output)
            )


def _running_rows():
    from sqlalchemy import select

    from lembas.db.models import Job

    return select(Job).where(Job.status == "running")


async def shutdown() -> None:
    """Cancel every watcher. The detached remote jobs are unaffected -- they run
    on, and a later start rehydrates them from their rows."""
    tasks = list(_WATCHERS.values())
    _WATCHERS.clear()
    for task in tasks:
        task.cancel()
    for task in tasks:
        with contextlib.suppress(asyncio.CancelledError, Exception):
            await task


__all__ = [
    "JOB_ROOT",
    "Completed",
    "JobState",
    "Reading",
    "clear",
    "for_chat",
    "get",
    "launch",
    "launch_and_wait_command",
    "new_id",
    "parse_completed",
    "read",
    "register",
    "stop",
    "valid_id",
]