//! Semantic control-logic security rules over the Structured Text AST.
//!
//! Each rule walks the parsed [`Pou`] and yields [`RuleHit`]s the scanner turns
//! into findings. Rules reason over structure (declarations, assignments, calls,
//! array accesses, division, jumps) rather than raw text, so they see through
//! formatting and comments.

use std::collections::{HashMap, HashSet};

use compliance_core::models::Severity;

use super::ast::*;

/// One rule match within a POU.
pub struct RuleHit {
    pub line: u32,
    pub severity: Severity,
    pub rule_id: &'static str,
    pub title: String,
    pub description: String,
    pub cwe: Option<&'static str>,
    pub remediation: &'static str,
}

/// Run every rule over a POU.
pub fn analyze(pou: &Pou) -> Vec<RuleHit> {
    let mut hits = Vec::new();
    let ctx = Ctx::build(pou);

    // Declaration-level rules.
    for v in &pou.vars {
        if let Some(init) = &v.init {
            check_credential_binding(&v.name, init, &pou.name, &mut hits);
            check_default_password(init, &v.name, &pou.name, &mut hits);
        }
    }

    // Body walk.
    walk(&pou.body, pou, &ctx, &GuardSet::default(), &mut hits);
    hits
}

/// Per-POU context precomputed once.
struct Ctx {
    /// Names declared in VAR_INPUT (untrusted / externally driven).
    input_vars: HashSet<String>,
    /// Array variable name → declared (lo, hi) bounds.
    arrays: HashMap<String, (i64, i64)>,
}

impl Ctx {
    fn build(pou: &Pou) -> Self {
        let mut input_vars = HashSet::new();
        let mut arrays = HashMap::new();
        for v in &pou.vars {
            if v.section == VarSection::Input {
                input_vars.insert(v.name.to_ascii_lowercase());
            }
            if let Some(b) = v.array_bounds {
                arrays.insert(v.name.to_ascii_lowercase(), b);
            }
        }
        Self { input_vars, arrays }
    }
}

/// Variables proven non-zero on the current control-flow path (from enclosing
/// `IF`/`WHILE` conditions), so guarded divisions aren't false-flagged.
#[derive(Default, Clone)]
struct GuardSet {
    nonzero: HashSet<String>,
}

impl GuardSet {
    fn with(&self, names: Vec<String>) -> Self {
        let mut g = self.clone();
        g.nonzero.extend(names);
        g
    }
    fn is_nonzero(&self, name: &str) -> bool {
        self.nonzero.contains(name)
    }
}

/// Variable names a condition proves non-zero (`v <> 0`, `v > 0`, `v >= 1`,
/// `v < 0`, and conjunctions thereof).
fn guards_from_cond(cond: &Expr) -> Vec<String> {
    let mut out = Vec::new();
    collect_nonzero(cond, &mut out);
    out
}

fn collect_nonzero(e: &Expr, out: &mut Vec<String>) {
    let Expr::Binary { op, lhs, rhs, .. } = e else {
        return;
    };
    let is_zero = |x: &Expr| {
        matches!(x, Expr::Int(0, _)) || matches!(x, Expr::Real(r, _) if r.abs() < f64::EPSILON)
    };
    let int_of = |x: &Expr| match x {
        Expr::Int(n, _) => Some(*n),
        _ => None,
    };
    match op {
        BinOp::And => {
            collect_nonzero(lhs, out);
            collect_nonzero(rhs, out);
        }
        BinOp::Ne => {
            if let (Some(v), true) = (lhs.as_ident(), is_zero(rhs)) {
                out.push(v.to_ascii_lowercase());
            }
            if let (true, Some(v)) = (is_zero(lhs), rhs.as_ident()) {
                out.push(v.to_ascii_lowercase());
            }
        }
        BinOp::Gt | BinOp::Lt => {
            // v > 0  or  v < 0
            if let (Some(v), true) = (lhs.as_ident(), is_zero(rhs)) {
                out.push(v.to_ascii_lowercase());
            }
        }
        BinOp::Ge => {
            // v >= n, n >= 1
            if let (Some(v), Some(n)) = (lhs.as_ident(), int_of(rhs)) {
                if n >= 1 {
                    out.push(v.to_ascii_lowercase());
                }
            }
        }
        _ => {}
    }
}

