mod watchdog;

use std::cell::RefCell;
use std::collections::HashMap;
use std::path::Path;
use std::rc::Rc;

use pyo3::prelude::*;
use pyo3::types::{PyDict, PyModule};

use crate::bindings as rustlib;
use crate::bindings::controller::Controller;
use crate::bindings::py_convert;
use crate::cli::Args;
use crate::common::game_constants::MAX_TURNS;
use crate::common::Team;
use crate::game::Game;
use crate::map_loader;
use crate::replay::recorder::GameDiff;

/// Returns the calling thread's cumulative CPU time in nanoseconds.
/// Uses CLOCK_THREAD_CPUTIME_ID — immune to noisy neighbors and OS scheduling.
/// Returns 0 when the `tle` feature is disabled.
#[cfg(feature = "tle")]
pub fn thread_cpu_time_ns() -> u64 {
    let mut ts = libc::timespec {
        tv_sec: 0,
        tv_nsec: 0,
    };
    unsafe {
        libc::clock_gettime(libc::CLOCK_THREAD_CPUTIME_ID, &mut ts);
    }
    ts.tv_sec as u64 * 1_000_000_000 + ts.tv_nsec as u64
}

#[cfg(not(feature = "tle"))]
pub fn thread_cpu_time_ns() -> u64 {
    0
}

/// Percent of the per-turn time limit that can be banked as extra time.
const ADAPTIVE_TIME_PERCENT: u64 = 5;

struct GameRunner {
    game: Rc<RefCell<Game>>,
    player_classes: [Py<PyAny>; 2],
    unit_runners: HashMap<i32, Py<PyAny>>,
    watchdog: watchdog::Watchdog,
    turn_timeout_ms: u64,
    /// Banked extra time per bot (nanoseconds), used to absorb CPU jitter.
    bot_extra_time_ns: HashMap<i32, u64>,
    // Cached module references to avoid py.import() during gameplay
    sys_mod: Py<PyModule>,
    io_mod: Py<PyModule>,
    gc_mod: Py<PyModule>,
}

impl GameRunner {
    fn run(&mut self, py: Python) -> PyResult<()> {
        let sys = self.sys_mod.bind(py);
        sys.call_method1("setswitchinterval", (self.turn_timeout_ms as f64 / 1000.0,))?;
        let gc = self.gc_mod.clone_ref(py).into_bound(py);
        gc.call_method0("disable")?;
        for i in 0..MAX_TURNS {
            self.run_turn(py, &gc)?;
            if i % 100 == 0 {
                println!("Completed turn {}", i);
            }
            if self.game.borrow_mut().winner_team().is_some() {
                break;
            }
        }
        gc.call_method0("enable")?;
        Ok(())
    }

    fn run_turn(&mut self, py: Python, gc: &Bound<'_, PyModule>) -> PyResult<()> {
        self.game.borrow_mut().new_turn();

        self.perform_unit_actions(py)?;
        let mut game = self.game.borrow_mut();
        game.distribute_resources();
        game.update_cooldowns();
        let players = game.players.clone();
        game.replay_recorder
            .append(GameDiff::UpdatePlayers { players });
        game.turn += 1;
        drop(game);
        self.cleanup_unit_runners();
        gc.call_method0("collect")?;
        Ok(())
    }

