#!/usr/bin/env python3
"""
pi_wake_agent.py — Reusable multi-agent wake-up timer with self-cron/daemon/succession

Usage:
    pi_wake_agent.py --install --interval 1h --session cc_UV_dev0_Fb --msg "Operator says CONTINUE. Pi here!"
    pi_wake_agent.py --once --interval 2h --session cc_UV_dev0_Fb --msg "Time's up!"
    pi_wake_agent.py --daemon --interval 1h --session cc_UV_dev0_Fb
    pi_wake_agent.py --succession --count 3 --interval 1h --session cc_UV_dev0_Fb --msg "Scheduled wake"
    pi_wake_agent.py --remove --session cc_UV_dev0_Fb --interval 1h
    pi_wake_agent.py --list
    pi_wake_agent.py --status
    pi_wake_agent.py --validate --session cc_UV_dev0_Fb
"""

import argparse
import logging
import os
import re
import shlex
import signal
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import List, Optional

# ─── Constants ────────────────────────────────────────────────────────────
SCRIPT_PATH = Path(__file__).resolve()
LOG_FILE = Path("/tmp/pi_wake_agent.log")
LOG_MAX_SIZE = 10 * 1024 * 1024  # 10 MB
LOG_MAX_FILES = 5
CRON_COMMENT_PREFIX = "pi_wake_agent"
AGENT_NICK = "pi_nvnemo"
H5I_AGENT = "pi_nvnemo"
H5I_BUS_ROOT = Path("/mnt/dolphinng5_predict")
DEFAULT_INTERVAL = "1h"

# ─── Logging Setup ───────────────────────────────────────────────────────
def setup_logging(debug: bool = False) -> logging.Logger:
    log_rotate()
    logger = logging.getLogger("pi_wake_agent")
    logger.setLevel(logging.DEBUG if debug else logging.INFO)
    
    fh = logging.FileHandler(LOG_FILE)
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(logging.Formatter("[%(asctime)s] [%(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S"))
    logger.addHandler(fh)
    
    ch = logging.StreamHandler(sys.stderr)
    ch.setLevel(logging.WARNING)
    ch.setFormatter(logging.Formatter("[%(asctime)s] [%(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S"))
    logger.addHandler(ch)
    
    return logger


def log_rotate() -> None:
    if not LOG_FILE.exists():
        return
    size = LOG_FILE.stat().st_size
    if size < LOG_MAX_SIZE:
        return
    for i in range(LOG_MAX_FILES - 1, 0, -1):
        src = LOG_FILE.with_suffix(f".log.{i}") if i > 1 else LOG_FILE.with_suffix(".log.1")
        if src.exists():
            dst = LOG_FILE.with_suffix(f".log.{i + 1}")
            src.rename(dst)
    LOG_FILE.rename(LOG_FILE.with_suffix(".log.1"))


# ─── Interval Parsing ────────────────────────────────────────────────────
def parse_interval(interval: str) -> int:
    """Parse interval string (e.g., '1h', '30m', '90m', '2h', '10s') to seconds."""
    match = re.match(r"^(\d+)([hms])$", interval)
    if not match:
        raise ValueError(f"Invalid interval format '{interval}'. Use like 1h, 30m, 90m, 2h, 10s")
    value, unit = int(match.group(1)), match.group(2)
    if unit == "h":
        return value * 3600
    elif unit == "m":
        return value * 60
    elif unit == "s":
        return value
    raise ValueError(f"Unknown unit: {unit}")


def interval_to_cron(interval: str) -> str:
    """Convert interval to cron schedule."""
    match = re.match(r"^(\d+)([hm])$", interval)
    if not match:
        raise ValueError(f"Cron only supports minutes/hours intervals: {interval}")
    value, unit = int(match.group(1)), match.group(2)
    if unit == "h":
        return f"0 */{value} * * *"
    elif unit == "m":
        if value >= 60:
            raise ValueError(f"For minutes >= 60, use hours (e.g., 1h not 60m)")
        return f"*/{value} * * * *"
    raise ValueError(f"Cron only supports minutes/hours intervals: {interval}")


def interval_to_human(interval: str) -> str:
    match = re.match(r"^(\d+)([hms])$", interval)
    if not match:
        return interval
    value, unit = match.group(1), match.group(2)
    if unit == "h":
        return f"{value} hour(s)"
    elif unit == "m":
        return f"{value} minute(s)"
    elif unit == "s":
        return f"{value} second(s)"
    return interval


