# steward/ansible/executor.py
from __future__ import annotations
import asyncio
import json
import logging
import os
import re
import shutil
import tempfile
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

_OUTPUT_CAP_BYTES = 1024 * 1024  # 1 MB DB cap (full output also goes to a log artifact)
_FLUSH_LINES = 50
_FLUSH_SECS = 5.0
_MEM_LINES_CAP = 5000   # bounded per-run in-memory replay buffer (SSE late-joiners)
_MEM_TTL_SECS = 300     # drop in-memory state this long after a run completes
_FAILURE_CAP = 50       # max failed-task lines captured into structured results

# Full run logs persist here (survive restart + the 1 MB DB cap). Override via env.
ARTIFACT_DIR = os.environ.get("STEWARD_ANSIBLE_LOG_DIR", "/data/ansible/runs")


def artifact_path(run_id: str) -> str:
    return os.path.join(ARTIFACT_DIR, f"{run_id}.log")


@dataclass
class _Done:
    status: str


class _Aborted(Exception):
    """Raised internally when a queued run is cancelled before it launches."""


# Per-run broadcast state (lives until shortly after completion, then GC'd)
_run_lines: dict[str, list[str]] = {}          # bounded tail of output lines
_run_done: dict[str, _Done] = {}               # set when run completes
_run_listeners: dict[str, list[asyncio.Queue]] = {}  # active SSE client queues

# Running subprocesses by run_id (for cancellation) + run_ids an operator cancelled.
_run_procs: dict[str, asyncio.subprocess.Process] = {}
_cancelled: set[str] = set()


def register_listener(run_id: str) -> asyncio.Queue:
    """Create a listener queue for a run. Replays existing lines immediately."""
    q: asyncio.Queue = asyncio.Queue()
    for line in _run_lines.get(run_id, []):
        q.put_nowait(line)
    done = _run_done.get(run_id)
    if done:
        q.put_nowait(done)
    else:
        _run_listeners.setdefault(run_id, []).append(q)
    return q


def deregister_listener(run_id: str, q: asyncio.Queue) -> None:
    listeners = _run_listeners.get(run_id, [])
    if q in listeners:
        listeners.remove(q)


def _schedule_cleanup(run_id: str) -> None:
    """Drop a finished run's in-memory state after a grace window."""
    async def _cleanup() -> None:
        await asyncio.sleep(_MEM_TTL_SECS)
        _run_lines.pop(run_id, None)
        _run_done.pop(run_id, None)
    try:
        asyncio.create_task(_cleanup())
    except RuntimeError:
        pass  # no running loop (e.g. tests) — fine to skip


def _broadcast(run_id: str, item: str | _Done) -> None:
    if isinstance(item, str):
        buf = _run_lines.setdefault(run_id, [])
        buf.append(item)
        if len(buf) > _MEM_LINES_CAP:  # keep only the most recent lines in memory
            del buf[: len(buf) - _MEM_LINES_CAP]
    else:
        _run_done[run_id] = item
    for q in _run_listeners.get(run_id, []):
        q.put_nowait(item)
    if isinstance(item, _Done):
        _run_listeners.pop(run_id, None)
        _schedule_cleanup(run_id)


def build_ansible_command(
    playbook_path: str,
    inventory_path: str,
    params: dict | None = None,
) -> list[str]:
    """Build the ansible-playbook argv from a run's optional params.

    Pure (no IO). Everything is passed as argv — never shell-interpolated — so
    extra-var values and limits can contain arbitrary characters safely.
    params keys: extra_vars (list of "key=value"), limit, tags, check (bool).
    """
    cmd = ["ansible-playbook", playbook_path, "-i", inventory_path]
    params = params or {}
    for ev in params.get("extra_vars") or []:
        cmd += ["-e", ev]
    if params.get("limit"):
        cmd += ["--limit", params["limit"]]
    if params.get("tags"):
        cmd += ["--tags", params["tags"]]
    if params.get("check"):
        cmd += ["--check", "--diff"]
    return cmd