    fn perform_unit_actions(&mut self, py: Python) -> PyResult<()> {
        let cpu_budget_ns = self.turn_timeout_ms * 1_000_000;
        let max_extra_ns = cpu_budget_ns * ADAPTIVE_TIME_PERCENT / 100;
        let units = { self.game.borrow().unit_order.clone() };
        for unit_id in units {
            let game = self.game.borrow();
            let Some(entity) = game.entity(unit_id) else {
                continue;
            };
            let team = entity.team;
            drop(game);
            self.ensure_unit_runner(py, unit_id, team)
                .expect("ensure_unit_runner");

            let player = self
                .unit_runners
                .get(&unit_id)
                .expect("runner missing")
                .clone_ref(py);

            let extra_ns = *self
                .bot_extra_time_ns
                .entry(unit_id)
                .or_insert(max_extra_ns);
            let effective_budget_ns = cpu_budget_ns + extra_ns;

            let cpu_deadline_ns = thread_cpu_time_ns() + effective_budget_ns;
            let controller = Py::new(
                py,
                Controller::new(self.game.clone(), unit_id, cpu_deadline_ns),
            )
            .expect("Py::new controller");

            // Wall-clock timeout: 2x the effective CPU budget.
            // Gives headroom for GIL contention and OS scheduling jitter.
            let wall_timeout = effective_budget_ns as f64 / 1_000_000_000.0 * 1.9;

            loop {
                match self.watchdog.clear_async_exc.call0(py) {
                    Ok(_) => break,
                    Err(ref e) if e.is_instance_of::<pyo3::exceptions::PySystemExit>(py) => {
                        continue
                    }
                    Err(e) => panic!("clear_async_exc failed: {e}"),
                }
            }

            // Capture Python stdout for this unit's turn.
            let sys = self.sys_mod.bind(py);
            let io = self.io_mod.bind(py);
            let old_stdout = sys.getattr("stdout")?.unbind();
            let stdout_buf = io.call_method0("StringIO")?.unbind();
            sys.setattr("stdout", stdout_buf.clone_ref(py))?;

            loop {
                match self.watchdog.arm.call1(py, (wall_timeout,)) {
                    Ok(_) => break,
                    Err(ref e) if e.is_instance_of::<pyo3::exceptions::PySystemExit>(py) => {
                        continue
                    }
                    Err(e) => panic!("arm failed: {e}"),
                }
            }
            let cpu_start = thread_cpu_time_ns();
            let result = player.call_method1(py, "run", (controller,));
            let cpu_elapsed_ns = thread_cpu_time_ns() - cpu_start;

            // Disarm the watchdog first so no new async exceptions can fire
            loop {
                match self.watchdog.disarm.call0(py) {
                    Ok(_) => break,
                    Err(ref e) if e.is_instance_of::<pyo3::exceptions::PySystemExit>(py) => {
                        continue
                    }
                    Err(e) => panic!("disarm failed: {e}"),
                }
            }

            let stdout: String = stdout_buf
                .call_method0(py, "getvalue")
                .and_then(|s| s.extract(py))
                .unwrap_or_default();
            sys.setattr("stdout", old_stdout.clone_ref(py))?;

            self.game
                .borrow_mut()
                .replay_recorder
                .append(GameDiff::BotOutput {
                    id: unit_id,
                    stdout,
                    exec_time_us: (cpu_elapsed_ns / 1000) as u32,
                    tled: cpu_elapsed_ns > effective_budget_ns,
                });

            let cpu_tle = cpu_elapsed_ns > effective_budget_ns;

            if let Err(ref err) = result {
                if err.is_instance_of::<pyo3::exceptions::PyKeyboardInterrupt>(py) {
                    return Err(err.clone_ref(py));
                } else if err.is_instance_of::<pyo3::exceptions::PySystemExit>(py) {
                    // Watchdog or check_deadline fired — skip this unit's turn.
                } else {
                    err.print(py);
                    self.game.borrow_mut().destroy_entity(unit_id);
                }
            } else if cpu_tle {
                // Bot returned normally but exceeded CPU budget (tight loop
                // that finished just before the watchdog fired). Skip turn.
            }

            // Update banked extra time: bank unused time, debit overuse.
            let new_extra = if cpu_elapsed_ns <= cpu_budget_ns + extra_ns {
                (extra_ns + cpu_budget_ns).saturating_sub(cpu_elapsed_ns)
            } else {
                0
            };
            self.bot_extra_time_ns
                .insert(unit_id, new_extra.min(max_extra_ns));
        }
        Ok(())
    }

    fn ensure_unit_runner(&mut self, py: Python, unit: i32, team: Team) -> PyResult<()> {
        if self.unit_runners.contains_key(&unit) {
            return Ok(());
        }
        let player_cls = self.player_classes[team.index()].clone_ref(py);
        let player = player_cls.call0(py)?;
        self.unit_runners.insert(unit, player);
        Ok(())
    }