# ─── Cron Management ────────────────────────────────────────────────────
def cron_comment(sessions: List[str], interval: str) -> str:
    sessions_str = ",".join(sessions)
    return f"{CRON_COMMENT_PREFIX}:{sessions_str}:{interval}"


def install_cron(sessions: List[str], interval: str, message: str, logger: logging.Logger) -> None:
    cron_sched = interval_to_cron(interval)
    comment = cron_comment(sessions, interval)
    sessions_arg = f"'{','.join(sessions)}'"
    cmd = f"cd {H5I_BUS_ROOT} && export H5I_AGENT={H5I_AGENT} && {SCRIPT_PATH} --run --sessions {sessions_arg} --msg {shlex.quote(message)}"
    
    result = subprocess.run(["crontab", "-l"], capture_output=True, text=True)
    existing = result.stdout if result.returncode == 0 else ""
    lines = [line for line in existing.splitlines() if comment not in line]
    lines.append(f"{cron_sched} {cmd} # {comment}")
    new_cron = "\n".join(lines) + "\n"
    subprocess.run(["crontab", "-"], input=new_cron, text=True, check=True)
    logger.info(f"Installed cron: {cron_sched} -> {sessions} every {interval}")


def remove_cron(sessions: List[str], interval: str, logger: logging.Logger) -> None:
    comment = cron_comment(sessions, interval)
    result = subprocess.run(["crontab", "-l"], capture_output=True, text=True)
    existing = result.stdout if result.returncode == 0 else ""
    lines = [line for line in existing.splitlines() if comment not in line]
    new_cron = "\n".join(lines) + ("\n" if lines else "")
    subprocess.run(["crontab", "-"], input=new_cron, text=True, check=True)
    logger.info(f"Removed cron for {sessions} ({interval})")


def list_cron(logger: logging.Logger) -> None:
    print("=== pi_wake_agent cron entries ===")
    result = subprocess.run(["crontab", "-l"], capture_output=True, text=True)
    existing = result.stdout if result.returncode == 0 else ""
    found = False
    for line in existing.splitlines():
        if CRON_COMMENT_PREFIX in line:
            print(f"  {line}")
            found = True
    if not found:
        print("  (none)")


# ─── Zellij Operations ──────────────────────────────────────────────────
def zellij_session_exists(session: str) -> bool:
    try:
        result = subprocess.run(["zellij", "list-sessions"], capture_output=True, text=True, timeout=5)
        clean = re.sub(r'\x1b\[[0-9;]*m', '', result.stdout)
        for line in clean.splitlines():
            if line.startswith(session + " ") or line == session:
                return True
    except Exception:
        pass
    return False


def zellij_write_chars(session: str, text: str) -> bool:
    try:
        subprocess.run(["zellij", "--session", session, "action", "write-chars", text], 
                       capture_output=True, timeout=5)
        return True
    except Exception:
        return False


def zellij_write_enter(session: str) -> bool:
    try:
        subprocess.run(["zellij", "--session", session, "action", "write", "13"], 
                       capture_output=True, timeout=5)
        return True
    except Exception:
        return False


# ─── Wake Action ────────────────────────────────────────────────────────────────────
    logger.info(f"One-shot timer set for {interval} ({interval_to_human(interval)})")
    
    def _wake():
        time.sleep(interval_seconds)
        run_wake(sessions, message, logger)
    
    thread = threading.Thread(target=_wake, daemon=True)
    thread.start()
    logger.info(f"Background timer started (thread: {thread.ident})")
    pid_file = Path(f"/tmp/pi_wake_agent_{sessions[0]}.pid")
    pid_file.write_text(str(os.getpid()))


def run_wake(sessions: List[str], message: str, logger: logging.Logger) -> None:
    sessions_str = ",".join(sessions)
    logger.info(f"Waking {sessions_str} with message: {message}")
    
    for session in sessions:
        zellij_write_chars(session, f"[{AGENT_NICK} via zellij] {message} Run: h5i-bus msg inbox")
        for _ in range(5):
            zellij_write_enter(session)
            time.sleep(1)
    
    # Fire-and-forget h5i bus message (non-blocking)
    def _send_bus():
        try:
            subprocess.run(
                ["h5i", "msg", "send", "Fable", f"{message} (timer wakeup)"],
                cwd=H5I_BUS_ROOT,
                env={**os.environ, "H5I_AGENT": H5I_AGENT},
                capture_output=True,
                timeout=5
            )
        except Exception:
            pass  # Silently ignore - fire and forget
    
    threading.Thread(target=_send_bus, daemon=True).start()
    
    logger.info(f"Wake sent to {sessions_str}")