// ── the walker ─────────────────────────────────────────────────────
fn walk(stmts: &[Stmt], pou: &Pou, ctx: &Ctx, guards: &GuardSet, hits: &mut Vec<RuleHit>) {
    for s in stmts {
        match s {
            Stmt::Assign {
                target,
                value,
                line,
            } => {
                check_safety_bypass(target, value, *line, &pou.name, hits);
                // A string bound to a secret-looking target is a credential.
                if let Some(name) = flatten_ident(target) {
                    check_credential_binding(&name, value, &pou.name, hits);
                    check_default_password(value, &name, &pou.name, hits);
                }
                walk_expr(target, pou, ctx, guards, hits);
                walk_expr(value, pou, ctx, guards, hits);
            }
            Stmt::Call { callee, args, line } => {
                check_insecure_comm(callee, args, *line, &pou.name, hits);
                check_credentials_in_call(callee, args, *line, &pou.name, hits);
                for a in args {
                    walk_expr(&a.value, pou, ctx, guards, hits);
                }
            }
            Stmt::Jump { label, line } => hits.push(RuleHit {
                line: *line,
                severity: Severity::Medium,
                rule_id: "plc-unstructured-jump",
                title: "Unstructured jump (JMP) in control logic".to_string(),
                description: format!(
                    "POU `{}` uses `JMP {label}`. Unstructured jumps make control flow hard to \
                     verify and can bypass safety interlocks or leave outputs in an undefined \
                     state on unexpected paths.",
                    pou.name
                ),
                cwe: Some("CWE-691"),
                remediation: "Replace JMP with structured constructs (IF/CASE/loops); reserve \
                              jumps for well-reviewed state machines only.",
            }),
            Stmt::If {
                branches,
                else_body,
                ..
            } => {
                for (cond, body) in branches {
                    walk_expr(cond, pou, ctx, guards, hits);
                    let child = guards.with(guards_from_cond(cond));
                    walk(body, pou, ctx, &child, hits);
                }
                if let Some(b) = else_body {
                    walk(b, pou, ctx, guards, hits);
                }
            }
            Stmt::Case {
                selector,
                arms,
                else_body,
                ..
            } => {
                walk_expr(selector, pou, ctx, guards, hits);
                for (labels, body) in arms {
                    for l in labels {
                        walk_expr(l, pou, ctx, guards, hits);
                    }
                    walk(body, pou, ctx, guards, hits);
                }
                if let Some(b) = else_body {
                    walk(b, pou, ctx, guards, hits);
                }
            }
            Stmt::For {
                from, to, by, body, ..
            } => {
                walk_expr(from, pou, ctx, guards, hits);
                walk_expr(to, pou, ctx, guards, hits);
                if let Some(b) = by {
                    walk_expr(b, pou, ctx, guards, hits);
                }
                walk(body, pou, ctx, guards, hits);
            }
            Stmt::While { cond, body, .. } => {
                walk_expr(cond, pou, ctx, guards, hits);
                let child = guards.with(guards_from_cond(cond));
                walk(body, pou, ctx, &child, hits);
            }
            Stmt::Repeat { body, until, .. } => {
                walk(body, pou, ctx, guards, hits);
                walk_expr(until, pou, ctx, guards, hits);
            }
            Stmt::Return { .. } | Stmt::Exit { .. } | Stmt::Label { .. } => {}
        }
    }
}

fn walk_expr(e: &Expr, pou: &Pou, ctx: &Ctx, guards: &GuardSet, hits: &mut Vec<RuleHit>) {
    match e {
        Expr::Index { base, index, line } => {
            check_array_bounds(base, index, *line, ctx, &pou.name, hits);
            walk_expr(base, pou, ctx, guards, hits);
            walk_expr(index, pou, ctx, guards, hits);
        }
        Expr::Binary { op, lhs, rhs, line } => {
            if matches!(op, BinOp::Div | BinOp::Mod) {
                check_division(rhs, *line, &pou.name, guards, hits);
            }
            walk_expr(lhs, pou, ctx, guards, hits);
            walk_expr(rhs, pou, ctx, guards, hits);
        }
        Expr::Unary { expr, .. } => walk_expr(expr, pou, ctx, guards, hits),
        Expr::Member { base, .. } => walk_expr(base, pou, ctx, guards, hits),
        Expr::Call { args, .. } => {
            for a in args {
                walk_expr(&a.value, pou, ctx, guards, hits);
            }
        }
        _ => {}
    }
}