def build_credentials(creds: dict | None, tmpdir: str) -> tuple[list[str], list[tuple[str, str]]]:
    """Given global Ansible creds + a temp dir, return (extra argv, files to write).

    Pure: decides flags + file contents; the caller writes the files (mode 0600).
    Secrets are always passed via files — never on argv / the process list.
    creds keys: ssh_private_key, vault_password, become_password.
    """
    args: list[str] = []
    files: list[tuple[str, str]] = []
    creds = creds or {}

    key = creds.get("ssh_private_key")
    if key:
        path = os.path.join(tmpdir, "id_key")
        files.append((path, key if key.endswith("\n") else key + "\n"))
        args += ["--private-key", path]

    vault = creds.get("vault_password")
    if vault:
        path = os.path.join(tmpdir, "vault_pass")
        files.append((path, vault + "\n"))
        args += ["--vault-password-file", path]

    become = creds.get("become_password")
    if become:
        path = os.path.join(tmpdir, "become.yml")
        # A JSON string is valid YAML; keeps the password out of argv.
        files.append((path, "ansible_become_password: " + json.dumps(become) + "\n"))
        args += ["-e", "@" + path]

    return args, files


def ansible_env(creds: dict | None, base_env) -> dict:
    """Return a copy of base_env with ANSIBLE_HOST_KEY_CHECKING set from creds."""
    env = dict(base_env)
    env["ANSIBLE_HOST_KEY_CHECKING"] = "True" if (creds or {}).get("host_key_checking") else "False"
    return env


# ── Structured results (parsed from the streamed output) ──────────────────────

_RECAP_RE = re.compile(
    r"^(?P<host>\S+)\s*:\s*ok=(?P<ok>\d+)\s+changed=(?P<changed>\d+)\s+"
    r"unreachable=(?P<unreachable>\d+)\s+failed=(?P<failed>\d+)"
    r"(?:\s+skipped=(?P<skipped>\d+))?"
)


def parse_recap(lines: list[str]) -> dict[str, dict[str, int]]:
    """Parse Ansible's PLAY RECAP into {host: {ok,changed,unreachable,failed,skipped}}.

    Streaming-friendly: we keep the default stdout callback (live output) and
    derive the per-host summary from the recap rather than swapping to the json
    callback (which would buffer everything to the end and kill live streaming).
    """
    hosts: dict[str, dict[str, int]] = {}
    in_recap = False
    for line in lines:
        if "PLAY RECAP" in line:
            in_recap = True
            continue
        if not in_recap:
            continue
        m = _RECAP_RE.match(line.strip())
        if m:
            hosts[m.group("host")] = {
                "ok": int(m.group("ok")),
                "changed": int(m.group("changed")),
                "unreachable": int(m.group("unreachable")),
                "failed": int(m.group("failed")),
                "skipped": int(m.group("skipped") or 0),
            }
    return hosts


# ── Cancellation ──────────────────────────────────────────────────────────────

async def _set_status(app: "Quart", run_id: str, status) -> None:
    from sqlalchemy import update
    from steward.models.ansible import AnsibleRun
    async with app.db_sessionmaker() as session:
        async with session.begin():
            await session.execute(
                update(AnsibleRun).where(AnsibleRun.id == run_id).values(status=status)
            )


async def _terminate_then_kill(proc: asyncio.subprocess.Process, grace: float = 10.0) -> None:
    try:
        proc.terminate()
    except ProcessLookupError:
        return
    try:
        await asyncio.wait_for(proc.wait(), timeout=grace)
    except asyncio.TimeoutError:
        try:
            proc.kill()
        except ProcessLookupError:
            pass


def cancel_run(run_id: str) -> bool:
    """Request cancellation of a run. Terminates the process if it's running;
    a queued run is flagged so start_run aborts before launch. Returns False
    only when the run isn't tracked in this process (already finished)."""
    if run_id not in _run_procs and run_id not in _cancelled:
        # Not obviously in-flight. Still flag it (covers a just-queued run whose
        # proc hasn't been registered yet); caller should gate on run status.
        _cancelled.add(run_id)
        return run_id in _run_procs or True
    _cancelled.add(run_id)
    proc = _run_procs.get(run_id)
    if proc is not None:
        asyncio.create_task(_terminate_then_kill(proc))
    return True


# ── Run execution ─────────────────────────────────────────────────────────────

