#![feature(rustc_private)]

extern crate rustc_driver;
extern crate rustc_hir;
extern crate rustc_interface;
extern crate rustc_middle;
extern crate rustc_span;

use std::env;
use std::process::ExitCode;

use rustc_driver::{Callbacks, Compilation};
use rustc_hir as hir;
use rustc_hir::def::Res;
use rustc_hir::intravisit::{self, Visitor};
use rustc_interface::interface::Compiler;
use rustc_middle::ty::{TyCtxt, TypeckResults};
use rustc_span::{Span, symbol::Symbol};

/// Crates whose purpose is to run the actor engine or turn external I/O,
/// process, telemetry, and provider API streams into actor observations.
/// Adding an entry changes the architecture.
const EXECUTION_OWNERS: &[&str] = &[
    "dashboard",
    "iroh-driver",
    "swactor",
    "swactor-engine",
    "swactor-process",
    "swactor-process-context",
    "swactor-transport",
    "swactor-vastai",
    "telemetry",
];

/// Test and benchmark harness packages that must not be workspace dependencies.
const TEST_SUPPORT_OWNERS: &[&str] = &[
    "actor-control-flow-lint-tests",
    "myelin-e2e-fuzz",
    "swactor-benchmarks",
];

/// Packages whose integration tests coordinate real child processes and file
/// locks to prove cross-process admission safety. Unit builds stay governed;
/// only `cargo test` builds of these packages use this exemption.
const TEST_COORDINATION_OWNERS: &[&str] = &["provisioning"];

/// Policy-bearing crates that must never enter an execution owner's dependency
/// closure.
const DOMAIN_CONTROL_CRATES: &[&str] = &["myelin", "provisioning", "xtask"];

#[derive(Clone, Copy)]
struct Capability {
    label: &'static str,
    resolution: &'static str,
    test_wait: bool,
    paths: &'static [&'static str],
}

const MOVE_TO_OWNER: &str =
    "move stream mechanics into an approved execution owner or move the decision into an actor";
const USE_ACTOR_TIMER: &str = "schedule a typed actor message through the engine; the receiving actor owns the deadline decision";

/// Stable resolved item paths. These are deliberately compiler identities, not
/// spellings found in source, so re-exports, renamed imports, and local wrappers
/// cannot evade the boundary.
const CAPABILITIES: &[Capability] = &[
    Capability {
        label: "asynchronous task spawning",
        resolution: MOVE_TO_OWNER,
        test_wait: false,
        paths: &[
            "tokio::runtime::Handle::spawn",
            "tokio::runtime::Runtime::spawn",
            "tokio::spawn",
            "tokio::task::spawn",
            "tokio::task::spawn_local",
        ],
    },
    Capability {
        label: "blocking task spawning",
        resolution: MOVE_TO_OWNER,
        test_wait: false,
        paths: &[
            "tokio::runtime::Handle::spawn_blocking",
            "tokio::runtime::Runtime::spawn_blocking",
            "tokio::task::spawn_blocking",
        ],
    },
    Capability {
        label: "engine task scheduling",
        resolution: MOVE_TO_OWNER,
        test_wait: false,
        paths: &["swactor_engine::EngineHandle::spawn"],
    },
    Capability {
        label: "OS thread creation",
        resolution: MOVE_TO_OWNER,
        test_wait: false,
        paths: &[
            "std::thread::Builder::spawn",
            "std::thread::Builder::spawn_unchecked",
            "std::thread::spawn",
        ],
    },
    Capability {
        label: "thread sleeping",
        resolution: USE_ACTOR_TIMER,
        test_wait: true,
        paths: &[
            "std::thread::park",
            "std::thread::park_timeout",
            "std::thread::sleep",
        ],
    },
    Capability {
        label: "direct timer driving",
        resolution: USE_ACTOR_TIMER,
        test_wait: false,
        paths: &[
            "swactor_engine::EngineHandle::interval",
            "swactor_engine::EngineHandle::timer",
            "swactor_engine::EngineHandle::timeout",
            "tokio::time::interval",
            "tokio::time::interval_at",
            "tokio::time::sleep",
            "tokio::time::sleep_until",
            "tokio::time::timeout",
            "tokio::time::timeout_at",
        ],
    },
    Capability {
        label: "runtime construction or driving",
        resolution: "the actor engine owns runtime construction and progression",
        test_wait: false,
        paths: &[
            "swactor::runtime::SingleThreadRuntime::tick",
            "swactor::runtime::SingleThreadRuntime::try_tick",
            "swactor::runtime::SingleThreadRuntime::has_work",
            "tokio::runtime::Builder::new_current_thread",
            "tokio::runtime::Builder::new_multi_thread",
            "tokio::runtime::Handle::block_on",
            "tokio::runtime::Runtime::block_on",
            "tokio::runtime::Runtime::new",
        ],
    },
    Capability {
        label: "blocking receive used as a controller",
        resolution: "receive observations in an actor; tests may use a bounded observation wait",
        test_wait: true,
        paths: &[
            "crossbeam_channel::channel::Receiver::recv",
            "crossbeam_channel::channel::Receiver::recv_deadline",
            "crossbeam_channel::channel::Receiver::recv_timeout",
            "std::sync::mpsc::Receiver::recv",
            "std::sync::mpsc::Receiver::recv_deadline",
            "std::sync::mpsc::Receiver::recv_timeout",
            "tokio::sync::mpsc::bounded::Receiver::blocking_recv",
            "tokio::sync::oneshot::Receiver::blocking_recv",
        ],
    },
    Capability {
        label: "process creation",
        resolution: "send a command to the process I/O owner and return exit/output observations to an actor",
        test_wait: false,
        paths: &[
            "std::process::Child::kill",
            "std::process::Child::try_wait",
            "std::process::Child::wait",
            "std::process::Child::wait_with_output",
            "std::process::Command::output",
            "std::process::Command::spawn",
            "std::process::Command::status",
            "tokio::process::Child::kill",
            "tokio::process::Child::start_kill",
            "tokio::process::Child::try_wait",
            "tokio::process::Child::wait",
            "tokio::process::Child::wait_with_output",
            "tokio::process::Command::output",
            "tokio::process::Command::spawn",
            "tokio::process::Command::status",
        ],
    },
];