def run_once(sessions: List[str], interval: str, message: str, logger: logging.Logger) -> None:
    interval_seconds = parse_interval(interval)
    logger.info(f"One-shot timer set for {interval} ({interval_to_human(interval)})")
    
    def _wake():
        time.sleep(interval_seconds)
        run_wake(sessions, message, logger)
    
    thread = threading.Thread(target=_wake, daemon=True)
    thread.start()
    logger.info(f"Background timer started (thread: {thread.ident})")
    pid_file = Path(f"/tmp/pi_wake_agent_{sessions[0]}.pid")
    pid_file.write_text(str(os.getpid()))
    return 0


def run_daemon(sessions: List[str], interval: str, message: str, logger: logging.Logger) -> None:
    interval_seconds = parse_interval(interval)
    logger.info("=== DAEMON START ===")
    logger.info(f"Interval: {interval} ({interval_to_human(interval)})")
    logger.info(f"Sessions: {sessions}")
    logger.info(f"Message: {message}")
    
    run_wake(sessions, message, logger)
    
    while True:
        logger.debug(f"Sleeping for {interval_seconds}s...")
        time.sleep(interval_seconds)
        run_wake(sessions, message, logger)


def run_succession(sessions: List[str], interval: str, count: int, message: str, logger: logging.Logger) -> None:
    """Run wake N times at interval, then self-clean (remove cron if installed)."""
    interval_seconds = parse_interval(interval)
    logger.info(f"=== SUCCESSION START === Count: {count}, Interval: {interval} ({interval_to_human(interval)})")
    logger.info(f"Sessions: {sessions}")
    logger.info(f"Message: {message}")
    
    for i in range(1, count + 1):
        logger.info(f"Succession {i}/{count}")
        run_wake(sessions, message, logger)
        if i < count:
            logger.debug(f"Sleeping for {interval_seconds}s until next succession...")
            time.sleep(interval_seconds)
    
    # Self-clean: remove any cron entry for this session/interval combo
    try:
        remove_cron(sessions, interval, logger)
        logger.info("Self-cleanup complete (cron removed)")
    except Exception as e:
        logger.warning(f"Self-cleanup failed: {e}")
    
    logger.info("=== SUCCESSION COMPLETE ===")


def status(logger: logging.Logger) -> None:
    print("=== pi_wake_agent Status ===")
    print(f"Script: {SCRIPT_PATH}")
    print(f"Log: {LOG_FILE}")
    print(f"Log size: {LOG_FILE.stat().st_size if LOG_FILE.exists() else 'N/A'} bytes")
    print(f"Agent: {AGENT_NICK}")
    print()
    list_cron(logger)
    print()
    print("=== Active one-shot timers ===")
    found = False
    for pid_file in Path("/tmp").glob("pi_wake_agent_*.pid"):
        found = True
        try:
            pid = int(pid_file.read_text().strip())
            os.kill(pid, 0)
            print(f"  PID {pid} (active)")
        except (ProcessLookupError, ValueError):
            print(f"  PID {pid_file.read_text().strip()} (dead, cleaning up)")
            pid_file.unlink(missing_ok=True)
    if not found:
        print("  (none)")


def validate_sessions(sessions: List[str], logger: logging.Logger) -> None:
    logger.info(f"Validating sessions: {sessions}")
    for session in sessions:
        if zellij_session_exists(session):
            logger.info(f"  {session}: EXISTS")
        else:
            logger.warning(f"  {session}: NOT FOUND (may be dead/EXITED)")


# ─── Session Parsing ─────────────────────────────────────────────────────
def parse_sessions(raw) -> List[str]:
    if not raw:
        return []
    if isinstance(raw, list):
        return [s.strip() for s in raw]
    return [s.strip() for s in raw.split(",") if s.strip()]


# ─── Usage ──────────────────────────────────────────────────────────────
def create_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="pi_wake_agent.py — Reusable multi-agent wake-up timer with self-cron/daemon/succession",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
EXAMPLES:
  # Install recurring 1-hour timer for one session
  pi_wake_agent.py --install --interval 1h --session cc_UV_dev0_Fb

  # Install 30-minute timer for multiple sessions
  pi_wake_agent.py --install --interval 30m --sessions "cc_UV_dev0_Fb,cc_UV_dev1_48" --msg "Wake up!"

  # One-shot wake in 2 hours (no cron)
  pi_wake_agent.py --once --interval 2h --session cc_UV_dev0_Fb --msg "Time's up!"

  # Run N times at interval, then self-clean
  pi_wake_agent.py --succession --count 3 --interval 1h --session cc_UV_dev0_Fb --msg "Scheduled wake"

  # Run as daemon (long-lived process, no cron)
  pi_wake_agent.py --daemon --interval 1h --session cc_UV_dev0_Fb

  # Remove timer
  pi_wake_agent.py --remove --session cc_UV_dev0_Fb --interval 1h

  # List all timers
  pi_wake_agent.py --list

  # Show status
  pi_wake_agent.py --status

  # Validate sessions
  pi_wake_agent.py --validate --session cc_UV_dev0_Fb