    fn cleanup_unit_runners(&mut self) {
        let game = self.game.borrow();
        self.unit_runners
            .retain(|unit, _| game.entity(*unit).is_some());
        self.bot_extra_time_ns
            .retain(|unit, _| game.entity(*unit).is_some());
    }
}

pub fn run(args: Args) -> PyResult<()> {
    Python::with_gil(|py| {
        ensure_sys_path(py, &args.engine_root)?;
        register_rust_module(py)?;

        // --- Phase 1: Initialize type cache (before sandbox, while imports work) ---
        py_convert::init_type_cache(py)?;

        // --- Phase 2: Initialize watchdog (needs ctypes + threading) ---
        // Must happen before the sandbox blocks ctypes imports.
        let watchdog = watchdog::Watchdog::new(py);

        // --- Phase 3: Install audit hook sandbox ---
        // Blocks ctypes, subprocess, socket, os.fork, gc introspection, etc.
        // for all subsequent Python code. The watchdog's set_async_exc reference
        // was captured above (before the hook) and triggers zero audit events.
        // Returns a lock function to tighten file access after bot loading.
        let sandbox_lock = if args.sandboxed {
            Some(install_sandbox(py)?)
        } else {
            None
        };

        // --- Phase 4: Validate + load bot code (runs under sandbox) ---
        // Bot imports/file reads still work here (fs lock not yet engaged).
        // Reject forbidden except handlers before loading bot code.
        check_except_handlers(py, &args.player_a)?;
        check_except_handlers(py, &args.player_b)?;
        // Exit codes: 10 = bot A failed, 11 = bot B failed, 12 = both failed.
        let a_result = load_player_class(py, &args.player_a);
        let b_result = load_player_class(py, &args.player_b);
        match (&a_result, &b_result) {
            (Err(e), Ok(_)) => {
                eprintln!("Bot A failed to load: {e}");
                std::process::exit(10);
            }
            (Ok(_), Err(e)) => {
                eprintln!("Bot B failed to load: {e}");
                std::process::exit(11);
            }
            (Err(ea), Err(eb)) => {
                eprintln!("Both bots failed to load: A={ea}, B={eb}");
                std::process::exit(12);
            }
            _ => {}
        }
        let player_classes = [a_result.unwrap(), b_result.unwrap()];

        // --- Phase 5: Lock filesystem access ---
        // After bots are loaded, block file opens to bot directories so
        // bots can't read each other's source code during gameplay.
        if let Some(lock_fn) = sandbox_lock {
            lock_fn.call0(py)?;
        }

        let (env, cores) = map_loader::load_map(&args.map).map_err(|err| {
            pyo3::exceptions::PyIOError::new_err(format!(
                "failed to load map {}: {}",
                args.map, err
            ))
        })?;
        let game = Game::new(env, cores, args.seed, args.suppress_indicators);

        // Seed Python's random module so bot RNG is reproducible.
        let random_mod = py.import("random")?;
        random_mod.call_method1("seed", (args.seed,))?;

        let sys_mod = PyModule::import(py, "sys")?.unbind();
        let io_mod = PyModule::import(py, "io")?.unbind();
        let gc_mod = PyModule::import(py, "gc")?.unbind();

        let mut runner = GameRunner {
            game: Rc::new(RefCell::new(game)),
            player_classes,
            unit_runners: HashMap::new(),
            watchdog,
            turn_timeout_ms: args.turn_timeout_ms,
            bot_extra_time_ns: HashMap::new(),
            sys_mod,
            io_mod,
            gc_mod,
        };
        runner.run(py)?;
        let mut game = runner.game.borrow_mut();
        let winner = game.winner_team();
        game.replay_recorder
            .write_to_path(&args.replay, winner)
            .map_err(|err| {
                pyo3::exceptions::PyIOError::new_err(format!(
                    "failed to write replay {}: {}",
                    args.replay, err
                ))
            })?;

        Ok(())
    })
}