async def start_run(
    app: "Quart",
    run_id: str,
    playbook_path: str,
    inventory_path: str,
    source_path: str,
    params: dict | None = None,
    inventory_content: str | None = None,
) -> None:
    """Execute ansible-playbook as a subprocess and update the DB run row.

    Respects a global concurrency cap (app._ansible_semaphore): if no slot is
    free the run shows as 'queued' until one frees. Supports cancellation, a
    persistent full-log artifact, and a parsed per-host result summary.
    """
    from steward.models.ansible import AnsibleRunStatus

    _run_lines[run_id] = []
    db_output = ""
    truncated = False
    lines_since_flush = 0
    last_flush_at = time.monotonic()
    failures: list[str] = []

    async def _flush(final_status: str | None = None, results: dict | None = None) -> None:
        nonlocal lines_since_flush, last_flush_at
        from sqlalchemy import update
        from steward.models.ansible import AnsibleRun

        values: dict = {"output": db_output}
        if results is not None:
            values["results"] = results
        if final_status is not None:
            values["status"] = AnsibleRunStatus(final_status)
            values["finished_at"] = datetime.now(timezone.utc)
        async with app.db_sessionmaker() as session:
            async with session.begin():
                await session.execute(
                    update(AnsibleRun).where(AnsibleRun.id == run_id).values(**values)
                )
        lines_since_flush = 0
        last_flush_at = time.monotonic()

    # ── Concurrency gate (semaphore lazily bound to the running loop so the
    # cap is honoured in prod's single loop without breaking multi-loop tests).
    loop = asyncio.get_running_loop()
    sem = getattr(app, "_ansible_semaphore", None)
    if sem is None or getattr(app, "_ansible_sem_loop", None) is not loop:
        cap = int((app.config.get("ANSIBLE") or {}).get("max_concurrent_runs", 3) or 3)
        sem = asyncio.Semaphore(max(1, cap))
        app._ansible_semaphore = sem
        app._ansible_sem_loop = loop
    queued = sem.locked()
    if queued:
        await _set_status(app, run_id, AnsibleRunStatus.queued)
        _broadcast(run_id, "[queued — waiting for an Ansible run slot]")
    await sem.acquire()

    final_status = "failed"
    logf = None
    try:
        if run_id in _cancelled:  # cancelled while queued — abort before launching
            raise _Aborted
        if queued:
            await _set_status(app, run_id, AnsibleRunStatus.running)

        cwd = source_path if Path(source_path).exists() else None
        creds = app.config.get("ANSIBLE", {})
        try:
            os.makedirs(ARTIFACT_DIR, exist_ok=True)
            logf = open(artifact_path(run_id), "w", encoding="utf-8")
        except OSError:
            logf = None  # artifact is best-effort; the DB still holds capped output

        # Temp dir holds the ephemeral inventory (if any) + credential files.
        tmpdir = tempfile.mkdtemp(prefix="steward-ansible-")
        os.chmod(tmpdir, 0o700)
        try:
            if inventory_content:
                inv_target = os.path.join(tmpdir, "inventory")
                with open(inv_target, "w", encoding="utf-8") as invf:
                    invf.write(inventory_content)
            else:
                inv_target = inventory_path
            cmd = build_ansible_command(playbook_path, inv_target, params)

            cred_args, cred_files = build_credentials(creds, tmpdir)
            for cred_path, content in cred_files:
                with open(cred_path, "w", encoding="utf-8") as cf:
                    cf.write(content)
                os.chmod(cred_path, 0o600)

            proc = await asyncio.create_subprocess_exec(
                *cmd, *cred_args,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
                cwd=cwd,
                env=ansible_env(creds, os.environ),
            )
            _run_procs[run_id] = proc
            assert proc.stdout is not None

            async for raw_line in proc.stdout:
                line = raw_line.decode(errors="replace").rstrip("\n\r")
                _broadcast(run_id, line)
                if logf is not None:
                    logf.write(line + "\n")
                if (line.startswith("fatal:") or "FAILED!" in line) and len(failures) < _FAILURE_CAP:
                    failures.append(line[:500])

                if not truncated:
                    candidate = (db_output + "\n" + line) if db_output else line
                    if len(candidate.encode()) > _OUTPUT_CAP_BYTES:
                        db_output += "\n[output truncated — download the full log]"
                        truncated = True
                    else:
                        db_output = candidate

                lines_since_flush += 1
                elapsed = time.monotonic() - last_flush_at
                if lines_since_flush >= _FLUSH_LINES or elapsed >= _FLUSH_SECS:
                    await _flush()

            await proc.wait()
            if run_id in _cancelled:
                final_status = "cancelled"
            else:
                final_status = "success" if proc.returncode == 0 else "failed"
        finally:
            shutil.rmtree(tmpdir, ignore_errors=True)

    except _Aborted:
        final_status = "cancelled"
    except Exception:
        logger.exception("Ansible run %s failed with exception", run_id)
        final_status = "cancelled" if run_id in _cancelled else "failed"
    finally:
        if logf is not None:
            logf.close()
        _run_procs.pop(run_id, None)
        _cancelled.discard(run_id)
        sem.release()

    results = {"hosts": parse_recap(_run_lines.get(run_id, [])), "failures": failures}
    await _flush(final_status, results=results)
    _broadcast(run_id, _Done(status=final_status))