// ── individual rules ───────────────────────────────────────────────
const SECRET_HINTS: &[&str] = &[
    "password",
    "passwd",
    "pwd",
    "secret",
    "apikey",
    "api_key",
    "token",
    "credential",
    "privkey",
    "private_key",
    "passphrase",
];

const DEFAULT_PASSWORDS: &[&str] = &[
    "admin",
    "administrator",
    "password",
    "passwd",
    "1234",
    "12345",
    "123456",
    "0000",
    "1111",
    "root",
    "default",
    "admin123",
    "changeme",
    "letmein",
    "guest",
    "user",
    "system",
    "plc",
    "codesys",
];

const COMM_FB_HINTS: &[&str] = &[
    "modbus", "tcp", "udp", "socket", "mqtt", "opcua", "opc_ua", "ethernet", "ethip", "enip",
    "dnp3", "ftp", "telnet", "http", "send", "connect", "sock", "comm", "profinet", "s7",
];

/// Insecure cleartext service ports.
const INSECURE_PORTS: &[i64] = &[21, 23, 80, 502, 20000, 44818, 102];

fn check_credential_binding(var_name: &str, value: &Expr, pou: &str, hits: &mut Vec<RuleHit>) {
    let name = var_name.to_ascii_lowercase();
    let looks_secret = SECRET_HINTS.iter().any(|h| name.contains(h));
    if looks_secret {
        if let Expr::Str(s, line) = value {
            if !s.is_empty() {
                hits.push(RuleHit {
                    line: *line,
                    severity: Severity::High,
                    rule_id: "plc-hardcoded-credential",
                    title: "Hardcoded credential in PLC program".to_string(),
                    description: format!(
                        "POU `{pou}` binds a hardcoded secret to `{var_name}`. Credentials \
                         embedded in control logic are extracted trivially from a project export \
                         or a firmware dump and cannot be rotated without a redeploy."
                    ),
                    cwe: Some("CWE-798"),
                    remediation: "Store secrets outside the program (secure parameter store / \
                                  operator-entered, retained-but-protected memory); never commit \
                                  them to the POU.",
                });
            }
        }
    }
}

fn check_default_password(value: &Expr, var_name: &str, pou: &str, hits: &mut Vec<RuleHit>) {
    if let Expr::Str(s, line) = value {
        let lower = s.to_ascii_lowercase();
        if DEFAULT_PASSWORDS.contains(&lower.as_str()) {
            hits.push(RuleHit {
                line: *line,
                severity: Severity::Critical,
                rule_id: "plc-default-password",
                title: "Default/weak password in PLC program".to_string(),
                description: format!(
                    "POU `{pou}` uses the well-known default/weak password `{s}` (bound to \
                     `{var_name}`). Default PLC credentials are the first thing an attacker tries."
                ),
                cwe: Some("CWE-1393"),
                remediation:
                    "Require a strong, unique, operator-set password; block commissioning \
                              until the default is changed.",
            });
        }
    }
}

fn check_credentials_in_call(
    callee: &str,
    args: &[CallArg],
    line: u32,
    pou: &str,
    hits: &mut Vec<RuleHit>,
) {
    for a in args {
        if let Some(name) = &a.name {
            let n = name.to_ascii_lowercase();
            if SECRET_HINTS.iter().any(|h| n.contains(h)) {
                if let Expr::Str(s, l) = &a.value {
                    if !s.is_empty() {
                        hits.push(RuleHit {
                            line: *l,
                            severity: Severity::High,
                            rule_id: "plc-hardcoded-credential",
                            title: "Hardcoded credential passed to a function block".to_string(),
                            description: format!(
                                "POU `{pou}` passes a hardcoded secret as `{name}` to `{callee}`."
                            ),
                            cwe: Some("CWE-798"),
                            remediation: "Supply credentials from protected configuration at \
                                          runtime, not as a literal argument.",
                        });
                    }
                }
            }
        }
    }
    let _ = line;
}