include!("timing_policy.rs");
struct ActorControlFlowCallbacks {
    package: String,
    test_build: bool,
    trace: bool,
}

impl Callbacks for ActorControlFlowCallbacks {
    fn after_analysis<'tcx>(&mut self, _compiler: &Compiler, tcx: TyCtxt<'tcx>) -> Compilation {
        let role_exempt = TEST_SUPPORT_OWNERS.contains(&self.package.as_str())
            || EXECUTION_OWNERS.contains(&self.package.as_str())
            || (self.test_build && TEST_COORDINATION_OWNERS.contains(&self.package.as_str()));
        if EXECUTION_OWNERS.contains(&self.package.as_str()) {
            check_owner_dependencies(tcx, &self.package);
        }

        for owner in tcx.hir_body_owners() {
            let typeck = tcx.typeck(owner);
            let body = tcx.hir_body_owned_by(owner);
            let mut visitor = CapabilityVisitor {
                tcx,
                typeck,
                package: &self.package,
                caller: tcx.def_path_str(owner.to_def_id()),
                role_exempt,
                test_build: self.test_build,
                trace: self.trace,
                timing_counts: Vec::new(),
            };
            visitor.visit_body(body);
        }

        Compilation::Continue
    }
}

fn check_owner_dependencies(tcx: TyCtxt<'_>, package: &str) {
    for &crate_num in tcx.crates(()) {
        let dependency_symbol = tcx.crate_name(crate_num);
        let dependency = dependency_symbol.as_str();
        if DOMAIN_CONTROL_CRATES.contains(&dependency) {
            tcx.dcx().err(format!(
                "actor control-flow policy: execution owner `{package}` depends on domain-control crate `{dependency}`; execution owners must remain below domain policy in the dependency graph"
            ));
        }
    }
}

struct CapabilityVisitor<'a, 'tcx> {
    tcx: TyCtxt<'tcx>,
    typeck: &'tcx TypeckResults<'tcx>,
    package: &'a str,
    caller: String,
    role_exempt: bool,
    test_build: bool,
    trace: bool,
    timing_counts: Vec<(String, usize)>,
}

impl<'tcx> Visitor<'tcx> for CapabilityVisitor<'_, 'tcx> {
    fn visit_expr(&mut self, expr: &'tcx hir::Expr<'tcx>) {
        match &expr.kind {
            hir::ExprKind::MethodCall(..) => {
                if let Some(def_id) = self.typeck.type_dependent_def_id(expr.hir_id) {
                    self.check(def_id, expr.span);
                }
            }
            hir::ExprKind::Path(qpath) => {
                if let Res::Def(_, def_id) = self.typeck.qpath_res(qpath, expr.hir_id) {
                    self.check(def_id, expr.span);
                }
            }
            _ => {}
        }
        intravisit::walk_expr(self, expr);
    }
}

