#!/usr/bin/env python3
"""Pueue backfill watcher with rich Telegram progress reports.

Stdlib-only (no requests/httpx) — works in any Python 3.13 venv.

Usage:
    source ~/opendeviationbar-py/deploy/opendeviationbar.local.env
    export OPENDEVIATIONBAR_TELEGRAM_TOKEN OPENDEVIATIONBAR_TELEGRAM_CHAT_ID
    python pueue-watcher.py
"""

import contextlib
import json
import os
import re
import subprocess
import sys
import time
import traceback
import urllib.error
import urllib.request
from datetime import UTC, datetime
from pathlib import Path

POLL_S = 300


def _esc(text: str) -> str:
    """Escape text for Telegram HTML."""
    return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
TOKEN = os.environ.get("OPENDEVIATIONBAR_TELEGRAM_TOKEN", "")
CHAT = os.environ.get("OPENDEVIATIONBAR_TELEGRAM_CHAT_ID", "")
OVERRIDE = Path(
    "~/.config/systemd/user/opendeviationbar-kintsugi.service.d/backfill-guard.conf"
).expanduser()


def tg(msg: str, *, silent: bool = False) -> None:
    """Send Telegram message using stdlib urllib (no requests dependency)."""
    payload = json.dumps({
        "chat_id": CHAT,
        "text": msg,
        "parse_mode": "HTML",
        "disable_notification": silent,
    }).encode()
    req = urllib.request.Request(
        f"https://api.telegram.org/bot{TOKEN}/sendMessage",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        urllib.request.urlopen(req, timeout=10)
    except urllib.error.HTTPError as e:
        print(f"Telegram HTTP {e.code}: {e.read().decode()[:200]}", file=sys.stderr)
        raise


def pueue_tasks() -> dict:
    r = subprocess.run(
        ["pueue", "status", "--json"],
        capture_output=True, text=True, timeout=30, check=False,
    )
    if r.returncode != 0:
        msg = f"pueue: {r.stderr}"
        raise RuntimeError(msg)
    return json.loads(r.stdout).get("tasks", {})


_MAX_TQDM_TOTAL = 400


def get_task_progress(task_id: int) -> tuple[int, int] | None:
    """Parse tqdm N/TOTAL from a running task's log. Returns (current, total) or None."""
    try:
        r = subprocess.run(
            ["pueue", "log", "--lines", "20", str(task_id)],
            capture_output=True, text=True, timeout=10, check=False,
        )
        matches = re.findall(r"(\d+)/(\d+)", r.stdout)
        if matches:
            current, total = int(matches[-1][0]), int(matches[-1][1])
            if 0 < total <= _MAX_TQDM_TOTAL and 0 <= current <= total:
                return (current, total)
    except Exception:
        return None  # pueue log unavailable — non-critical
    return None


def get_failure_reason(task_id: int) -> str | None:
    """Extract failure reason from a failed task's log."""
    try:
        r = subprocess.run(
            ["pueue", "log", "--lines", "30", str(task_id)],
            capture_output=True, text=True, timeout=10, check=False,
        )
        lines = r.stdout.strip().splitlines()
        # Look for error/traceback/signal info
        for line in reversed(lines):
            line_stripped = line.strip()
            if not line_stripped or line_stripped.startswith("BTCUSDT:"):
                continue
            if any(kw in line_stripped.lower() for kw in [
                "error", "exception", "killed", "signal", "oom",
                "memory", "traceback", "failed", "timeout",
            ]):
                return line_stripped[:120]

        # Fallback: check exit code from pueue status
        r2 = subprocess.run(
            ["pueue", "status", "--json"],
            capture_output=True, text=True, timeout=10, check=False,
        )
        tasks = json.loads(r2.stdout).get("tasks", {})
        task = tasks.get(str(task_id))
        if task:
            status = task.get("status", {})
            if "Done" in status:
                result = status["Done"].get("result", {})
                if result == "Killed":
                    return "killed by system or user (likely OOM — check dmesg)"
                if isinstance(result, dict) and "Failed" in result:
                    code = result["Failed"]
                    if isinstance(code, int):
                        exit_map = {
                            134: "exit 134 (SIGKILL — likely OOM or pueue kill)",
                            137: "exit 137 (SIGKILL — OOM killer)",
                            139: "exit 139 (SIGSEGV — segfault)",
                        }
                        return exit_map.get(code, f"exit {code}")
                    return f"Failed: {code}"
    except Exception:
        return None  # pueue log unavailable — non-critical
    return None


def parse_status(t: dict) -> tuple[str, str | None, str | None, str | None]:
    """Returns (state, start_iso, end_iso, result)."""
    s = t.get("status", {})
    if "Done" in s:
        d = s["Done"]
        result = d.get("result", "Unknown")
        if isinstance(result, dict) and "Failed" in result:
            result = "Failed"
        return ("Done", d.get("start"), d.get("end"), str(result))
    if "Running" in s:
        d = s["Running"]
        return ("Running", d.get("start"), None, None)
    if "Queued" in s:
        return ("Queued", None, None, None)
    if "Stashed" in s:
        return ("Stashed", None, None, None)
    return (str(list(s.keys())), None, None, None)


def fmt_duration(start_iso: str | None, end_iso: str | None = None) -> str:
    if not start_iso:
        return ""
    s = datetime.fromisoformat(start_iso.replace("Z", "+00:00"))
    e = (
        datetime.fromisoformat(end_iso.replace("Z", "+00:00"))
        if end_iso
        else datetime.now(UTC)
    )
    d = (e - s).total_seconds()
    h, rem = divmod(int(d), 3600)
    m, _ = divmod(rem, 60)
    return f"{h}h{m:02d}m" if h else f"{m}m"


def fmt_report(tasks: dict, *, final: bool = False) -> str:  # noqa: PLR0915
    groups: dict[str, list] = {}
    for t in tasks.values():
        g = t.get("group", "default")
        groups.setdefault(g, []).append(t)

    now = datetime.now(UTC).strftime("%H:%M:%S UTC")
    title = (
        "\U0001f4ca <b>Backfill FINAL Report</b>"
        if final
        else "\U0001f4ca <b>Backfill Progress</b>"
    )
    lines = [f"{title} \u2014 {now}", ""]

    total_r, total_q, total_s, total_f = 0, 0, 0, 0
    all_durations: list[float] = []
    weighted_done = 0.0
    total_jobs = 0

    for g in sorted(groups):
        jobs = sorted(groups[g], key=lambda t: t.get("id", 0))
        running, queued, success, failed = [], [], [], []
        for t in jobs:
            state, start, end, result = parse_status(t)
            if state == "Running":
                running.append(t)
            elif state in ("Queued", "Stashed"):
                queued.append(t)
            elif state == "Done" and result == "Success":
                success.append(t)
                if start and end:
                    sa = datetime.fromisoformat(start.replace("Z", "+00:00"))
                    ea = datetime.fromisoformat(end.replace("Z", "+00:00"))
                    all_durations.append((ea - sa).total_seconds())
            elif state == "Done":
                failed.append(t)

        total_r += len(running)
        total_q += len(queued)
        total_s += len(success)
        total_f += len(failed)

        if failed and not running and not queued:
            icon = "\u26a0\ufe0f"
        elif not running and not queued and not failed:
            icon = "\u2705"
        elif running:
            icon = "\u23f3"
        else:
            icon = "\U0001f4cb"
        n = len(success) + len(failed) + len(running) + len(queued)
        lines.append(
            f"{icon} <b>{g}</b> ({len(success) + len(failed)}/{n}): "
            f"{len(success)}\u2713 {len(failed)}\u2717 "
            f"{len(running)}\u23f3 {len(queued)}\U0001f4cb"
        )

        for t in jobs:
            state, start, end, result = parse_status(t)
            label = _esc(t.get("label", f"task-{t.get('id', '?')}"))
            dur = fmt_duration(start, end)
            total_jobs += 1
            if state == "Running":
                prog = get_task_progress(t.get("id", -1))
                if prog:
                    cur, tot = prog
                    pct_job = cur / tot * 100
                    weighted_done += cur / tot
                    lines.append(
                        f"  \U0001f504 {label} "
                        f"({cur}/{tot} days, {pct_job:.0f}% \u2014 {dur} elapsed)"
                    )
                else:
                    lines.append(f"  \U0001f504 {label} ({dur} elapsed)")
            elif state == "Done" and result != "Success":
                weighted_done += 1.0
                reason = get_failure_reason(t.get("id", -1))
                reason_str = _esc(reason) if reason else "unknown error (check: pueue log &lt;id&gt;)"
                lines.append(f"  \u274c {label} ({dur}) \u2014 {reason_str}")
            elif state == "Done":
                weighted_done += 1.0
                lines.append(f"  \u2705 {label} ({dur})")
            else:
                lines.append(f"  \u23f8 {label}")
        lines.append("")

    total = total_r + total_q + total_s + total_f
    done = total_s + total_f
    pct = weighted_done / max(total_jobs, 1) * 100
    bar_len = 20
    filled = int(bar_len * pct / 100)
    bar = "\u2588" * filled + "\u2591" * (bar_len - filled)
    lines.append(
        f"<b>Progress</b>: {bar} {pct:.0f}% "
        f"({done}/{total} done + {total_r} in-flight)"
    )
    lines.append(
        f"  {total_s}\u2713 done | {total_f}\u2717 failed | "
        f"{total_r}\u23f3 running | {total_q}\U0001f4cb queued"
    )

    if total_r + total_q > 0 and all_durations:
        avg_s = sum(all_durations) / len(all_durations)
        remaining_per_group: dict[str, int] = {}
        for t in tasks.values():
            g = t.get("group", "default")
            state = parse_status(t)[0]
            if state in ("Running", "Queued", "Stashed"):
                remaining_per_group[g] = remaining_per_group.get(g, 0) + 1
        max_remaining = max(remaining_per_group.values()) if remaining_per_group else 0
        eta_s = avg_s * max_remaining
        eta_h = eta_s / 3600
        eta_time = datetime.now(UTC).timestamp() + eta_s
        eta_str = datetime.fromtimestamp(
            eta_time, tz=UTC
        ).strftime("%H:%M UTC %b %d")
        lines.append(f"  \u23f1 ETA: ~{eta_h:.1f}h ({eta_str})")
        lines.append(f"  avg {avg_s / 60:.0f}m/job, longest queue: {max_remaining} jobs")

    if total_f > 0:
        lines.append("")
        lines.append(
            "<i>Failed jobs resume from checkpoint</i> \u2014 "
            "<code>pueue restart &lt;id&gt;</code> to retry"
        )

    return "\n".join(lines)


def remove_override() -> None:
    override_path = Path(OVERRIDE)
    if override_path.exists():
        override_path.unlink()
    with contextlib.suppress(OSError):
        override_path.parent.rmdir()
    subprocess.run(
        ["systemctl", "--user", "daemon-reload"], check=True, timeout=30,
    )
    subprocess.run(
        ["systemctl", "--user", "restart", "opendeviationbar-kintsugi"],
        check=True, timeout=30,
    )


def main() -> int:
    try:
        import setproctitle

        setproctitle.setproctitle(
            f"odb-pueue-watcher: poll-interval={POLL_S}s telegram-alerts"
        )
    except ImportError:
        pass  # Optional on dev machines

    if not TOKEN or not CHAT:
        print(
            "FATAL: Set OPENDEVIATIONBAR_TELEGRAM_TOKEN and "
            "OPENDEVIATIONBAR_TELEGRAM_CHAT_ID",
            file=sys.stderr,
        )
        return 1

    tasks = pueue_tasks()
    report = fmt_report(tasks)
    tg(
        f"\U0001f50d <b>pueue-watcher started</b>\n"
        f"Polling every {POLL_S}s with per-job progress.\n\n{report}"
    )
    print(report)

    while True:
        try:
            tasks = pueue_tasks()
            active = sum(
                1 for t in tasks.values()
                if parse_status(t)[0] in ("Running", "Queued", "Stashed")
            )

            report = fmt_report(tasks, final=(active == 0))
            print(report)
            tg(report, silent=(active > 0))

            if active == 0:
                remove_override()
                tg(
                    "\U0001f513 <b>Kintsugi override removed + restarted</b>\n"
                    "Now repairing ALL thresholds (250/500/750/1000).\n"
                    "pueue-watcher exiting."
                )
                return 0

        except Exception:
            tb = traceback.format_exc()
            print(tb, file=sys.stderr)
            try:
                tg(
                    f"\U0001f534 <b>pueue-watcher ERROR</b>\n"
                    f"<pre>{tb[-500:]}</pre>\nRetrying in {POLL_S}s."
                )
            except Exception:
                print("Failed to send error to Telegram!", file=sys.stderr)

        time.sleep(POLL_S)


if __name__ == "__main__":
    sys.exit(main())