fn check_safety_bypass(target: &Expr, value: &Expr, line: u32, pou: &str, hits: &mut Vec<RuleHit>) {
    let Some(name) = flatten_ident(target) else {
        return;
    };
    let n = name.to_ascii_lowercase();
    let safety = [
        "safety",
        "estop",
        "e_stop",
        "emergency",
        "interlock",
        "guard",
        "permit",
    ]
    .iter()
    .any(|h| n.contains(h));
    let watchdog = n.contains("watchdog") || n.contains("wdt");
    // A safety enable / interlock / watchdog signal driven to FALSE or 0 in
    // application logic is a bypass (e.g. `Safety_Enable := FALSE`, `Watchdog_Kick := 0`).
    let disabling = matches!(value, Expr::Bool(false, _)) || matches!(value, Expr::Int(0, _));
    if (safety || watchdog) && disabling {
        hits.push(RuleHit {
            line,
            severity: Severity::Critical,
            rule_id: "plc-safety-bypass",
            title: "Safety interlock / watchdog disabled in logic".to_string(),
            description: format!(
                "POU `{pou}` disables a safety-related signal (`{name}`) in program logic. \
                 Bypassing interlocks or watchdogs in code defeats the plant's protective \
                 functions and is a direct hazard."
            ),
            cwe: Some("CWE-1384"),
            remediation: "Never disable safety functions from application logic; safety must be \
                          handled by a certified safety controller / hard-wired circuit.",
        });
    }
}

fn check_array_bounds(
    base: &Expr,
    index: &Expr,
    line: u32,
    ctx: &Ctx,
    pou: &str,
    hits: &mut Vec<RuleHit>,
) {
    // Only reason about arrays we know the bounds of.
    let Some(arr_name) = base.as_ident() else {
        return;
    };
    if !ctx.arrays.contains_key(&arr_name.to_ascii_lowercase()) {
        return;
    }
    // Index by an untrusted input variable → potential out-of-bounds access.
    if let Some(idx_name) = index.as_ident() {
        if ctx.input_vars.contains(&idx_name.to_ascii_lowercase()) {
            hits.push(RuleHit {
                line,
                severity: Severity::High,
                rule_id: "plc-array-unchecked-index",
                title: "Array indexed by unvalidated input".to_string(),
                description: format!(
                    "POU `{pou}` indexes array `{arr_name}` with the input variable `{idx_name}` \
                     without a validated bounds check. An out-of-range index corrupts adjacent \
                     memory or faults the PLC (loss of control)."
                ),
                cwe: Some("CWE-129"),
                remediation: "Clamp or validate the index against the array bounds (e.g. \
                              `LIMIT`/explicit `IF idx >= lo AND idx <= hi`) before the access.",
            });
        }
    }
}

fn check_division(
    divisor: &Expr,
    line: u32,
    pou: &str,
    guards: &GuardSet,
    hits: &mut Vec<RuleHit>,
) {
    // A divisor proven non-zero by an enclosing guard is safe.
    if let Expr::Ident(name, _) = divisor {
        if guards.is_nonzero(&name.to_ascii_lowercase()) {
            return;
        }
    }
    // Flag division by a variable (could be zero); nonzero literals are fine.
    let risky = matches!(
        divisor,
        Expr::Ident(_, _) | Expr::Member { .. } | Expr::Index { .. } | Expr::Int(0, _)
    );
    if risky {
        hits.push(RuleHit {
            line,
            severity: Severity::Medium,
            rule_id: "plc-division-by-zero",
            title: "Division by a variable without a zero-guard".to_string(),
            description: format!(
                "POU `{pou}` divides by a variable that is not proven non-zero. A zero divisor \
                 raises a PLC exception and can halt the scan cycle (denial of control)."
            ),
            cwe: Some("CWE-369"),
            remediation: "Guard the divisor (`IF d <> 0 THEN …`) or use a safe-divide helper that \
                          returns a defined value for a zero denominator.",
        });
    }
}