/// Install the Python audit hook sandbox. Blocks dangerous operations
/// (ctypes memory manipulation, process spawning, network, gc introspection,
/// etc.) and removes ctypes modules from sys.modules.
///
/// Returns a lock function that, when called, blocks ALL Python-level file
/// opens (preventing /proc/self/mem reads, weight file loads, lazy imports,
/// etc.). Must be called AFTER bot loading so import-time reads still work.
fn install_sandbox(py: Python) -> PyResult<Py<PyAny>> {
    let locals = PyDict::new(py);
    py.run(
        c"
import sys

def _make_sandbox():
    blocked = frozenset({
        # ctypes memory manipulation
        'ctypes.addressof', 'ctypes.call_function', 'ctypes.cdata',
        'ctypes.cdata/buffer', 'ctypes.create_string_buffer',
        'ctypes.create_unicode_buffer', 'ctypes.dlopen', 'ctypes.dlsym',
        'ctypes.dlsym/handle', 'ctypes.memoryview_at', 'ctypes.string_at',
        'ctypes.wstring_at', 'ctypes.PyObj_FromPtr', 'ctypes._CData',
        # Process creation
        'subprocess.Popen', 'os.system', 'os.exec', 'os.fork', 'os.forkpty',
        'os.spawn', 'os.posix_spawn', 'os.kill', 'os.killpg',
        '_posixsubprocess.fork_exec',
        # Network
        'socket.__new__', 'socket.bind', 'socket.connect',
        'socket.getaddrinfo', 'socket.gethostbyname', 'socket.gethostbyaddr',
        'socket.sendmsg', 'socket.sendto',
        # GC introspection (could be used to find and tamper with sandbox state)
        'gc.get_objects', 'gc.get_referrers', 'gc.get_referents',
        # Threading (would bypass per-thread CPU time metering)
        '_thread.start_new_thread',
    })

    # fs_locked is a mutable list inside the closure — inaccessible to bot
    # code because the hook function is registered via sys.addaudithook()
    # (no public API to list hooks) and gc introspection is blocked above.
    fs_locked = [False]

    def hook(event, args):
        if event in blocked:
            raise RuntimeError('Sandbox violation: ' + event)
        if fs_locked[0] and event == 'open':
            raise RuntimeError('Sandbox violation: file access blocked during gameplay')

    def lock():
        fs_locked[0] = True

    return hook, lock

_hook, _lock_sandbox = _make_sandbox()
sys.addaudithook(_hook)

# Remove ctypes modules so bots can't re-import them.
for _mod_name in list(sys.modules):
    if 'ctypes' in _mod_name:
        del sys.modules[_mod_name]

# Remove _watchdog from sys.modules so bots can't access the watchdog.
sys.modules.pop('_watchdog', None)
",
        None,
        Some(&locals),
    )?;
    let lock_fn = locals
        .get_item("_lock_sandbox")?
        .expect("_lock_sandbox not found")
        .unbind();
    Ok(lock_fn)
}