impl CapabilityVisitor<'_, '_> {
    fn check(&mut self, def_id: rustc_hir::def_id::DefId, span: Span) {
        let path = self.tcx.def_path_str(def_id);
        let normalized_path = normalize_def_path(&path);
        let item_name = self.tcx.item_name(def_id);
        if self.trace && is_candidate_name(item_name) {
            eprintln!(
                "actor-control-flow trace: package=`{}` caller=`{}` callee=`{path}`",
                self.package, self.caller
            );
        }

        if GUARDED_TIMING_ITEMS.contains(&item_name.as_str()) {
            let occurrence = self.next_timing_occurrence(&normalized_path);
            let allowance = TIMING_ALLOWANCES.iter().find(|allowance| {
                allowance.package == self.package
                    && allowance.caller == self.caller
                    && allowance.callee == normalized_path
            });
            match allowance {
                Some(allowance) if occurrence <= allowance.calls => return,
                Some(allowance) => {
                    self.tcx.dcx().span_err(
                        span,
                        format!(
                            "unapproved timing primitive: `{normalized_path}` occurrence {occurrence} exceeds the {} audited call(s) in `{}` (`{}`); remove the timer or update the central timing registry after architectural review",
                            allowance.calls, self.caller, allowance.purpose
                        ),
                    );
                }
                None => {
                    self.tcx.dcx().span_err(
                        span,
                        format!(
                            "unapproved timing primitive: `{normalized_path}` is forbidden in workspace package `{}` at `{}`; use typed engine scheduling or register an exact external-resource boundary",
                            self.package, self.caller
                        ),
                    );
                }
            }
            return;
        }

        let Some(capability) = CAPABILITIES
            .iter()
            .find(|capability| capability.paths.contains(&normalized_path.as_str()))
        else {
            return;
        };

        if self.role_exempt || (self.test_build && capability.test_wait) {
            return;
        }

        self.tcx.dcx().span_err(
            span,
            format!(
                "actor control-flow violation: `{}` is forbidden in workspace crate `{}`; {}",
                capability.label, self.package, capability.resolution
            ),
        );
    }

    fn next_timing_occurrence(&mut self, callee: &str) -> usize {
        if let Some((_, count)) = self
            .timing_counts
            .iter_mut()
            .find(|(known, _)| known == callee)
        {
            *count += 1;
            return *count;
        }
        self.timing_counts.push((callee.to_owned(), 1));
        1
    }
}

fn normalize_def_path(path: &str) -> String {
    let mut normalized = String::with_capacity(path.len());
    let mut cursor = 0;
    while let Some(relative_start) = path[cursor..].find("::<") {
        let start = cursor + relative_start;
        normalized.push_str(&path[cursor..start]);
        let generic_start = start + 3;
        let mut depth = 1_usize;
        let mut end = path.len();
        for (offset, character) in path[generic_start..].char_indices() {
            match character {
                '<' => depth += 1,
                '>' => {
                    depth -= 1;
                    if depth == 0 {
                        end = generic_start + offset + character.len_utf8();
                        break;
                    }
                }
                _ => {}
            }
        }
        cursor = end;
    }
    normalized.push_str(&path[cursor..]);
    normalized
}

fn is_candidate_name(name: Symbol) -> bool {
    GUARDED_TIMING_ITEMS.contains(&name.as_str())
        || matches!(
            name.as_str(),
            "block_on"
                | "blocking_recv"
                | "has_work"
                | "interval"
                | "interval_at"
                | "new_current_thread"
                | "new_multi_thread"
                | "recv"
                | "spawn"
                | "spawn_blocking"
                | "spawn_local"
                | "tick"
                | "timer"
                | "try_tick"
        )
}

fn main() -> ExitCode {
    let args: Vec<String> = env::args().collect();
    let package = env::var("MYELIN_ACTOR_LINT_PACKAGE").unwrap_or_else(|_| "unknown".to_owned());
    let test_build = env::var_os("MYELIN_ACTOR_LINT_TEST_BUILD").is_some();
    let trace = env::var_os("MYELIN_ACTOR_LINT_TRACE").is_some();
    let mut callbacks = ActorControlFlowCallbacks {
        package,
        test_build,
        trace,
    };
    rustc_driver::catch_with_exit_code(|| rustc_driver::run_compiler(&args, &mut callbacks))
}