use pyo3::prelude::*;
use pyo3::types::PyAny;

pub struct Watchdog {
    pub arm: Py<PyAny>,
    pub disarm: Py<PyAny>,
    pub clear_async_exc: Py<PyAny>,
}

#[cfg(feature = "tle")]
mod tle_impl {
    use super::*;

    impl Watchdog {
        /// Create a new watchdog.
        ///
        /// Single long-lived watchdog thread that reuses a Condition variable.
        /// `arm(wall_timeout)` wakes it; `disarm()` cancels before it fires.
        /// No thread spawn/join per unit turn.
        ///
        /// The watchdog sleeps for the wall-clock timeout, then fires
        /// a SystemExit async exception on the main thread. CPU time
        /// enforcement is handled separately by `check_deadline` in the
        /// controller, which runs on the main thread at every API call.
        pub fn new(py: Python) -> Watchdog {
            let watchdog_mod = pyo3::types::PyModule::from_code(
                py,
                c"
import ctypes
import threading
import time

# Capture references before bot code can monkey-patch them.
_monotonic = time.monotonic

tid = threading.get_ident()
set_async_exc = ctypes.pythonapi.PyThreadState_SetAsyncExc
cond = threading.Condition()

_wall_deadline = None
shutdown_flag = False

def watchdog_loop():
    global _wall_deadline, shutdown_flag
    with cond:
        while True:
            while _wall_deadline is None and not shutdown_flag:
                cond.wait()
            if shutdown_flag:
                return
            remaining = _wall_deadline - _monotonic()
            if remaining > 0:
                cond.wait(timeout=remaining)
            if _wall_deadline is None:
                continue
            set_async_exc(ctypes.c_long(tid), ctypes.py_object(SystemExit))
            _wall_deadline = None

thread = threading.Thread(target=watchdog_loop, daemon=True)
thread.start()

def arm(wall_timeout):
    global _wall_deadline
    with cond:
        _wall_deadline = _monotonic() + wall_timeout
        cond.notify()

def disarm():
    global _wall_deadline
    with cond:
        _wall_deadline = None
        cond.notify()

def shutdown():
    global shutdown_flag
    with cond:
        shutdown_flag = True
        cond.notify()

def clear_async_exc():
    while True:
        try:
            set_async_exc(ctypes.c_long(tid), ctypes.c_long(0))
            return
        except SystemExit:
            pass
",
                c"_watchdog",
                c"_watchdog",
            )
            .expect("watchdog module");

            Watchdog {
                arm: watchdog_mod.getattr("arm").expect("arm attr").unbind(),
                disarm: watchdog_mod
                    .getattr("disarm")
                    .expect("disarm attr")
                    .unbind(),
                clear_async_exc: watchdog_mod
                    .getattr("clear_async_exc")
                    .expect("clear_async_exc attr")
                    .unbind(),
            }
        }
    }
}

#[cfg(not(feature = "tle"))]
impl Watchdog {
    /// No-op watchdog when TLE enforcement is disabled.
    pub fn new(py: Python) -> Watchdog {
        let noop = py
            .eval(c"lambda *a, **kw: None", None, None)
            .expect("noop lambda")
            .unbind();
        Watchdog {
            arm: noop.clone_ref(py),
            disarm: noop.clone_ref(py),
            clear_async_exc: noop,
        }
    }
}