fn check_insecure_comm(
    callee: &str,
    args: &[CallArg],
    line: u32,
    pou: &str,
    hits: &mut Vec<RuleHit>,
) {
    let c = callee.to_ascii_lowercase();
    let is_comm = COMM_FB_HINTS.iter().any(|h| c.contains(h));
    if !is_comm {
        return;
    }
    // Auth/encryption explicitly disabled.
    for a in args {
        if let Some(name) = &a.name {
            let n = name.to_ascii_lowercase();
            let security_flag = ["auth", "secure", "encrypt", "tls", "ssl", "authentication"]
                .iter()
                .any(|h| n.contains(h));
            if security_flag && matches!(a.value, Expr::Bool(false, _)) {
                hits.push(RuleHit {
                    line,
                    severity: Severity::High,
                    rule_id: "plc-insecure-comm",
                    title: "Network communication with security disabled".to_string(),
                    description: format!(
                        "POU `{pou}` calls `{callee}` with `{name} := FALSE`, disabling \
                         authentication/encryption on an industrial network link."
                    ),
                    cwe: Some("CWE-319"),
                    remediation: "Enable authentication + transport encryption; segment OT \
                                  networks and restrict the endpoint to trusted peers.",
                });
            }
        }
        // Well-known cleartext port literal.
        if let Expr::Int(p, _) = &a.value {
            if INSECURE_PORTS.contains(p) {
                hits.push(RuleHit {
                    line,
                    severity: Severity::Medium,
                    rule_id: "plc-insecure-protocol-port",
                    title: "Cleartext industrial protocol port".to_string(),
                    description: format!(
                        "POU `{pou}` opens `{callee}` on port {p}, a well-known cleartext OT \
                         protocol port with no built-in authentication or encryption."
                    ),
                    cwe: Some("CWE-319"),
                    remediation: "Front the protocol with a secure gateway/VPN, or use the \
                                  authenticated/encrypted variant; never expose it to untrusted \
                                  networks.",
                });
            }
        }
    }
    let _ = line;
}

/// The dotted/base identifier of an lvalue expression (`a`, `a.b` → `a.b`,
/// `a[i]` → `a`), for name-based rules.
fn flatten_ident(e: &Expr) -> Option<String> {
    match e {
        Expr::Ident(n, _) => Some(n.clone()),
        Expr::Member { base, field, .. } => flatten_ident(base).map(|b| format!("{b}.{field}")),
        Expr::Index { base, .. } => flatten_ident(base),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::pipeline::plc::parser;

    const VULN: &str = r#"
FUNCTION_BLOCK CommCtrl
VAR_INPUT
    cmdIndex : INT;
END_VAR
VAR
    Password : STRING := 'admin123';
    buffer : ARRAY[0..15] OF INT;
    Safety_Enable : BOOL := TRUE;
    divisor : INT;
    result : INT;
END_VAR
    Safety_Enable := FALSE;
    result := 100 / divisor;
    buffer[cmdIndex] := 1;
    Modbus_Connect(IP := '192.168.0.10', PORT := 502, AUTH := FALSE);
    IF cmdIndex > 100 THEN
        JMP fault;
    END_IF;
fault:
    result := 0;
END_FUNCTION_BLOCK
"#;

    fn rule_ids(src: &str) -> Vec<&'static str> {
        parser::parse(src)
            .iter()
            .flat_map(analyze)
            .map(|h| h.rule_id)
            .collect()
    }

    #[test]
    fn vulnerable_program_triggers_every_rule() {
        let ids = rule_ids(VULN);
        for expected in [
            "plc-hardcoded-credential",
            "plc-default-password",
            "plc-safety-bypass",
            "plc-division-by-zero",
            "plc-array-unchecked-index",
            "plc-insecure-comm",
            "plc-insecure-protocol-port",
            "plc-unstructured-jump",
        ] {
            assert!(
                ids.contains(&expected),
                "expected rule {expected}, got {ids:?}"
            );
        }
    }

    #[test]
    fn clean_program_has_no_findings() {
        let clean = r#"
PROGRAM Clean
VAR
    a : INT := 5;
    b : INT := 3;
    total : INT;
END_VAR
    IF b <> 0 THEN
        total := a / b;
    END_IF;
END_PROGRAM
"#;
        assert!(rule_ids(clean).is_empty(), "clean program should be quiet");
    }
}