"""
    )
    parser = argparse.ArgumentParser(
        description="pi_wake_agent.py — Reusable multi-agent wake-up timer with self-cron/daemon/succession",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--install", action="store_const", const="install", dest="mode", help="Install recurring cron timer")
    parser.add_argument("--once", action="store_const", const="once", dest="mode", help="One-shot wake (no cron)")
    parser.add_argument("--daemon", action="store_const", const="daemon", dest="mode", help="Run as long-lived daemon (no cron)")
    parser.add_argument("--succession", action="store_const", const="succession", dest="mode", help="Run N times at interval, then self-clean")
    parser.add_argument("--count", type=int, default=1, help="Number of successions (for --succession mode)")
    parser.add_argument("--run", action="store_const", const="run", dest="mode", help="Internal: run wake action (called by cron)")
    parser.add_argument("--remove", action="store_const", const="remove", dest="mode", help="Remove cron timer")
    parser.add_argument("--list", action="store_const", const="list", dest="mode", help="List active cron timers")
    parser.add_argument("--status", action="store_const", const="status", dest="mode", help="Show status (cron + one-shot timers)")
    parser.add_argument("--validate", action="store_const", const="validate", dest="mode", help="Validate sessions exist in zellij")
    parser.add_argument("--interval", default=DEFAULT_INTERVAL, help="Interval (default: 1h). Formats: 30m, 1h, 90m, 2h, etc.")
    parser.add_argument("--session", action="append", dest="session_list", help="Zellij session name (can repeat)")
    parser.add_argument("--sessions", help="Comma-separated list of sessions")
    parser.add_argument("--msg", default=f"Operator says CONTINUE. {AGENT_NICK} here, saying hi!", help="Wake message")
    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
    parser.add_argument("--dry-run", action="store_true", help="Show what would be done without executing")
    parser.set_defaults(mode="install")
    return parser


# ─── Main ───────────────────────────────────────────────────────────────
def main() -> int:
    parser = create_parser()
    args = parser.parse_args()
    
    # Combine sessions
    sessions: List[str] = []
    if args.sessions:
        sessions.extend(parse_sessions(args.sessions))
    if args.session_list:
        sessions.extend(args.session_list)
    
    # Setup logging
    logger = setup_logging(args.debug)
    
    # Validate
    if len(sessions) == 0 and args.mode not in ("list", "status"):
        logger.error("--session or --sessions required")
        parser.print_help()
        return 1
    
    if args.mode == "succession" and args.count < 1:
        logger.error("--count must be >= 1 for succession mode")
        return 1
    
    logger.info(f"=== pi_wake_agent {args.mode} ===")
    logger.info(f"Interval: {args.interval} ({interval_to_human(args.interval)})")
    logger.info(f"Sessions: {sessions}")
    logger.info(f"Message: {args.msg}")
    if args.mode == "succession":
        logger.info(f"Count: {args.count}")
    if args.dry_run:
        logger.info("DRY RUN: no actions will be executed")
    
    if args.dry_run:
        logger.info("DRY RUN: would execute mode '%s' with sessions=%s, interval=%s, message='%s'", args.mode, sessions, args.interval, args.msg)
        return 0
    
    try:
        if args.mode == "install":
            install_cron(sessions, args.interval, args.msg, logger)
        elif args.mode == "once":
            run_once(sessions, args.interval, args.msg, logger)
        elif args.mode == "daemon":
            run_daemon(sessions, args.interval, args.msg, logger)
        elif args.mode == "succession":
            run_succession(sessions, args.interval, args.count, args.msg, logger)
        elif args.mode == "run":
            run_wake(sessions, args.msg, logger)
        elif args.mode == "remove":
            remove_cron(sessions, args.interval, logger)
        elif args.mode == "list":
            list_cron(logger)
        elif args.mode == "status":
            status(logger)
        elif args.mode == "validate":
            validate_sessions(sessions, logger)
        else:
            logger.error(f"Unknown mode: {args.mode}")
            return 1
    except Exception as e:
        logger.error(f"Error: {e}")
        if args.debug:
            import traceback
            traceback.print_exc()
        return 1
    
    return 0


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