/// Validate except handlers in all bot .py files.
///
/// Only allows handlers of the form `except Name:` or `except (Name, ...):`
/// where every name is in the whitelist of known-safe builtin exceptions plus
/// GameError. This guarantees the type expression is always a valid exception
/// type, preventing the TypeError side-channel that could be used to evade TLE.
///
/// Excluded from the whitelist: BaseException, SystemExit, KeyboardInterrupt
/// (used for TLE/engine control)
fn check_except_handlers(py: Python, main_py_path: &str) -> PyResult<()> {
    let locals = PyDict::new(py);
    locals.set_item("main_py_path", main_py_path)?;
    py.run(
        c"
import ast, os

ALLOWED = {
    # All builtin exceptions minus BaseException, SystemExit, KeyboardInterrupt
    'ArithmeticError', 'AssertionError', 'AttributeError',
    'BaseExceptionGroup', 'BlockingIOError', 'BrokenPipeError', 'BufferError',
    'BytesWarning', 'ChildProcessError', 'ConnectionAbortedError',
    'ConnectionError', 'ConnectionRefusedError', 'ConnectionResetError',
    'DeprecationWarning', 'EOFError', 'EncodingWarning', 'EnvironmentError',
    'Exception', 'ExceptionGroup', 'FileExistsError', 'FileNotFoundError',
    'FloatingPointError', 'FutureWarning', 'GeneratorExit', 'IOError',
    'ImportError', 'ImportWarning', 'IndentationError', 'IndexError',
    'InterruptedError', 'IsADirectoryError', 'KeyError', 'LookupError',
    'MemoryError', 'ModuleNotFoundError', 'NameError', 'NotADirectoryError',
    'NotImplementedError', 'OSError', 'OverflowError',
    'PendingDeprecationWarning', 'PermissionError', 'ProcessLookupError',
    'RecursionError', 'ReferenceError', 'ResourceWarning', 'RuntimeError',
    'RuntimeWarning', 'StopAsyncIteration', 'StopIteration', 'SyntaxError',
    'SyntaxWarning', 'SystemError', 'TabError', 'TimeoutError', 'TypeError',
    'UnboundLocalError', 'UnicodeDecodeError', 'UnicodeEncodeError',
    'UnicodeError', 'UnicodeTranslateError', 'UnicodeWarning', 'UserWarning',
    'ValueError', 'Warning', 'ZeroDivisionError',
    # Game-specific
    'GameError',
}

bot_dir = os.path.dirname(os.path.abspath(main_py_path))

for root, dirs, files in os.walk(bot_dir):
    for fname in files:
        if not fname.endswith('.py'):
            continue
        fpath = os.path.join(root, fname)
        with open(fpath) as f:
            source = f.read()
        tree = ast.parse(source, filename=fpath)
        for node in ast.walk(tree):
            if not isinstance(node, ast.ExceptHandler):
                continue
            lineno = node.lineno
            ty = node.type
            if ty is None:
                raise ValueError(f'{fpath}:{lineno}: bare `except:` is not allowed; use a specific exception type')
            # Only Name or Tuple-of-Names are permitted — anything else
            # (Attribute, Call, Constant, List, etc.) is rejected.
            if isinstance(ty, ast.Name):
                names = [ty.id]
            elif isinstance(ty, ast.Tuple):
                if not all(isinstance(e, ast.Name) for e in ty.elts):
                    raise ValueError(f'{fpath}:{lineno}: except handler types must be plain names')
                names = [e.id for e in ty.elts]
            else:
                raise ValueError(f'{fpath}:{lineno}: except handler types must be plain names')
            for name in names:
                if name not in ALLOWED:
                    raise ValueError(f'{fpath}:{lineno}: `{name}` is not an allowed exception type')
",
        None,
        Some(&locals),
    )?;
    Ok(())
}

fn load_player_class(py: Python, path: &str) -> PyResult<PyObject> {
    let importlib = PyModule::import(py, "importlib.util")?;
    let spec = importlib
        .call_method1("spec_from_file_location", ("player_mod", path))?
        .unbind();
    let module = importlib
        .call_method1("module_from_spec", (spec.clone_ref(py),))?
        .unbind();
    spec.getattr(py, "loader")?
        .call_method1(py, "exec_module", (module.clone_ref(py),))?;
    let player_cls = module.getattr(py, "Player")?;
    Ok(player_cls)
}

fn ensure_sys_path(py: Python, engine_root: &Path) -> PyResult<()> {
    let sys = PyModule::import(py, "sys")?;
    let path = sys.getattr("path")?;
    let py_dir = engine_root.join("py");
    path.call_method1("insert", (0, py_dir.to_str().unwrap()))?;
    path.call_method1("insert", (1, engine_root.to_str().unwrap()))?;
    Ok(())
}

fn register_rust_module(py: Python) -> PyResult<()> {
    let rust_mod = pyo3::wrap_pymodule!(rustlib::controller::controller_mod)(py);
    let controller = rust_mod.bind(py).getattr("Controller")?;
    let cambc = PyModule::import(py, "cambc")?;
    cambc.setattr("Controller", controller)?;
    Ok(())
}

